Dynamic Routing Between Capsules を読んだ

TL;DR

  • neuron をまとめてベクトル化した capsule という概念を導入。各成分は特徴量の表出を parametrize
  • routing という手法を導入し、in と out の内積値で低次カプセルから高次カプセルへの値の受け渡しを記述
  • CapsNet というモデルを具体的に構成し、特に MNIST の画像を重ねた MultiMNIST で高い性能を発揮
  • capsule は CNN では効率的に扱えてない平行移動以外のアフィン変換を効率的に扱い得るもので、発展性がある

モチベーション

まずもって論文の書き出しからしてちょっと何が言いたいのか分からない。議論についていけなくてすみませんという感じだ。

とは言うものの頑張って読んでみるとしよう。 まず、人間の視覚は物体の特定領域をうまく選び出してそこを注視していくことで、それ以外の irrelevant な詳細はうまく落とす。 物体を適切に認識するという高次機能は、これらの部分的な情報をうまく組み合わせて全体像を形成することで実現する。 この論文では、必要な部分的情報がニューロンをまとめた capsule という単位が activate することで拾われていくという構造を提案している。

これを構文木とのアナロジーでも説明していて、子ノードが activate された capsule に対応していて、多層のニューラルネットワークから学習によってこの構文木の構造が削り出されてくるということになる。 ただ、この解釈はどうにも論文の記述とマッチしないように思える。 各 fixation (これが画像の特徴的な部分に対応しているはず) ごとに構文木のような構造を構築するという感じのことを書いていて、fixation の方が親ノード的な書き方に見える。 上の解釈的には抽象的な概念が親ノードになるのかと思っているのだが、記述としてはそうなっているようには見えない。

うーん、自分がちゃんと理解できてないところがあるんだろうな、少し寝かせておこう。

capsule でまとめられた neuron の活性は画像中の特定の性質（位置や変形度合いやテクスチャなど）を表現していて、個々の出力を logistic として扱うことでその性質の存在確率と解釈する。 この性質に関しては、論文を見る限りは画像に関する特徴量何でもという印象を受ける。 そんなに解釈可能性が高いものが well-separated に encode されるような気はしないのだが、後で見るように MNIST での実験では理解が容易な性質が現れることが確認できる。

この論文では CNN のアーキテクチャを踏襲している。 つまり、ある部分で学習した結果得られた良い特徴量を他でも使い回す。 これが画像分析でうまくいくことを知っているためである。 ただしニューロンのスカラー出力をベクトル出力に置き換え、かつ max pooling を dynamic routing というものに置き換える。

この dynamic routing とは、子カプセルの出力に重みを掛けたものと親カプセルの出力の内積を取り、カプセルが揃っている親の方へ出力を受け渡すようにするものである。 ちょっとややこしいのは重みに関しては back propagation で学習するが、子カプセルと親カプセルの結びつきに関しては dynamic という言葉からも分かるように自動的に update される点にある。 この論文の技術的な部分はほとんどこれに尽きているのだが、これは具体的な数式を見るのが早いだろう。

capsule と dynamic routing の理解

肝になってるのは下記の routing algorithm である。こいつを理解しよう。

初期化している bijb_{ij} は子カプセルと親カプセルの結びつきの強さを表現するものである。

まずはこれを softmax をとって規格化する。 ここでちょっと気を付けたいのは softmax がベクトル表記になっていることである。 これは子カプセルの足iに対して親カプセルjを成分に持つようなベクトルとして表現しているからであるが、これはあくまでカプセル同士の結びつきの話でカプセルの中身の要素とは関係がない。

次に sj=∑icijWijui\rm{s_j} = \sum_i c_{ij} \rm{W_{ij}} \rm{u_i} で親カプセルの出力を計算する。 ui\rm{u_i} は子カプセルの出力で、Wij\rm{W_{ij}} は back propagation で学習する重み行列、sj\rm{s_j} が親カプセルの出力である。 このベクトル表記は成分がカプセルの要素になっているのが上のベクトル表記との違いである。 重み行列は子カプセルの次元から親カプセルの次元に重み付きで変換していることになる。

これで計算された出力は方向を保ったままノルムを [0,1] に収めるように squashing という変換をかます。

最後に子カプセルの出力を重み行列で変換したものと squash した親カプセルの出力の内積を取り、その値に基づいてカプセル同士の結びつきの係数を update する。 これを指定回数だけ繰り返すということをしている。

これはカプセルがベクトルになっているからこそ可能なアルゴリズムであることは注目に値する。 カプセルに encode された性質が、子カプセル（に重み行列を掛けたもの）と親カプセルで大きさも方向もマッチする場合に結びつきが強くなるということになっている。

この routing は２つの layer 間で完結している。 重み行列に関しては back propagation で学習するので、出力層で loss function を定義する必要がある。それが以下である。

TkT_k がデータがクラスkの場合は1でそれ以外は0になる。 他にもパラメタがあるが、この論文ではこれを採用したというものでしょう。 この辺を掘っていくのは一つの方向性になるのではないだろうか（そこまで面白いことが起こりそうな感じはしないが）？

