はじめに

SELF-LABELLING VIA SIMULTANEOUS CLUSTERING AND REPRESENTATION LEARNINGを読んだのでメモ． K-Meansクラスタリングによりデータにラベル付け(self-labeling)を行うDeepClusterに対し，self-labelingで生成されたラベルとモデルの出力間のクロスエントロピーを最適輸送問題として考えることでself-labelingを行う手法．

Method

事前準備として色々定義する．まず，データ $I$ を適当な空間マッピングするニューラルネット $\boldsymbol{x}=\Phi(I)$ を考える．ただし， $\boldsymbol{x}\in\mathbb{R}^D$ は特徴ベクトル．さらに，特徴ベクトルを入力とし識別を行うclassification head $h:\mathbb{R}^D\rightarrow\mathbb{R}^K$ を考える．ただし， $K$ はカテゴリ数を表す．データは $N$ 点 $I_1,\dots,I_N$ あるとし，それぞれのデータに対応するラベルを $y_1,\dots,y_N\in\{1,\dots,K\}$ とする．最終的な分類スコアはsoftmax関数を用いて下記のように計算される．（論文の通りに書いたが，ここまでの定義に従えば $\Phi$ を挟んでいるので $\boldsymbol{x}_i$ は $I$ となる気がする．）

$\displaystyle p(y=\cdot|\boldsymbol{x}_i)=\mathrm{softmax}(h\circ\Phi(\boldsymbol{x}_i))$

モデルは下記のクロスエントロピーを最小化するように学習が進められる．

$\displaystyle E(p|y_1,\dots,y_N)=-\frac{1}{N}\sum_{i=1}^N\log p(y_i|\boldsymbol{x}_i)$

上記の学習は，ラベルが利用できない場合にはself-labelingの方法によって自動的にラベル付けを行う必要がある．

半教師あり学習などの枠組みではモデル $h\circ\Phi$ とラベル $y_1,\dots,y_N$ の同時最適化となるが，完全な教師なし学習では全てを同一ラベルにアサインすることによってコストを簡単に最小化してしまう．そのためここではラベルを事後分布 $q(y|\boldsymbol{x}_i)$ として下記のようにクロスエントロピーを書き直す．

$\displaystyle E(p,q)=-\frac{1}{N}\sum_{i=1}^N\sum_{y=1}^Kq(y|\boldsymbol{x}_i)\log p(y|\boldsymbol{x}_i)$

仮に事後分布をクロネッカーのデルタ $q(y|\boldsymbol{x}_i)=\delta(y-y_i)$ として決定的に定義をしたのならば上記の式は最初に定義したクロスエントロピーの式に等しい([E(p,q)=E(p|y_1,\dots,y_N)])．この式でも自明な解にたどり着くため，下記のように制約を与えることでこれを避ける．

$\displaystyle \min_{p,q}E(p,q)\ \text{subject to}\ \forall y:q(y|\boldsymbol{x}_i)\in\{0,1\}\ \text{and}\ \sum_{i=1}^Nq(y|\boldsymbol{x}_i)=\frac{N}{K}$

この制約は，各データには一つのラベルしか割り当てが行われず，各カテゴリに属するデータ数が均一であることを意味する．するとこの問題は最適輸送理論の問題として考えることが可能となる．より具体化するため $P_{yi}=p(y|\boldsymbol{x}_i)\frac{1}{N}$ をモデルから予測された $K\times N$ の同時確率行列とし，同様に $Q_{yi}=q(y|\boldsymbol{x}_i)\frac{1}{N}$ を $K\times N$ の同時確率行列とする．ここでは $Q$ を次のような輸送多面体の要素(すなわち $\{0,1\}$ ではなく $(0,1)$ の連続値)に緩和する．

$\displaystyle U(r,c):=\{Q\in\mathbb{R}^{K\times N}_+|Q\mathbb{1}=r,Q^\top\mathbb{1}=c\}$

$\mathbb{1}$ は全要素が1の列ベクトルで， $Q\mathbb{1},Q^\top\mathbb{1}$ は $Q$ を列方向，行方向へ積分することを意味している．ここで $Q$ が条件付き確率となることを要請する．すなわち

$\displaystyle r=\frac{1}{K}\cdot\mathbb{1},\ c=\frac{1}{N}\cdot\mathbb{1}$

これにより元の目的関数(クロスエントロピー)は次のように書き直される．

$\displaystyle E(p,q)+\log N=\langle Q,-\log P\rangle$

