【論文読み】 Improved Knowledge Distillation via Teacher Assistant

以下の論文の日本語訳（一部のみ）を行い、要約しました。
[\[1902\.03393\] Improved Knowledge Distillation via Teacher Assistant](https://arxiv.org/abs/1902.03393)

## Abstract
ディープニューラルネットワークは強力なモデルであり、多くのタスクで魅力的な結果を達成しているにもかかわらず、スマートフォンや組み込みセンサーノードのようなエッジデバイスに展開するには大きすぎます。これらのネットワークを圧縮する努力がなされてきたが、一般的な方法としては、事前に学習した大きな（教師）ネットワークを用いて、より小さな（学生）ネットワークを学習させる知識蒸留法がある。しかし、本論文では、**教師と生徒の間のギャップが大きくなると、生徒ネットワークの性能が低下することを示した。**学生ネットワークが固定されていると、任意の大きさの教師を採用することはできません。本研究では、この欠点を解消するために、**中間規模のネットワーク（教師アシスタント=TA）を用いて、学生と教師の間のギャップを埋める多段階知識蒸留法を導入する。**さらに、TAのサイズの影響を調べ、多段階知識蒸散へとフレームワークを拡張する。CIFAR-10,100とImageNetのデータセット、CNNとResNetアーキテクチャを用いた理論解析と広範な実験により、提案したアプローチの有効性が実証された。

## Introduction
教師モデルと生徒モデルのギャップ（サイズ差）が大きい場合、蒸留がうまくいくとは限らない。より多くのパラメータ（および精度）を備えた教師から蒸留された生徒モデルは、より少ないパラメータを備えた教師から蒸留された生徒モデルよりもパフォーマンスが悪いことを実験的に示した。

この観察に触発されて、教師と生徒モデルのギャップを埋めるために教師と生徒の間に教師アシスタント（TA）としての中間モデルを導入する、教師アシスタント知識蒸留（TAKD）と呼ばれる新しい蒸留フレームワークを提案する（図1）。TAモデルは教師から蒸留され、学生はTAからのみ蒸留される。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/4bb2d963-f165-2baa-e3c8-a4b82e5523aa.png)

私たちの貢献は、(1)教師と生徒の間のギャップの大きさ(キャパシティ)が重要であることを示す。(2) 本研究では、教師と学生の間のギャップを研究した最初の研究者であり、最大の教師では蒸留性能が最大にならないことを検証した。 (3) 本研究では、このフレームワークを拡張し、教師から学生への知識伝達をさらに向上させるために、教師から学生への複数のTAの連鎖を含むようにし、最適なTAを見つけるための洞察を提供した。

## Related Work
（省略）

## Assistant based Knowledge Distillation
### Background and Notations
（省略）

### The Gap Between Student and Teacher
以下の図2は、生徒モデルを2層CNNに固定し、教師モデルを4から10層CNNに変化させたときの教師モデルと生徒モデルの精度の変化を示している。教師モデルの層数が増えるほど、教師モデルの精度（赤破線）は高くなる。一方で、8層または10層CNNの教師モデルから蒸留された生徒モデルの精度は、6層CNNを教師モデルに用いた場合と比較して低いことが分かる。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/970de498-fe7e-0725-d72f-4237b3d34167.png)
この結果から、以下のことが言える。

1. 教師のパフォーマンスが向上し、その結果、教師はより良い予測器として学生のためにより良い教師を提供している。
2. 教師が非常に複雑化しているため、生徒はヒントを受けているにもかかわらず、教師を真似するのに十分な能力や仕組みを持っていない。
3. データに関する教師の確実性が高まるため、そのロジット（ソフトターゲット）が弱くなる。これにより、ソフトターゲットとのマッチングを介して行われる知識伝達が弱くなる。（＝教師が強力すぎるので、ソフトターゲット≒ハードターゲットとなってしまい、蒸留が機能しない）

1は蒸留性能を上げる方向だが、2と3は下げる方向である。最初は教師の規模が大きくなるにつれて1が優勢であるが、大きくなるにつれて2と3が優勢になる。

同様に、教師モデルを10層CNNに固定した場合、蒸留により最も精度が高くなる生徒モデルは何層CNNか？を考える。以下の図3から明らかな通り、教師モデルのサイズに対し、最適なサイズの生徒モデルの存在することが分かる（生徒モデルは大きければ大きいほど良いわけではない）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/cc08a922-9881-19f8-7db9-7cee1c9a45e5.png)

### Teacher Assistant Knowledge Distillation (TAKD)
事前学習された大規模な教師モデルを小規模な生徒モデルに蒸留するシナリオを考える。ここで、教師と生徒モデルのサイズを選択（変更）することはできず、両方とも固定されている状況とする。

