📝

データセットを要約するMMD-criticの計算量削減メモ

に公開

機械学習の解釈性の文脈で、MMD-criticを勉強した。

MMD-criticは,データ分布を理解するために代表例(prototypes)と例外例(criticisms)を同時に抽出する手法であり、分布間距離を測るMMDを基盤としている。代表例で全体像を示しつつ、説明しきれない部分を例外例として提示できるため、データ理解やモデル解釈性の向上に有効で、サンプル要約や異常検知にも応用可能。

実装にあたって最適化まわりの計算を整理したのでメモしておく。なお、手法の詳細は元論文やわかりやすい解説があるのでここでは扱わない。

プロトタイプ選択の効率化

プロトタイプ集合 S を評価する関数は次で与えられる。

J(S) = \frac{2}{n|S|} \sum_{i \in [n], j \in S} K_{i,j} - \frac{1}{|S|^2} \sum_{i, j \in S} K_{i, j}

ここで K はカーネル k(x_i,x_j) を要素に持つグラム行列。

論文ではこの関数のgreedy最適化を提案している。プロトタイプの候補を c とすると、新しいプロトタイプは上記のスコアが最大になるように選択する。つまり、

c = \arg\max_{c \in [n]\backslash S} J(S \cup \left\{c\right\}) - J(S)

論文中 Algorithm 1 では全ての候補 [n]\backslash S について逐次スコアを計算しているが、このまま実装すると非効率である。

そこで式を整理し、全候補のスコアを一括で計算できる形に変形する。

あらかじめ、以下の量を定義する。

  • グラム行列の列平均: a_j = \frac{1}{n} \sum_{i=1}^{n} K_{i,j}
  • 選択済み集合 S に対する合計: A(S) = \sum_{j \in S} a_j
  • プロトタイプ間総和: Q(S) = \sum_{i,j \in S} K_{i,j}
  • 各点と S の要素との類似度合計: r_i(S) = \sum_{j \in S} K_{i,j}

これらを使って、まずは J(S) の式の第1項を変換する。

\sum_{i \in [n], j \in S} K_{i,j} = \sum_{j \in S}\sum_{i=1}^{n} K_{i,j} = \sum_{j \in S} (n a_j) = n A(S)

新しい集合を S^{\prime} = S \cup \left\{c\right\}, |S^{\prime}|=m+1 とする。すると、

\sum_{i \in [n], j \in S^{\prime}} K_{i,j} = \sum_{i \in [n], j \in S} K_{i,j} + \sum_{i \in [n]} K_{i,c} = n A(S) + n a_c

第2項は Q(S) そのもの。cを追加したことで増えた分を考えると、

\sum_{i,j \in S^{\prime}} K_{i,j} = \sum_{i,j \in S} K_{i,j} + \sum_{i \in S} K_{i, c} + \sum_{j \in S} K_{c, j} + K_{c, c}

グラム行列は対象なので、\sum_{i \in S} K_{i, c} = \sum_{j \in S} K_{c, j} = r_c(S)。よって第2項は、

Q(S) + 2 r_c(S) + K_{c, c}

となる。

これらを合わせて完成。

J(S^{\prime}) = \frac{2}{m+1} (A(S) + a_c) - \frac{1}{(m+1)^2}(Q(S) + 2r_c(S) + K_{c, c})

スコアを整理したことで効率的な実装が可能になる。J(S^{\prime})が最大となるcを選択し、その後各統計を更新する。

1. J(S')を計算
2. cを選択
3. A(S), Q(S), r_i(S)を更新

これですべての候補についてのループ処理を排除できた。

criticism 選択の効率化

criticism 集合をCとすると、 C を評価するスコアは、

L(C) + r(K, C) = \sum_{l \in C} \left| \frac{1}{n} \sum_{i \in [n]} K_{i, l} - \frac{1}{m} \sum_{j \in S} K_{j, l} \right| + \log\det K_{C, C}

であり、これもgreedyに最適化する。

新しい criticism l を追加したときのスコアの増分を求めると

\Delta(l|C) = ( \frac{1}{n} \sum_{i \in [n]} K_{i, l} - \frac{1}{m} \sum_{j \in S} K_{j, l} ) + (\log\det K_{C \cup \left\{l\right\}, C \cup \left\{l\right\}} - \log\det K_{C, C})

となる。

正則化項(\log\detの項)の計算を簡単にしたい。Cに新しくlを追加したとき、部分行列をブロック行列で表現する。

K_{C \cup \left\{l\right\}, C \cup \left\{l\right\}} = \begin{pmatrix} K_{C, C} & K_{C, l} \\ K_{l, C} & K_{l, l} \\ \end{pmatrix}

K_{C, C} は正則なので、ブロック行列の行列式は以下のように変形できる。

\det( K_{C \cup \left\{l\right\}, C \cup \left\{l\right\}} ) = \det(K_{C, C})(K_{l, l} - K_{C, l}^T K_{C, C}^{-1} K_{C, l})

よって、正則化項の増分は、

\log\det K_{C \cup \left\{l\right\}, C \cup \left\{l\right\}} - \log\det K_{C, C} = \log\frac{\det(K_{C, C})(K_{l, l} - K_{C, l}^T K_{C, C}^{-1} K_{C, l})}{\det(K_{C, C})} = \log(K_{l, l} - K_{C, l}^T K_{C, C}^{-1} K_{C, l})

となる。対数の中身をs_lとする。

s_l(C) = K_{l, l} - K_{C, l}^T K_{C, C}^{-1} K_{C, l}

K_{C, C}のコレスキー分解K_{C, C}=LL^Tが手元にあれば、以下のように2次形式の計算はサボれる。

v = L^{-1} K_{C, l} \\ s_l(C) = K_{l, l} - v^Tv

これを利用してスコアの増分を計算し、最大となるlを選択する。

lを選択した後は,C \cup \left\{c\right\}についての新しいコレスキー分解L^{\prime}を求めたい。いったん、

L^{\prime} = \begin{pmatrix} L & 0 \\ u^T & \alpha \\ \end{pmatrix}

として、形を比較する。

L^{\prime} L^{\prime T} = \begin{pmatrix} LL^T & Lu \\ u^T L^T & u^Tu + \alpha^2 \\ \end{pmatrix}

これがブロック行列

\begin{pmatrix} K_{C, C} & K_{C, l} \\ K_{l, C} & K_{l, l} \\ \end{pmatrix}

と一致してほしい。

左上のK_{C, C}=LL^Tはもともと成り立っている。右上のK_{C, l}=Luをuについて解くと、これはスコアを計算する過程ですでに計算している(v = L^{-1} K_{C, l}を思い出す)。
ここから、\alphaが\sqrt{s_l(C)}であることが簡単に求まる。

以上をまとめてアルゴリズムが完成。

1. Lu=K_{C,l}を解いてs_l(C)を計算
2. L(C)由来の増分およびr(K,C)由来の増分を計算
3. lを選択
4. コレスキー分解Lを更新

criticism の選択も一括計算できることがわかった。

(余談) criticism 選択の正則化が多様性を促すことの確認

criticism 選択の正則化に注目する。選択されやすいcriticismはs_l(C)に登場する2次の項K_{C, l}^T K_{C, C}^{-1} K_{C, l}の値が小さい。ここで、K_{C, l}は選択済みcriticismと候補lの類似度であり、現在のcriticism集合となるべく似ていない候補の選択を促す効果が確認できる。

Discussion