Tensorflowにおけるコレスキー分解を援用したGMMの実装

概要

DAGMMという異常値検知の手法をTensorflowを用いて実装した際、その中に用いられる正規混合分布(GMM)の対数尤度(エネルギー)算出部分にコレスキー分解を用いました。割とよく用いられるやり方だと思うのですが、Tensorflowでこれを実現する方法について意外に明文化された情報が少ないため、備忘録代わりに記載しておきます。間違っていたらご指摘ください。なお DAGMM の実装については、過去のポスト異常値検出手法 DAGMM の実装を参照してください。

DAGMM論文においては、多次元空間上の混合正規分布(GMM)の負の対数尤度をエネルギーと定義しています。$K$個のクラスが混合したGMMにおいて、クラス $k=1,2,..,K$ の混合比,平均,分散共分散行列をそれぞれ $\hat{\phi}_k, \hat{\mu}_k, \hat{\Sigma}_k$ とすると、このGMMに従う変数 $\boldsymbol{z}$ のエネルギー（負の対数尤度）は、次のようになります。

E(\boldsymbol{z}) = -\log \left(
\sum_{k=1}^K \hat{\phi}_k
\frac{
  \exp \left( -\frac12
    (\boldsymbol{z} - \hat{\mu}_k)^T
    \hat{\Sigma}_k^{-1}
    (\boldsymbol{z} - \hat{\mu}_k)
  \right)
}{
  \sqrt{|2 \pi \hat{\Sigma}_k|}
} \right)

ここで、$|\hat{\Sigma}_k|$ とは、分散共分散行列の行列式です。

いま、データから $\hat{\phi}_k, \hat{\mu}_k, \hat{\Sigma}_k$ が求まっている前提とします。シンプルには、分散共分散行列 $\hat{\Sigma}_k$ の逆行列を計算することが筋道ですが、逆行列の算出は不安定です。分散共分散行列を使った多次元正規分布の対数尤度の算出には、コレスキー分解を用いた算出が利用できます。

コレスキー分解による算出

コレスキー分解(Cholesky Decomposition)は、分散共分散行列のような(実)対称行列$A$に対して、下三角行列 $L$ を用いた $A=L L^T$ のような形式の分解です。また、このような分解は一意であることが示されています。

分散共分散行列 $\hat{\Sigma}_k$ をコレスキー分解し、$\hat{\Sigma}_k = L_k L_k^T$ のようにしておくと、先のエネルギーの式の $\exp$ の中身は、

\begin{align}
- \frac12
  (\boldsymbol{z} - \hat{\mu}_k)^T
  \hat{\Sigma}_k^{-1}
  (\boldsymbol{z} - \hat{\mu}_k)
&= - \frac12
    (\boldsymbol{z} - \hat{\mu}_k)^T
    (L_k L_k^T)^{-1}
    (\boldsymbol{z} - \hat{\mu}_k) \\
&= - \frac12
  (\boldsymbol{z} - \hat{\mu}_k)^T
  (L_k^{-1})^{T} L_k^{-1}
  (\boldsymbol{z} - \hat{\mu}_k) \\
&= - \frac12
  \left(
    L_k^{-1} (\boldsymbol{z} - \hat{\mu}_k)
  \right)^T \left(
    L_k^{-1} (\boldsymbol{z} - \hat{\mu}_k)
  \right)
\end{align}

と書けます。ここで、

L_k \boldsymbol{y}_k = \boldsymbol{z} - \hat{\mu}_k

となる $y_k$ を連立方程式を解いて求めておけば、先のエネルギーの式の $\exp$ の中身は単に、$- \frac12 |\boldsymbol{y}_k|^2$ となり、非常にシンプルになります。この変形を用いて、さらにエネルギーの式で $\exp$ の外に出ているものを $\exp$ の中に入れてしまうと、

E(\boldsymbol{z})
= -\log \left[
  \sum_k \exp \left(
    \log \hat{\phi}_k - \frac{K}{2} \log(2\pi)
    - \frac12 \log |\hat{\Sigma}_k|
    - \frac12 |\boldsymbol{y}_k|^2
  \right)
\right]

ここで

