1. はじめに
事前にクラスタ数が分からない場合でもクラスタリングとクラスタ数の推定を同時に行うモデルの一種として、ノンパラメトリックベイズモデルであるディリクレ過程混合モデル（Dirichlet Process Mixture Model, DPM）があります。通常のパラメトリックべイズモデルと比較して、ノンパラメトリックベイズモデルでは混合モデルを構成する要素分布（クラスタ）の数を無限個に拡張して考えます。
無限次元を計算上扱うことは出来ない為、実装する時はディリクレ過程（Dirichlet Process, DP）を積分消去した中華レストラン過程（Chinese Restaurant Process, CRP）や、ディリクレ過程を有限次元で近似する打切り棒折り過程（Truncated Stick Breaking Process, TSB1）などに置き換えて実装するといった工夫が必要です。
今回は前回の記事で紹介したpSGLDを使ってDPMの一種である無限混合正規分布を解くために、有限対称ディリクレ分布（Finite Symmetric Dirichlet distribution, FSD2）を使った実装を紹介します。
この記事の対象としている人
- ベイズ推定、DPMをなんとなく知っている。
- 混合モデルを実装したことがある。
- Tensorflow ProbabilityでDPMをどう実装しようか悩んでいる。
2. DPMの実装方法の検討
ノンパラベイズ、DPMの理論的背景については分かりやすい記事、書籍3が沢山ありますので、ここでは実装方法に注目して書きます。まず有限混合モデルについておさらいし、次にTSB、最後にFSDについて触れます。中華レストラン過程を使った実装はクラスタ数が動的に変わるため、Tensorflow Probabilityでは実装しにくいと考えて候補から外しています。
2.1. 有限混合モデル
無限混合正規分布を実装する場合、有限の場合と比べてどこが変わるのかまず確認します。有限混合モデルは以下の確率モデルで表すことができます。
\begin{array} { l }
{ \theta _ { k } \sim G _ { 0 } , \text { for } k = 1 , \ldots , K } \\
{ \boldsymbol{\pi} = \left( \pi _ { 1 } , \ldots , \pi _ { K } \right) \sim {\rm Dir} \left( \boldsymbol {\pi} ; \boldsymbol{\alpha}\right) },\ \boldsymbol{\alpha} = (\alpha _ { 1 } , \ldots , \alpha _ { K }) \\
z_i \sim {\rm Multi} \left(z;\boldsymbol{\pi}\right),
\text { and } x_i \sim p(x\,|\,\theta_{z_i}) \text { for } i = 1 , \ldots , N
\end{array}
クラスタ数が$K$の時、各要素分布$p$のパラメータ$\theta_k$は事前分布$G_0$からサンプルされます。また、各要素分布の混合比を表す$K$次元の$\boldsymbol{\pi}$はハイパーパラメータ$\boldsymbol{\alpha}$を持つディリクレ分布からサンプルされます。最後に、潜在変数$z_i$は混合比$\boldsymbol{\pi}$に基づいてサンプルされ、$z_i$に基づいて要素分布$p$から$x_i$がサンプルされます。
無限混合モデルではクラスタ数$K=\infty$としてモデル化します。しかし、無限次元の混合比$\boldsymbol{\pi} = \left( \pi _ { 1 } , \ldots , \pi _ { \infty } \right)$をサンプルするディリクレ過程（ただし基底分布が連続の時）＝無限事前のディリクレ分布4は実装上扱えないため、次に紹介する打切り棒折り過程、有限対象ディリクレ分布で有限次元で近似する必要があります。
2.2. 打切り棒折り過程（TSB）による実装
$K$が無限の場合、無限次元のディリクレ分布は下記に示す棒折り過程で示せることが分かっています。
\pi_k = v_k \prod_{j=1}^{K-1} (1-v_j),\,v_k \sim {\rm Beta}(v;1,\alpha)
クラスタ$k$の混合比$\pi_k$をベータ分布からサンプルされた変数$\{v_j\}_{j=k}^{K-1}$で表します。ただし$\pi_k$の合計は1であることから、その殆どは0であることが分かります。よって十分大きなクラスタ数であれば、有限次元で近似が行えます。$\alpha$は正の実数で小さいほどより少ない数の要素分布で混合モデルが表されることになります。
以下にTensorflow Probabilityで実装する場合の実装例を示します。
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.internal import distribution_util
import numpy as np
tfd = tfp.distributions
# Initial number of clusteres
K = 30
def truncated_stick_breaking_process(v):
sticks = \
tf.exp(tf.log(v)\
+distribution_util.pad(
tf.cumsum(tf.log(1-v), axis=-1)[..., :-1],
axis=-1,
front=True,
value=0))
# To make sure the sum of sticks becomes one,
# we set the last stick as one minus rest of the sticks.
enforce_sum_to_1 = \
1-tf.reduce_sum(
sticks[..., :-1],
axis=-1)
pi = tf.concat(
[sticks[..., :-1],
enforce_sum_to_1[..., tf.newaxis]],
axis=-1)
return pi
# alpha must be greater than zero
alpha = tf.nn.softplus(tf.get_variable(
'alpha',
initializer=np.ones([1], dtype=dtype)))
alpha_prior = tfd.InverseGamma(
concentration=np.ones([1], dtype=dtype),
rate=np.ones([1]),
name='rv_alpha')
# Inverse gamma distribution samples alpha
alpha_prior_log_prob = alpha_prior.log_prob(alpha)
# mixture probability
v = tf.nn.sigmoid(tf.get_variable(
'mix_probs',
initializer=np.ones([K], dtype)*(1./K)))
v_prior = tfd.Beta(
tf.ones(K, dtype=dtype),
alpha,
name='rv_beta')
# Beta distribution samples v
v_prior_log_prob = v_prior.log_prob(v)
# Apply TSB to v
pi = truncated_stick_breaking_process(v)
今回実装して分かった事ですが、pSGLDで最適化する場合、途中で上記実装中のpiやalphaの勾配がnan, infになり、収束しない場合が続出しました（たまに収束するが）。。上記の実装でHMC等他のサンプリング法を試した訳ではないのですが、TSBとpSGLDの相性が悪いのかもと感じています。参考までにPyMCにDPMをNUTSしてるサンプルがあります。
2.3. 有限対称ディリクレ分布（FSD）による実装
TSBと同じく有限次元で無限次元のディリクレ分布を近似する方法としてFSDがあり、下記式で表されます。
\boldsymbol{\pi}=\left( \pi _ { 1 } , \ldots , \pi _ { K } \right) \sim {\rm Dir}(\boldsymbol{\pi}; \frac{\alpha}{K},\ldots,\frac{\alpha}{K})
FSDで実装した場合はpSGLDでも問題なく収束しました。下記は実装です。TSBに比べてかなりシンプルに書けます。
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.internal import distribution_util
import numpy as np
tfd = tfp.distributions
# Initial number of clusteres
K = 30
dtype = np.float64
# alpha must be greater than zero
alpha = tf.nn.softplus(tf.get_variable(
'alpha',
initializer=np.ones([1], dtype=dtype)))
alpha_prior = tfd.InverseGamma(
concentration=np.ones([1], dtype=dtype),
rate=np.ones([1]),
name='rv_alpha')
# Inverse gamma distribution samples alpha
alpha_prior_log_prob = alpha_prior.log_prob(alpha)
# mixture probability
pi = tf.nn.softmax(tf.get_variable(
'mix_probs',
initializer=np.ones([K], dtype)*(1./K)))
pi_prior = tfd.Dirichlet(
concentration=np.ones(K, dtype)*alpha/K,
name='rv_fsd')
# Dirichlet distribution samples pi
pi_prior_log_prob = pi_prior.log_prob(pi)
[Kurihara+ 07]によると$K$が十分大きい場合、TSBとFSDは同じくらいのパフォーマンスになることが期待されるます。
3. 実験結果
前回の記事と同じく、共分散行列が対角行列で値が1、平均がそれぞれ[-4, 4], [0, 0], [4, 4]の２次元正規分布から5000点をサンプリングし、pSGLDとFSDを使った無限混合正規分布でクラスタリングします。
下記グラフはクラスタリング結果です。クラスタ数の初期値を30としましたが、最後は３つのクラスタに正しくクラスタリングできています。
クラスタリングの収束の過程をgifアニメーションにして見ました。最初は多数あったクラスタが、イテレーションが進むにつれて３つのクラスタに収束していることが分かります。
推定された要素分布の混合比を見ます。
３つのクラスタに混合比が集中しており、他のクラスタの混合比は殆ど０に近い値になっています。
真の平均（True loc）と上の３つのクラスタについて推定された平均（Estimated loc）、推定された共分散行列（Estimated precision）を見てみます。
True loc:
[[-4. -4.]
[ 0. 0.]
[ 4. 4.]]
Estimated loc of component 14, 17, 19:
[[-0.04240797 0.03894829]
[ 3.9698621 4.02131227]
[-4.05034999 -4.01412346]]
Estimated mix probability of component 14, 17, 19:
[3.39782416e-01 3.30332493e-01 3.29885091e-01]
Estimated precision of component 14, 17, 19:
[[1.00645757 0.99210128]
[0.99222625 0.97966251]
[1.00968866 0.99226589]]
正しく平均、共分散行列が推定できています。
4. まとめ
FSDで無限混合正規分布をモデル化し、pSGLDでサンプリングを行いました。今回色々な方の実装を参考にしましたが、一般的に使われるTSBでは上手くサンプリングが収束せず苦労しました。勾配をクリッピングしたり、値に値域を設定したりと色々試しましたが、FSDを使うことでようやく安定して収束させることが出来ました。
ノンパラベイズの実装は敷居が高めな印象でしたが、今回は割と簡単な実装でも動かすことが出来たのが驚きでした。
今回の実装はColaboratoryとしてアップしています。
時間がある時に今回実装した無限混合モデルを応用して何かやって見たいです。
5. 脚注
-
実質は打切りしていても、単に棒折り過程（Stick Breaking Process, SBP）と呼ばれている場合も見かけます。打切り棒折り過程のTSBという表記の仕方は[Kurihara+ 07]から引用しています。 ↩
-
元論文[Kurihara+ 07]ではFinite Symmetric Dirichlet Prior表記。 ↩
-
下に今回の記事を書くに当たって参考にした記事、書籍を紹介します。 ↩