今回は、Self-Distillation: Towards Efficient and Compact Neural Networks という論文を読んでみた。機械学習の1つであるニューラルネットワークについて、精度を向上させる新たなモデルである。

Self-Distillationとは？

"Distillation"とは、日本語で"蒸留"を意味する。
今回のこの"Self-Distillation"は、日本語にすると"自己蒸留"となる。
今回のSelf-Distillationは、既にあるアプローチであるKnowledge Distillationを基にしたものとなっている。まずはこれについて触れていこうと思う。

Knowledge Distillationとは？

"Knowledge Distillation"とは、日本語で"知識蒸留"を意味する。これは、まずパラメータが過剰になってしまったニューラルネットワークを教師とする。次に、この教師モデルの出力を模倣する小さな生徒モデルを訓練する。これにより、生徒モデルは教師の知識を継承しているため、パラメータが過剰となった教師モデルを置き換えていると言え、モデルの圧縮と高速な推論が可能になるというものである。

しかし、従来の知識蒸留は以下の2点が問題とされてきた。

教師モデルの選択
生徒モデルの精度は、教師モデルに大きく依存していることが分かってきた。しかも、最も高精度なモデルが必ずしも知識蒸留の教師モデルにふさわしいというわけではないことも明らかになった。そのため、適切な教師モデルを探索するためにかなりの試行が必要となり、非常に時間がかかることが問題となっていた。
知識伝達の効率
生徒モデルが教師モデルほどの精度を達成できるとは限らず、知識蒸留によって許容できないほど大幅に精度が下がる可能性もあることが分かった。知識蒸留によって、正確・効率的・コンパクトな生徒モデルを得ることは難しいとされる。

これらを解決した新しい知識蒸留法が、"Self-Distillation(自己蒸留)"である。

Self-Distillationの仕組み

Self-Distillation(自己蒸留)は、ニューラルネットワークの中間層の後に、深さが異なるAttentionベースの分類器を複数取り付けることで実現する。それぞれの分類器について、より深い分類器を教師モデルとし、生徒モデルを学習するために利用する。推論するときは、浅い分類器を無視して推論を行うためパラメータや計算の追加はない。
また、従来の知識蒸留とは異なり、教師モデルも生徒モデルも同じニューラルネットワークの分類器であるため、教師モデルの探索が必要なくなる。加えて、従来の知識蒸留は教師モデルの学習と生徒モデルの学習の2段階の学習が必要であったが、この自己蒸留は1段階で教師モデルと生徒モデルを同時に訓練することができる。

(自己蒸留を用いたResNet18のアーキテクチャ)

Self-Distillationのメリット

高精度
モデルに自己蒸留を導入することで、精度を向上することができる。CIFAR100というデータセットを用いた時は平均で3.47%、ImageNetというデータセットを用いたときは平均で2.84%の精度向上が実現した。

(ImageNetを用いた推論と正答率オレンジ色が自己蒸留を用いたもの)
高速化
自己蒸留を用いることで最終的なパラメータ数を減らすことができるため、演算量を減らすことができ結果として高速化が実現できる。
圧縮
自己蒸留で最終的に学習される分類器は、元の教師モデルより浅いネットワークとなっているため、パラメータを削減し圧縮したモデルとすることができる。

結論

Self-Distillation(自己蒸留)によって、モデルの精度向上・高速化・圧縮を同時に実現することができる。自己蒸留によって浅い分類器の精度を向上することができたため、他の推論法を大きく上回る学習ができると期待される。