Information Dropoutの数理とKeras実装

by Yoshihiro Nagano

※ This is a Japanese version. The English version is here.

Information Dropoutという手法がある。 元々は深層学習の過学習を避けるために提案されていたテクニックであるDropoutを、与えられたタスクに対する最適なデータの表現を学習するInformation Bottleneckの観点から一般化する、という内容である。

この論文は2016年11月4日にarXivに初稿がアップされ(arXiv:1611.01353)、ICLR2017に投稿されていた。 残念ながら採択はされなかったが、レビュー結果には

The authors all agree that the theory presented in the paper is of high quality and is promising but the experiments are not compelling.

とコメントされており、理論的側面は質が高いが実験が不十分だったという評価であることがわかる。 その後定期的に原稿が更新されているので、適切に実験が追加された段階でいずれかのトップカンファレンスに採択されると思っている。

この論文を読んだとき面白いなーと思ったんだけど、残念ながら実装を見つけることができなかった。 最近ちょうど深層学習フレームワークKerasの使い方を知りたいというモチベーションがあったから、年末年始に自分で実装した。 この記事ではInformation Dropoutの概要の説明とその実装について書く。

Information Bottleneck

データx\mathbf{x}が与えられたとき、タスクy\mathbf{y}を解くための最適な表現z\mathbf{z}というものを考える。 直感的には、人の様々な特徴(データ)からそれぞれの人の年収を予測したいとき(タスク)の表現としては、所属や経歴などの特徴がよい表現に思える。 一方で同じデータに対しても走る速さを予測したいタスクであれば、恐らく体力測定の結果が重要になるだろう。 このように、データに対する適切な表現はタスク依存で変わりうる。 これをより厳密に定義すると、以下の3つの条件を満たすz\mathbf{z}が最適な表現となる。

  1. z\mathbf{z}の分布はx\mathbf{x}のみに依存する。y\mathbf{y}への依存はx\mathbf{x}を介したマルコフ連鎖のみで与えられる。
  2. x\mathbf{x}からz\mathbf{z}への変換においてI(x;y)=I(z;y)I(\mathbf{x};\mathbf{y}) = I(\mathbf{z}; \mathbf{y})であり、タスクy\mathbf{y}に関する情報を失わない。
  3. 上記の2つの条件を満たすうち、最もI(x;z)I(\mathbf{x}; \mathbf{z})が小さくなるもの。

この3つの条件は、以下のような制約付き最小化問題として書くことができる。

minI(x;z)s.t.H(yz)=H(yx).\begin{aligned} \min &\:\: I(\mathbf{x}; \mathbf{z}) \\ \mathrm{s.t.} &\:\: H(\mathbf{y}|\mathbf{z}) = H(\mathbf{y}|\mathbf{x}). \end{aligned}

ここでIIは相互情報量を表す。 ラグランジュの未定乗数法を用いれば以下のInformation Bottleneck (IB) Lagrangian

L=H(yz)+βI(x;z),\mathcal{L} = H(\mathbf{y}|\mathbf{z}) + \beta I(\mathbf{x}; \mathbf{z}),

を得る。 データの情報をなるべく落としつつ、タスクに関する情報のみを抽出したものが最適な表現、というきもちがある。 p(x,y)p(\mathbf{x}, \mathbf{y})を真の分布とし、学習データ{xi,yi}i=1,2,...,N\{\mathbf{x}_i, \mathbf{y}_i\}_{i=1,2,...,N}が得られたときにpθ(zx)p_{\theta}(\mathbf{z}|\mathbf{x})pθ(yz)p_{\theta}(\mathbf{y}|\mathbf{z})を推定することを考えると、IB Lagrangianは標本平均を用いて

L(θ)=1Ni=1NEzpθ(zxi)[logpθ(yiz)]+βDKL(pθ(zxi)pθ(z)),\mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^N \mathbb{E}_{\mathbf{z} \sim p_{\theta} (\mathbf{z}|\mathbf{x}_i)} [ -\log p_{\theta} (\mathbf{y}_i|\mathbf{z}) ] + \beta D_{\mathrm{KL}} ( p_{\theta} (\mathbf{z}|\mathbf{x}_i) || p_{\theta}(\mathbf{z}) ),

のように近似する事ができる。 ここで1項目はcross-entropy誤差を表し、 2項目はx\mathbf{x}からz\mathbf{z}への情報の伝達に関する罰則項と見ることができる。

また、この目的関数でβ=1\beta=1の場合はVariational Auto-Encoderに式の意味で一致する。 つまりIB LagrangianはVAEのreconstruction errorとpenaltyのバランスを操作できる、より一般的な表式になっていると言える。

Information Dropout

ここからは上で求めたIB Lagrangianの具体的な実装を与えていく。 まずpθ(zx)p_{\theta}(\mathbf{z}|\mathbf{x})について、Information DropoutではDropoutを一般化する形で、決定的な写像f(x)f(\mathbf{x})に対してmultiplicativeなノイズϵ\epsilonを加えるモデル化