\begin{align}
\log |\hat{\Sigma}_k|
=& \log |L_k|  + \log |L_k^{-1}| \\
=& 2 \log |L_k| \\
=& 2 \sum_i \log L_{kii}
\end{align}

に注意すれば($L_{kii}$ は $L_k$ の$i$番目の対角成分)

E(\boldsymbol{z})
= -\log \left[
  \sum_k \exp \left[
    \log \hat{\phi}_k - \frac12 \left(
      |\boldsymbol{y}_k|^2
      + K \log(2\pi)
      + 2 \sum_i \log L_{kii}
    \right)
  \right]
\right]

のように "logsumexp" の形式に帰着します。

TensorflowでのGMMの実装

Tensorflow上では、次の便利関数が利用できます。

• tf.cholesky : 実対称行列のコレスキー分解を行う。特に、3次元以上のTensorの場合、
  最後の2次元が [..., M, M] のような shape となっていれば、コレスキー分解の対象となる
  対称行列であると考えて、すべての行列を分解する。
• tf.matrix_triangular_solve : 下三角行列が格納された行列 $A$ とベクトル $x$ に対して $A \boldsymbol{y} = \boldsymbol{x}$ を満たす $y$
  を求める。行列 $A$ としては tf.cholesky の戻り値である3次元以上のTensorをそのまま
  入れることができるし、右辺の $x$ が行列であっても構わない。

やることは通常のコレスキー分解ですが「3次元以上のTensorであっても突っ込める」というのがなかなかよろしくて、混合分布を適用するときに「かゆいところに手が届く感」です。これらを用いた具体的な算出方法を示します。以下の手順は、TensorflowにおけるGMMの実装 をかなり参考にしています。なお前提として

• n_samples : データの行数
• n_features : データの特徴次元
• n_comp : GMMのクラス数

とし、z が GMMに従うデータを格納した Tensor であり [n_samples,n_features]
の次元を持つとします。また、すでに次の値が Tensor として算出済みであるとします。

• mu : 各クラスの平均値。shape=[n_comp, n_features]
• sigma : 各クラスの分散共分散行列 shape=[n_comp, n_features, n_features]
• phi : GMMのクラスの混合比率。shape=[n_comp]

コレスキー分解は次のように実施します。

min_vals = tf.diag(tf.ones(n_features, dtype=tf.float32)) * 1e-6
L = tf.cholesky(sigma + min_vals[None,:,:])

sigma はすべてのクラスの分散共分散が一緒に格納されていますが、「最後の2次元が分解対象の行列だと思って計算する」という性質上、すべてのクラスの分散共分散行列のコレスキー分解を tf.cholesky の1回の実行でやってくれます。ただし、分散共分散が特異になるとコレスキー分解もできないので、これを避けるために小さな値を対角成分に加えます。

このLを用いて、GMMのデータごとのエネルギー（負の対数尤度）を算出します。若干ややこしいのは、tf.matrix_triangular_solve で与える右辺は、最後の2次元が [n_features, n_samples] のように並んでいる必要があるので、tf.transpose を使って次元を入れ替えます。ただ、このように並べ替えるだけですべてのデータに対してすべてのクラスの分散共分散に対応する算出を一斉にやってくれるので偉大です。

# 対象データを中心化
# z_cented.shape = [n_samples, n_comp, n_features]
z_centered = z[:,None,:] - mu[None,:,:]

# 連立方程式を解く
# v.shape = [n_comp, n_features, n_samples]
v = tf.matrix_triangular_solve(L, tf.transpose(z_centered, [1, 2, 0]))

# 分散共分散行列の行列式の対数も、Lを使って計算できる
# log_det_sigma.shape = [n_comp]
log_det_sigma = 2.0 * tf.reduce_sum(tf.log(tf.matrix_diag_part(L)), axis=1)

# log-sum-exp する前のカッコ内の値
# logits.shape = [n_comp, n_samples]
logits = tf.log(phi[:,None]) - 0.5 * (tf.reduce_sum(tf.square(v), axis=1)
    + n_features * tf.log(2.0 * np.pi) + log_det_sigma[:,None])

# log-sum-exp の適用 (n_comp の sum をとる)
# energies.shape = [n_samples]
energies = - tf.reduce_logsumexp(logits, axis=0)