Information Dropoutの数理とKeras実装

※ 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

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

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

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

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

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

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

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

\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項目は $\mathbf{x}$ から $\mathbf{z}$ への情報の伝達に関する罰則項と見ることができる。

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

Information Dropout

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

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

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

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

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

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

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

実装

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

元の論文では教師なし学習であるVariational Auto-Encoderの潜在変数 $\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関数を用いて抽象化する*1。 information_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 \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で用いる目的関数

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

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

個人的な解釈としては、 $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の実験も行い、その結果も記載している。

GitHub - ganow/keras-information-dropout: Keras implementation of the Information Dropout (arXiv:1611.01353) paper

usage: train.py [-h] [--img-size IMG_SIZE IMG_SIZE] [--beta BETA] [--kernel-size KERNEL_SIZE KERNEL_SIZE] [--nb-filters NB_FILTERS NB_FILTERS NB_FILTERS NB_FILTERS] [--model-path MODEL_PATH] [--weights-path WEIGHTS_PATH] [--lr LR] [--nb-epoch NB_EPOCH] [--batch-size BATCH_SIZE] [--name NAME] [--rootdir ROOTDIR] datapath Training All-CNN-96 network.

github.com

GitHub - ganow/keras-information-dropout: Keras implementation of the Information Dropout (arXiv:1611.01353) paper

おきもち

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