ϵpα(x)(ϵ)z=ϵf(x),\begin{aligned} \epsilon &\sim p_{\alpha(\mathbf{x})} (\epsilon) \\ \mathbf{z} &= \epsilon \odot f(\mathbf{x}), \end{aligned}

を行う。 計算の手続きを表す模式図を元の論文から引用する。

ここでpα(x)p_{\alpha(\mathbf{x})}にBernoulli分布を取った場合は通常のDropoutに完全に一致する。 pα(x)p_{\alpha(\mathbf{x})}には離散もしくは連続の任意の確率分布を取ってよいため、Dropoutのより一般的な定式化となっている。 注目すべき点は、式の通りpα(x)(ϵ)p_{\alpha(\mathbf{x})}(\epsilon)x\mathbf{x}に依存している点にある。 つまり、このモデルではDropoutで素子を落とす確率が素子ごとに入力データに依存して決まる。 これは素子の活動度に対するゲインやゲートと解釈することもできる。 また、入力を決定的な項f(x)f(\mathbf{x})と確率的な項ϵ\epsilonに分解するというアプローチはVariational Auto-Encoderにおけるreparameterization trickがベースとなっている。

次に問題となるのは目的関数2項目のKL距離による罰則項である。 ミニバッチを用いた決定的学習のためにKL距離を解析的に計算することを考える。 ここでは簡単のためにx,zx, zは1次元とする。 まず、確率変数ϵ\epsilonの分布は対数正規分布pα(x)(ϵ)=logN(0,αθ2(x))p_{\alpha(x)}(\epsilon) = \log \mathcal{N} (0, \alpha_{\theta}^2(x))に従うとする。 この時zzの事前分布を対数一様分布p(log(z))=cp(\log(z))=c、活性化関数ffをReLUf(x)=max(0,x)f(x)=\max(0, x)とするとKL距離が解析的に計算できて、

DKL(pθ(zx)pθ(z))=logαθ(x)+const.,D_{\mathrm{KL}} ( p_{\theta} (z|x) || p_{\theta}(z) ) = - \log \alpha_{\theta}(x) + \mathrm{const.},

という関係式を得る。 以上の計算からこのKL距離は罰則項としてlogαθ(x)\log \alpha_{\theta}(x)をなるべく大きくする効果がある。 解釈としては、この罰則項はDropoutによる素子のばらつきを大きくし、ロバストな空間を張る効果があると思われる。

実装

深層学習フレームワークKerasを用いた実装を以下に示す。

元の論文では教師なし学習であるVariational Auto-Encoderの潜在変数z\mathbf{z}の部分をInformation Dropoutに置き換えたものと、 教師あり学習の正則化層としてInformation Dropoutを途中の層に複数回用いた2つの実験を行っている。 ここでは後者のモデルについて実装を解説する。 以下にネットワーク構成の表を元の論文から引用する。

このモデルに対して、以下のような実装を天下り的に考える。

kernel_size = (3, 3)
nb_filters = [32, 64, 96, 192]

input_tensor = Input(shape=input_shape, name='input')
x = information_dropout_block(input_tensor, kernel_size, nb_filters[0], beta)
x = information_dropout_block(x, kernel_size, nb_filters[1], beta)
x = information_dropout_block(x, kernel_size, nb_filters[2], beta)
x = information_dropout_block(x, kernel_size, nb_filters[3], beta)

x = Convolution2D(192, 3, 3)(x)
x = BatchNormalization(axis=bn_axis)(x)
x = Activation('relu')(x)

x = Convolution2D(192, 1, 1)(x)
x = BatchNormalization(axis=bn_axis)(x)
x = Activation('relu')(x)

x = Convolution2D(10, 1, 1)(x)
x = BatchNormalization(axis=bn_axis)(x)
x = Activation('relu')(x)

x = GlobalAveragePooling2D()(x)
x = Lambda(lambda x: softmax(x))(x)

model = Model(input_tensor, x, name='All-CNN-96')

import等は省略していてKerasの関数群を利用しているため、Kerasユーザー以外はよくわからないかもしれないが、基本的に関数呼出しくらいしかしてないからぱっと見雰囲気はつかめると思う。 表の中で似たようなConvolutionとDropoutを行うブロックが4つあるので、 これをinformation_dropout_block関数を用いて抽象化する*1information_dropout_block関数は引数として 入力テンソル、カーネルサイズ、フィルタ数、information dropoutのハイパーパラメーターβ\betaを取ることにする。 以下にinformation_dropout_block関数の実装の詳細を示す。