$\langle\cdot\rangle$ はフロべニウス内積（要素積+総和）で $\log$ は行列の各要素ごとに適用されるものとする．結果として最適な $Q$ は次の最小化問題を解くことで得られる．

$\displaystyle \min_{Q\in U(r,c)}\langle Q,-\log P\rangle$

これは線形計画法で溶けるが，通常データ数がミリオンオーダーでクラス数が数千になるので解くことが困難である．ここではSiknhorn Distance(以前まとめたブログ記事)で導入された次のような正則化を導入することで高速に解く方法を提案する．

$\displaystyle \min_{Q\in U(r,c)}\langle Q,-\log P\rangle+\frac{1}{\lambda}\mathrm{KL}(Q\| rc^\top)$

$\mathrm{KL}$ はKullback-Leiblerダイバージェンスで $rc^\top$ は $K\times N$ の確率行列．この正則化の導入により上記の最小化問題のかいが次のように表現される．

$\displaystyle Q=\mathrm{diag}(\alpha)P^\lambda\mathrm{diag}(\beta)$

注意として $P^\lambda$ は行列の $\lambda$ 乗ではなく要素毎の $\lambda$ 乗を表す． $\alpha$ と $\beta$ は $Q$ が確率行列の性質を満たすためのスケーリング係数ベクトルで単純な繰り返し計算で得られる．面白い知見としては， $\lambda$ は大きな値を取れば元の最適輸送問題と等しくなるが，ここでは $\lambda$ を大きくして正確に解くよりも適当な固定の $\lambda$ を使う方が良い結果が得られた．おそらく $Q$ の緩和により，各クラスに属するデータが均一であるという制約も多少緩和されるためや，smooth labelを使うと良いという知られた知見によるものと考えられる．

最終的なアルゴリズムとしては，モデル $h\circ \Phi$ と $Q$ の最適化を交互に解くことで表現学習をする．より具体的には

現在得られている $Q$ に対し，クロスエントロピーの最小化により $h\circ\Phi$ のパラメータを更新．
現在得られているモデル $h\circ\Phi$ に対し $P_{yi}=\mathrm{softmax}_y(h\circ\Phi(\boldsymbol{x}_i))$ を計算し， $\alpha$ と $\beta$ を次のように更新することで $Q$ を得る．

$\displaystyle \forall y:\alpha_y\leftarrow[P^\lambda\beta]^{-1}_y\ \forall i:\beta_i\leftarrow[\alpha^\top P^\lambda]_i^{-1}$

$Q$ の更新は行列ベクトル積であり計算オーダーは $\mathcal{Q}(NK)$ であるため，線形計画法で解くよりもリーズナブルに計算が可能(とは言え $NK\approx10^9$ くらいなのでどうなのかという疑問がある)．論文によれば繰り返し計算はGPUで収束まで2分程とのこと．

解釈

ここまでのアルゴリズムが結局何をやっているのかを考える．簡単のため，データ $\boldsymbol{x}_i$ をそのインデックス $i$ として $p(y|\boldsymbol{x}_i)=p(y|i),q(y|\boldsymbol{x}_i)=q(y|i)$ と書き換える．すると

$\displaystyle E(p,q)+\log N=-\sum_{i=1}^N\sum_{y=1}^Kq(y,i)\log p(y,i)=H(q,p)$

とかける．この最小値は $p=q$ の時に得られる．ここで， $q(i)=1/N$ ，周辺化エントロピー $H_q(i)=\log N$ が定数， $q(y)=1/K$ から $H_q(y)=\log K$ もまた定数と仮定する．するとエントロピーから二つの定数を引くことで

$\displaystyle \min_pE(p,q)+\log N=E(q,q)+\log N=H_q(y,i)=H_q(y)+H_q(i)-I_q(y,i)=\mathrm{const.}-I_q(y,i)$

と相互情報量が導かれる．すなわち結果として元の問題は相互情報量の最大化に等しくなっている．

まとめ

DeepClusterとの関連や，data augmentationの活用，マルチタスクでの学習に関する説明は省略した．データ数が各クラス毎に均一という仮定を導入することで最適輸送問題としてラベリングを定式化する着眼点が非常に面白かった．一方で， $Q$ をデータ数に対して線形時間で得ることができると言っているがデータ数を考えるとそんな簡単な話でもない気がするがどうなのか．また $Q$ の更新の際に $P$ を計算する必要があるが，それもナイーブには全データを一回モデルに入力する必要があるので実装上どう実現されているか気になるところ(実験の章はまともに読んでいないので書いてあるかも…)．