はじめに

Learning deep representations by mutual information estimation and maximizationを読んだのでメモ．Abstractの最後の１文で"DIM opens new avenues for unsupervised learn- ing of representations and is an important step towards flexible formulations of representation-learning objectives catered towards specific end-goals."といっている非常に力強い論文．

気持ち

現状の表現学習はreconstruction lossを使った自己符号化器で行われるものが多いが，それは入力のデータに依存した目的関数になっているからよくないと論文では言っている．Introductionに書いてあるように"is it possible to learn representations with a training objective that is not defined in input space?"というのが論文の思い．

モチベーションとしては，reconstruction lossは相互情報量と以下の関係があるというところから，相互情報量に基づいた表現学習をしようというもの．

$\displaystyle \mathcal{I}_e(X,Y)=\mathcal{H}_e(Y)-\mathcal{D}_{KL}(\mathbb{Q}_d(Y|X)||\mathbb{P}_e(Y|X))-\mathcal{R}_{e,d}(X)$

$\mathcal{H}$ はエントロピーで下付きの $e,d$ はそれぞれencoder，decoderを表す．また， $X,Y$ は入力と中間表現の確率変数を示す．上記の式から分かるようにreconstruction loss $\mathcal{R}_{e,d}(X)$ を最小化することは相互情報量を大きくすることにつながる．すなわち，入力 $X$ と中間表現 $Y$ が強い関連を持つようになる．同様にBiGANのようなadversarialな表現学習も $D_{KL}(\mathbb{P}_d(X,Y)||\mathbb{Q}_e(X,Y))$ を最小化するという文脈で同様に相互情報量を大きくしている．

Reconstructionを目的関数とした表現学習の問題点は，reconstruction lossを個々のピクセルごとに計算する事で，解像度が高くなると背景のような意味のないピクセルに誤差が引っ張られてしまい真に意味のある表現を獲得する妨げとなる．なので直接的に相互情報量を多くしたいというのが話の流れ．

Deep INFOMAX

提案手法であるDeep INFOMAX (DIM)について．DIMは3つのロスの組み合わせからなる．ここでは入出力間の相互情報量を最大化するようにencoderを学習する問題設定を考える．パラメータ $\psi$ を持つニューラルネットを $E_\psi:\mathcal{X}\rightarrow\mathcal{Y}$ と定義する．ここで入力空間上の経験分布 $\mathbb{P}$ に従う学習サンプルの集合 $\mathcal{X}:\mathbb{X}:=\{x^{(i)}\in\mathcal{X}\}_{i=1}^{N}$ が与えられたと仮定する．ここで $\mathbb{U}_{\psi,\mathbb{P}}=E_\psi\sharp\mathbb{P}$ をpush-forward distributionとして定義する．Push-forward distributionは $\mathbb{P}$ に従うサンプルを $E_\psi$ に通す事で得られる分布のこと．以下，marginal, joint, product of marginal push-forward distributionを定義する．

$\displaystyle \mathbb{U}_{\psi,\mathbb{P}}(Y=y):=\mathbb{P}(\{x^{(i)}\in\mathbf{X}|E_\psi(x^{(i)})=y\}), \\ \displaystyle \mathbb{J}_{\psi,\mathbb{P}}(X=x,Y=y):=\mathbb{P}(X=x)\delta_y(E_\psi(x)), \\ \displaystyle \mathbb{M}_{\psi,\mathbb{P}}(X=x,Y=y):=\mathbb{P}(X=x)\mathbb{P}(\{x^{(i)}\in\mathbf{X}|E_\psi(x^{(i)})=y\}),$

ただし， $\delta_y(E_\psi(x))$ は $Y$ 上でのディラック測度．

Mutual information estimation and maximization

ここではMINEの枠組みに従って相互情報量を推定する．すなわち次のように相互情報量を求める．

$\displaystyle \mathcal{I}(X;Y):=\mathcal{D}_{KL}(\mathbb{J}||\mathbb{M})\geq\hat{\mathcal{I}}_\omega(X;Y):=\mathbb{E}_\mathbb{J}[T_\omega]-\log\mathbb{E}_\mathbb{M}[e^{T_\omega}]$

ただし， $T_\omega:\mathcal{X}\times\mathcal{Y}\rightarrow\mathbb{R}$ でパラメータ $\omega$ を持つニューラルネットを表す．MINEに関しては以前読んでメモした．この相互情報量を使ってencoder $E_\psi$ を次の目的関数に従って学習する．

$\displaystyle (\hat{\omega},\hat{\psi})=\mathrm{arg}\max_{\omega,\psi}\hat{\mathcal{I}}_\omega(X;E_\psi(X))$

ここにはMINEと重要な異なる点があり，ここではencoderとMINEが同じ目的関数で最適化される事とどちらもCNNによる計算が行われることから，両方のネットワークの初期レイヤーを $E_\psi=f_\psi\circ C_\psi,\:T_{\psi,\omega}=D_\omega\circ C_\psi$ のように組み合わせる．もう一つの違いは，相互情報量の推定をKLダイバージェンスではなくJensen-Shannonダイバージェンス (JSD)で行っている点．これらの違いにより最終的な目的関数は次のようになる．