この教師モデルと生徒モデルのギャップを埋めるために、教師と生徒モデルの中間的な規模と性能を有する教師アシスタント（TA）を提案する。教師→TA→生徒のように、間にTAを介することで、前節の2の要素を緩和する。TAから出力されるLogitは教師よりも自信のないソフトターゲットとなるので、前節の3の要素も緩和する。1の要素については、TAは精度を低下させるかもしれないが、実験や理論分析のセクションで見るように、経験的な結果と理論的な分析の両方がTAKDの有効性（パフォーマンスの向上）を実証している。これは、正の相関関係にある要因(2や3のような)を奨励することが、負の相関関係にある要因(1のような)によるパフォーマンスの低下を上回るからである。

任意の中間サイズの TA が常に知識蒸留性能を向上させることが実験で実証されている。しかし、最高の性能を得るための最適なTAサイズは何か？もしあるTAが蒸留結果を改善するのであれば、なぜ別の蒸留されたTAを介してこのTAも訓練しないのだろうか？あるいは、ゼロから訓練された TA は、我々のアプローチと同じくらい効果的なのだろうか? 以下のセクションでは、いくつかの理論的な直観を補完する経験的な視点からこれらの疑問に答えようとしている。

## Experimental Setup
データセット：CIFAR-10, CIFAR-100, ImageNet
学習条件：
- ネットワーク：プレーンCNN(Conv, BN, Max Pooling, FC層で構成されるCNN), ResNet
- 前処理：平均0、標準偏差0.5に変換
- 最適化関数：Momentum SGD (Momentum=0.9)
- 学習回数：150 epoch
- 学習率：初期値0.1、プレーンCNN(Conv, BN, Max Pooling, FC層で構成されるCNN)では初期値のまま最後まで学習。ResNetでは、80 epochで0.01、120 epochで0.001に低下させ学習。
- Weight Decay: 0.0001 (ResNetのみ)

## Results and Analysis
### Will TA Improve Knowledge Distillation?
NOKDは通常の学習、BLKDは通常の蒸留、TAKDは提案手法。データセットの規模（CIFAR-10/100, ImageNet）によらず、提案手法が最も優れている。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/c43b952d-f010-30f0-8ea5-6b4f38627ced.png)

### What is the Best TA Size?
表2より、教師サイズ10、生徒サイズ2の場合、TAのサイズは4がベスト。表3より、教師サイズ110、生徒サイズ8の場合、TAのサイズは14(CIFAR-10の場合)または20(CIFAR-100の場合)がベスト。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/282d6344-78c2-e3ec-47d9-2ddf0a3c7458.png)
なぜ中間的なサイズのモデル（教師サイズ10、生徒サイズ2の場合、TAサイズは6）がベストにならないのか？
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/1204ac1b-afe8-49ea-e532-5ffc3dbca772.png)
図4から、TAの最適なサイズは、ネットワークサイズが教師と生徒の中間となるように設定するのではなく、精度が教師と生徒の中間となるように設定するのが良さそうである。

### Why Limiting to 1-step TA?
教師と生徒の間に複数のTAを配置したらどうなるか？教師（サイズ10）から生徒（サイズ2）に直接蒸留するよりも、間にTAを1つ以上挟めば精度が向上する。最も精度が高いのは、10(→8→6→4)→2のように、間に3つのTAを挟んだ場合。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/850e32b1-7aae-3fc0-eb7a-a7e7c63cc753.png)

### Comparison with Other Distillation Methods
他の蒸留手法と比較して、提案手法(TAKD)が最も優れている。ただし、TAKDは他の蒸留手法と組み合わせることで、蒸留の精度を改善することができることに注意されたい。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/163232/ca483426-fddc-eb35-ced7-035d560c5e08.png)

## Why Does Distillation with TA work?
（省略）

## Summary
ニューラルネットワークの知識蒸留における重要な特性について研究を行った。その結果、生徒と教師のネットワーク間のギャップが知識蒸留の効果を左右すること、ギャップが大きくなると生徒のネットワークの性能が低下することを示した。本研究では、このような状況を改善するために、教師アシスタントの知識蒸留に基づいたフレームワークを提案した。本研究では、様々なシナリオにおいて本研究のアプローチの有効性を実証し、その特性を経験的・理論的に検討した。完全にデータ駆動型の自動TA選択を設計することは、今後の研究のための興味深い場である。

## Appendix
### The Best TA Sequence
リソースや時間の制約がない場合、容量の小さい順に、その間に可能なすべてのTAを順次使用することです。

原則1（知識蒸留）。知識蒸留は、分類損失のみを用いた学生学習のみの場合と比較して、学生ネットワークの精度を向上させます。

原則2（知識蒸留の性能）。同じ容量/複雑さの教師ネットワークを与えられた場合、知識蒸留フレームワークでは、精度の高い方が学生にとって良い教師となる。

原則3 (教師アシスタントに基づく知識蒸散)。生徒と教師の間に教師アシスタントを導入することで、ベースラインの知識蒸留よりも改善される。