CapsNet の構造

具体的に MNIST のデータに対して分類モデルを構築しよう。 まずは以下のアーキテクチャが何を意味してるのか理解するところから。

最初の層は普通の convolution と ReLU である。 kernel size とか channel 数が多い気もするが、この辺はどれくらい tune しているかは明らかにされてない。

次の PrimaryCaps で capsule が登場する。 奥行きの8の次元が capsule のベクトルの要素になっている。 前の層は普通の convolution+ReLU なので、ここには routing が適用されない。 単にconvolution によって capsule が切り出されている。

次の層が DigitCaps であり、これも capsule になっている。 16が capsule のベクトルの要素であり、10は MNIST のクラス数に対応している。 重み行列は PrimaryCaps の層のカプセルの次元から DigitCaps の次元へと変換する。 この2層間は routing algorithm によって結びつきが dynamical に update されることになる。

最後はカプセルの出力を用いて loss function を計算していて、ここから back propagation で重みを更新するのに使われる。

これでモデルは良さそうな気がするが、論文ではもうひとつ構造を加えている。 DigitCaps の情報から元の画像を再構築する MSE loss も加えることで、画像の特徴学習を enhance している。

再構築の loss は 0.0005 の重みで加えるなど tune している感満載だが、ともかくこれを入れることで性能向上が示される。

実験

MNIST と MNIST の画像1枚に対して1000枚他の数字を重ねて 60M のデータとした MultiMNIST で検証している。 後者に関しては予測対象が2つになるが、学習は1クラスずつなので、同じ画像を2回別々のクラスとして使うのか重ねる前のベースのクラスのみを使うのかは分かってない（後者かな？）。 予測の際は出力の大きな2つのカプセルを予測クラスとして採用している。

結果は以下の通りで、特に MultiMNIST で良い結果を実現している。 Baseline は計算量が同じ程度になるように調整した CNN である。

再構築の結果はこんな感じ。 良い感じに再構築できているように見える。 この辺は conditional VAE 的な雰囲気も感じる。

獲得した capsule の要素が結構面白い結果になっている。 DigitCaps の16次元の要素を1つずつ0.05刻みで変えた結果が以下の図である。 スケーリングとか太さ変換とかが性質として獲得されていることを確認することができる。 もっと複雑な物体になったらこんなに綺麗にはいかないとは思うし、類似の話もこれまであったわけだが、もっと研究を進めていくに足るような結果ではないだろうか。

最後に MultiMNIST の分離の結果を載せる。 R(,) が再構築をする際に使ったラベルであり、L(,) が正解のラベルを表す。 画像の上は MultiMNIST の画像で、数字が2つ重なっていることが分かる。 画像の下は再構築した画像で、R(,) のラベルを1つずつ使って再構築した画像を色違いで重ねているものになっている。

左4つはうまくいっている例、次の2つは再構築に他のラベルを使った場合、右の2つは右がどのように予測を間違えたか（そして間違えたラベルで再構築した結果）で左が正しいラベルで再構築した結果である。

ポイントは同じピクセルでも別々のラベルのために使われていて、物体が重なっているような状況でもうまく機能していることが見て取れる。

まとめ

ニューロンの機能を一次元拡張（スカラーからベクトルに）するような capsule という概念を導入し、高次の特徴を効率的に表現する方法を提案。 routing というアルゴリズムで capsule の input (capsule i から来た capsule j の input) と output (capsule j の output) の中身が揃っている場合に、i と j の結びつきが強くなるように結合を update する。 これの気持ちとしては、CNN では max pooling などで潰していた細かい情報を capsule に保持させておいて、その情報が重要な場合により高次の layer にその情報を運んでいくという感じでしょうか。 この時に一次元拡張していることが重要になり、具体的には l 番目の層の capsule は l-1 番目の層の異なる capsule の異なる特徴を同時に考慮することができる。 capsule の次元の異なる要素が異なる特徴に対応し、それが routing によって増強されていくからである。 論文中ではこれを parallel attention mechanism と呼んでいる。 このことは、capsule の概念を具体的なモデルである CapsNet で構築し、MNIST で数字が重なっている場合にそれらを適切に分離して認識するタスクで高い性能を発揮していることからも示唆されている。

色々と示唆深いことが書かれているのだが、なかなか全部は咀嚼しきれないところもある。 今後の発展に応じてこの論文に立ち返って来ることで、自分に見えてくる世界が広がってくるのではないかと期待したいところだ。

ちなみに GitHub で各種フレームワークにおける実装も公開されている。 実装としては TensorFlow実装 とか Keras実装 辺りを今度使ってみようかな。

capsule 自体は2011年に この論文 で Hinton が提案しているので結構前から存在する概念のようだ。 この論文を読むだけでは capsule がどれくらい広がりのある概念かは想像が容易まではいかないが、きっと Hinton らはこの辺のことをずっと考えていて、もっと広い世界が見えているのだろう。 ということで、正直自分が書いた解釈が正しいかどうかはちょっと自信がない。 世間ではかなり盛り上がっているようなので、今後の発展が楽しみにして自分の理解も深めたい。