def information_dropout_block(input_tensor, kernel_size, nb_filter, beta):

  x = Convolution2D(nb_filter, kernel_size[0], kernel_size[1])(input_tensor)
  x = BatchNormalization(axis=bn_axis)(x)
  x = Activation('relu')(x)

  x = Convolution2D(nb_filter, kernel_size[0], kernel_size[1])(x)
  x = BatchNormalization(axis=bn_axis)(x)
  x = Activation('relu')(x)

  f_x = Convolution2D(nb_filter, kernel_size[0], kernel_size[1], subsample=(2, 2))(x)
  f_x = BatchNormalization(axis=bn_axis)(f_x)
  f_x = Activation('relu')(f_x)
  logalpha = Convolution2D(nb_filter, kernel_size[0], kernel_size[1],
                           activity_regularizer=KLRegularizer(beta=beta),
                           subsample=(2, 2))(x)

  def sampling(args):
    f_x, logalpha = args

    epsilon = K.exp(K.random_normal(shape=K.shape(f_x), mean=0.,
                                    std=K.exp(logalpha)))
    return K.in_train_phase(f_x * epsilon, f_x)

  noise_x = Lambda(sampling)([f_x, logalpha])

  return noise_x

ポイントとしては確率的なサンプリングをsampling関数としてブロック内に隠蔽しているため、ネットワークの構築の際には確率的手法であることを考える必要がなくなっている。 また、Information Dropoutを導入することによる目的関数の罰則項の部分はRegularizerとして実装することで、こちらについても実装を隠蔽している。 KLRegularizerの実装は以下のとおりである。

class KLRegularizer(Regularizer):

  def __init__(self, beta=0.0):
    self.beta = K.cast_to_floatx(beta)

  def __call__(self, logalpha):
    regularization = 0
    regularization += - self.beta * K.mean(K.sum(logalpha, keepdims=0))
    return regularization

罰則項は前述の解析解を利用すれば、Dropoutの確率を決めるパラメーターであるα\alphaについてlogαθ(x)-\log \alpha_{\theta}(\mathbf{x})で罰則をかける形になる。 その実装は上のように非常にシンプルな形になる。

*1: 関数ブロックで抽象化するというアイデアはResNetのKeras公式サンプルを参考にした。

実験結果

今回の分類タスクでは、Cluttered MNISTという元の論文でも使用されたデータセットを用いた。 Cluttered MNISTは0から9までの手書き文字数字にランダムにノイズを加えたデータセットで、ノイズの大きさや量をパラメーターとしてコントロールできる。以下にサンプルを示す。

上に示すような画像について50000枚の訓練データ、10000枚のテストデータを用意して学習を行った。 学習後のテストデータに対するconfusion matrix (混同行列)は以下の図のようになった。

大まかには学習がうまくいっているように見える。 1と7、2と7、4と9あたりの組み合わせで少し間違えているが、 このあたりの数字は似た形なので人間の直感的にも合う間違え方に思える。

次に、元論文にあった実験で個人的に一番おもしろいと思ったところを追試した。

上に示す図は、Information Dropoutで用いる目的関数

L(θ)=1Ni=1NEzpθ(zxi)[logpθ(yiz)]+βDKL(pθ(zxi)pθ(z)),\mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^N \mathbb{E}_{\mathbf{z} \sim p_{\theta} (\mathbf{z}|\mathbf{x}_i)} [ -\log p_{\theta} (\mathbf{y}_i|\mathbf{z}) ] + \beta D_{\mathrm{KL}} ( p_{\theta} (\mathbf{z}|\mathbf{x}_i) || p_{\theta}(\mathbf{z}) ),

の第2項、KL距離を使った罰則項の部分を入力データ空間で可視化した結果である。 今回の実装では入力x\mathbf{x}からf(x)f(\mathbf{x})logαθ(x)\log\alpha_{\theta}(\mathbf{x})の2つに分岐させる実装を行ったので、後者の素子の活動度をそのまま出力した。 最左列に示す入力画像に対して、Information Dropoutを行った層を浅い順にDropout 0, 1, 2で並べている。 KL距離DKL(pθ(zxi)pθ(z))D_{\mathrm{KL}} ( p_{\theta} (\mathbf{z}|\mathbf{x}_i) || p_{\theta}(\mathbf{z}) )は表現z\mathbf{z}の事前分布pθ(z)p_{\theta}(\mathbf{z})に対してデータx\mathbf{x}が与えられた際にどのくらい分布の形が変化したかを定量化しているため、入力データの中で識別に重要と思われる部分を抽出している。

元の論文ではβ\betaに適切な値を設定した場合に、 データの中で識別に関係する領域、この入力データでは数字の5の部分のみが徐々に抽出されるという結果が現れていた。 しかし、今回の実験ではβ\betaを含むハイパーパラメーターのチューニングが足りなかったのか、ノイズの部分もそれなりに拾ってしまっており、完全に再現することはできなかった。

個人的な解釈としては、DKL(pθ(zxi)pθ(z))D_{\mathrm{KL}} ( p_{\theta} (\mathbf{z}|\mathbf{x}_i) || p_{\theta}(\mathbf{z}) )という形はsaliency mapとかBayesian surpriseに近い概念なのでは?と思っていて、その辺の指標を陽に教師データにして学習することなく中間層に表現が出てくることがおもしろいと思った。

Putting it all together

今回紹介したコードは以下のリポジトリから本質的な部分だけ抽出したものである。 加えて、以下のリポジトリではVariational Auto-Encoderの実験も行い、その結果も記載している。

おきもち

タイトルに入れた数理というキーワードはあおり気味だったと反省している(数理要素の少なさ)。