［論文紹介］Born-Again Neural Networks
こんにちは、Born Again Neural Networksというknowledge distillation系の論文で面白いものがあったのでそれを紹介していきます。
一部再現したコードをgithubにあげています。
ちなみにknowledge distillation(KD)については分かりやすい解説記事がたくさんあるのでここでは詳しい説明は省略します、、、
例えば
がめちゃくちゃオススメです。
生まれ変わって強くなる
KDでは通常の正解ラベル(=hard target)に加えて学習済み教師モデルの出力(=soft target)を正解ラベルとして使って生徒モデルを学習することで、普通に学習させるよりも精度を向上させることを目的としています。
ここで、通常教師モデルはパラメータ数の大きいネットワーク、生徒モデルはパラメータ数の小さいネットワークを使用します。つまり、KDの枠組みを利用して精度とメモリのトレードオフを改善しようという意図があるわけです。
一方Born Again Networksでは生徒モデルと教師モデルのネットワークを完全に同じにします。自分自身が教師になるというわけです。
あるいは一度学習させたパラメータを全て捨て去って、まっさらな状態から心機一転、学習し直す事で以前よりも高い精度を実現してしまおうという試みとも言えるかもしれません。
タイトルのつけ方とか、なんとなく厨二病っぽい感じが好きです。
さて、そんな虫のいい話なんてあるのか？と疑問に思う人も多いと思いますが、実は上手くいきます。
上の表はCIFAR-10に対する結果で、Teacherが元のネットワークの精度（test loss）でBANが生まれ変わったネットワークの精度です。ネットワークサイズが大きくなると生まれ変わりの効果は効かなくなってしまいます*が。。。
より難しいタスクであるCIFAR-100であればどのネットワークを使用してもBANの方が精度が高くなっています。
*おそらくネットワークが大きいと最初の学習の段階で十分「良く」学習できてしまうためそこでサチってしまうものと考えられます。
Dark Knowledge Under the Light*
*論文中の章のタイトル。これもかっこいい。
Soft targetのうち、本来の正解ラベルでない部分をDark knowledgeと言います。「教師モデルがどのように画像を認識しているか」を表すとされ、これがKDにとって本質的に大事な働きをしていると(Hinton et al., 2015)は言っています。
一方で「Dark knowledgeはそこまで重要ではなく、KDは『教師データがこのサンプルをどれぐらい自信を持って分類しているか』という指標を元にサンプル間の重み付けを変えることによってワークしているのではないか」、という仮説を立てることも可能です。
これにメスを入れているのがこのパートで、Dark knowledgeの効果を検証するために2つの実験を行なっています。
Confidence Weighted by Teacher Max(CWTM)
一つ目の実験では、soft targetのうち最大値以外をすべて0にしています。つまり、Dark knowledgeの寄与を消してサンプル間の重み付けの調整のためだけにsoft lossを使うとどうなるのか、を検証しています。
zを生徒モデルのlogits(softmax層の直前の出力値)、z'を教師モデルのlogitsとして、
sをbatchの大きさとすると、通常のKDではsoft lossの誤差逆伝播を計算するときに
となるところが
と変更されます(ここでyは本物のラベルで、Tは温度項を表しています）。
Dark Knowledge with Permuted Predictions(DKPP)
次はDark Knowledgeを完全に0にするのではなく、でたらめに入れ替えてしまうというセットアップを考えます。ただしsoft targetで最大値をとる部分については変えません。
上記の条件を満たしてランダムに入れ替える関数をφとおくと、計算式は以下のようになります。
以上2つの設定を入れた結果が以下のようになっています。 学習対象はCIFAR-100です。ちなみに右半分はBorn-Againを何回か繰り返した時の結果とそれらのアンサンブルの結果です。
通常のBorn-Againが一番精度が高くなるのは自然ですが、CWTMとDKPPも元のモデルより精度が向上していることが確認できます。なお図中のBAN+Lはsoft lossとhard lossを組み合わせた時の結果です。普通のKDではsoft lossのみより採用されやすいですが、この実験では純粋なBANより悪くなっています。
結果の再現
CIFAR-10を対象に、論文中のネットワークよりも軽いものを使ってBorn-Againの効果を確かめてみました。畳み込み層7層、プーリング層3層のVGGを軽量化したようなシンプルなCNNです。data augmentationはなしです。
Teacherのaccuracyが88.18%で、同じネットワークを学習させたところ89.10%でした。
またepoch毎に精度を記録すると、上図のように常にteacher modelを上回っていたため（たまたまいい結果が得られたわけではなく）born againの効果が確認できたと言えそうです。