$\displaystyle (\hat{\omega},\hat{\psi})_G=\mathrm{arg}\max_{\omega,\psi}\hat{\mathcal{I}}_{\omega,\psi}^{(JSD)}(X;E_\psi(X))=\mathrm{arg}\max_{\omega,\psi}\mathbb{E}_{\mathbb{E}_{\psi,\mathbb{P}}}[-\mathrm{sp}(-T_{\psi,\omega}(x,y(x)))]-\mathbb{E}_{\mathbb{M}_{\psi,\mathbb{P}}}[\mathrm{sp}(T_{\psi,\omega}(x',y(x)))]$

$x$ は入力のサンプルで， $y(x)$ はhigh-level representation， $x'$ は $y(x)$ に関係のないもう一方のサンプル， $\mathrm{sp}(\cdot)$ はsoftplus関数で $\mathrm{sp}(z)=\log(1+e^z)$ ．

このJSDに基づく推定器とKLに基づく推定器は似た振る舞いをするが，JSDベースの方が同時最適するのには適しているとのこと．というのも，相互情報量はunboundedであるから不安定になるのに対し，Jensen-Shannonは上界が $\log 2$ になるため，MINEで行われていたclippingの必要がない．もう一つの理由が，KLのDonsker-Varadhan表現(MINEの論文またはメモを参照)における $\log$ の期待値がバイアスの乗った勾配を導いてしまうのに対しJSDではunbiasな勾配を求めることができる．あとは単純に実験的にJSDが優れていたということらしい．ただ，JSDの場合の相互情報量の式の導出がいまいちわからなかった．

Local mutual information maximization

前節の目的関数は入出力間の相互情報量を最大化することだったが，タスクによってはあまり意味のない表現を獲得する可能性がある．というのも画像の分類問題などでは局所的に見れば意味のないピクセルが存在していて，そのような領域からくる情報は表現獲得の邪魔をする可能性がある．そこで画像分類問題のためにhigh-level representationと画像のlocal patch間の平均相互情報量を最大化することを考える．

Local mutual information maximization frameworkの概要は論文のFigure 3に書いてある．まず初めに画像を特徴マップにエンコードする $C_\psi(x):=\{C_\psi^{(i)}(x)\}_{i=1}^{M\times M}$ ( $i$ はパッチのこと)．次にこの特徴マップを大局的な特徴ベクトル $y(x)=f_\psi\circ C_\psi(x):=E_\psi(x)$ へとエンコードする．この特徴ベクトルを $\{[C_\psi^{(i)}(x),y(x)]\}_{i=1}^{M\times M}$ のようにlower-level feature mapと結合し， $1\times 1$ convolutional discriminatorによりペアスコアを算出する．

$\displaystyle T_{\psi,\omega}^{(i)}(x,y(x))=D_\omega([C_\psi^{(i)}(x),y(x)])$

上の手順で作った特徴ベクトルをrealとし，fakeは異なる画像 $x'$ を使って次のように作る．

$\displaystyle T_{\psi,\omega}^{(i)}(x',y(x))=D_\omega([C_\psi^{(i)}(x'),y(x)])$

つまりlower-level feature mapを別な画像から作るということ．この場合には次の平均JSD相互情報量を最大化することで学習を行う．

$\displaystyle (\hat{\omega},\hat{\psi})_L=\mathrm{arg}\max_{\omega,\psi}\frac{1}{M^2}\sum_{i=1}^{M^2}\hat{\mathcal{I}}^{(JSD)}_{\omega,\psi}(X^{(i)};E_\psi(X))\\ \displaystyle =\mathrm{arg}\max_{\omega,\psi}\frac{1}{M^2}\sum_{i=1}^{M^2}\mathbb{E}_{\mathbb{J}_{\psi,\mathbb{P}}}[-\mathrm{sp}(-T_{\psi,\omega}^{(i)}(x,y(x)))]-\mathbb{E}_{\mathbb{M}_{\psi,\mathbb{P}}}[\mathrm{sp}(T_{\psi,\omega}^{(i)}(x',y(x)))]$

Matching representations to a prior distribution

ここではpush-forward distribution $\mathbb{U}_{\psi,\mathbb{P}}$ をprior $\mathbb{V}$ にマッチするように学習する事を考える．priorに合わせるように学習すれば潜在表現をコントロールしやすくなって色々な場面で便利だろうという事．基本的な戦略としては次のmin-max gameによって学習する．

$\displaystyle (\hat{\omega},\hat{\psi})_P=\mathrm{arg}\min_\psi\mathrm{arg}\max_\phi\hat{\mathcal{D}}_\phi(\mathbb{V}||\mathbb{U}_{\psi,\mathbb{P}})=\mathbb{E}_\mathbb{V}[\log D_\phi(y)]+\mathbb{E}_\mathbb{P}[\log(1-D_\phi(E_\psi(x)))]$

この学習方法はadversarial autoencodersやBiGANに似ているが，generatorがないということが大きな違い．

ここまでで，最初に定式化したglobal mutual information mazimizationとパッチを使うlocal mutual information maximization，最後のprior matchingの3つの目的関数を提案した．ただ，これは全部を同時に使うこともできて，これをここではdeep INFOMAX (DIM)として提案する．

$\displaystyle \mathrm{arg}\max_{\omega_1,\omega_2,\psi}(\alpha\hat{\mathcal{I}}_{\omega_1,\psi}(X;E_\psi(X))+\frac{\beta}{M^2}\sum_{i=1}^{M^2}\hat{\mathcal{I}}_{\omega_2,\psi}(X^{(i)};E_\psi(X)))+\mathrm{arg}\min_\psi\mathrm{arg}\max_\phi\gamma\hat{\mathcal{D}}_\phi(\mathbb{V}||\mathbb{U}_{\psi,\mathbb{P}})$

$\omega_1,\omega_2$ はlocalとglobalの目的関数に関するdiscriminatorのパラメータで $\alpha,\beta,\gamma$ はハイパーパラメータ．(実験では $\alpha=0$ にしたlocal-onlyと $\beta=0$ にしたglobal-only， $\alpha=0.5,\beta=0.1$ の両方入りを使って比較してたが，local-onlyが一番うまく動いているよう．)