[論文レビュー] NORM: Knowledge Distillation via N-to-One Representation Matching

Xiaolong Liu, Lujun Li+2 more|arXiv (Cornell University)|May 23, 2023

Advanced Neural Network Applications被引用数 19

ひとこと要約

NORMは、学生ネットワークの最後の畳み込み層の後に軽量な線形特徴変換を挿入することで、多対一の表現マッチング機構を導入します。学生の特徴をN倍のチャネルへ拡張し、Nセグメントに分割して、それらを教師の表現と共同でマッチングさせ、推論時オーバーヘッドなしにより多くの転送経路を可能にします。

ABSTRACT

Existing feature distillation methods commonly adopt the One-to-one Representation Matching between any pre-selected teacher-student layer pair. In this paper, we present N-to-One Representation (NORM), a new two-stage knowledge distillation method, which relies on a simple Feature Transform (FT) module consisting of two linear layers. In view of preserving the intact information learnt by the teacher network, during training, our FT module is merely inserted after the last convolutional layer of the student network. The first linear layer projects the student representation to a feature space having N times feature channels than the teacher representation from the last convolutional layer, and the second linear layer contracts the expanded output back to the original feature space. By sequentially splitting the expanded student representation into N non-overlapping feature segments having the same number of feature channels as the teacher's, they can be readily forced to approximate the intact teacher representation simultaneously, formulating a novel many-to-one representation matching mechanism conditioned on a single teacher-student layer pair. After training, such an FT module will be naturally merged into the subsequent fully connected layer thanks to its linear property, introducing no extra parameters or architectural modifications to the student network at inference. Extensive experiments on different visual recognition benchmarks demonstrate the leading performance of our method. For instance, the ResNet18|MobileNet|ResNet50-1/4 model trained by NORM reaches 72.14%|74.26%|68.03% top-1 accuracy on the ImageNet dataset when using a pre-trained ResNet34|ResNet50|ResNet50 model as the teacher, achieving an absolute improvement of 2.01%|4.63%|3.03% against the individually trained counterpart. Code is available at https://github.com/OSVAI/NORM

研究の動機と目的

• One-to-One Representation Matching (ORM)を超える形で改良された二段階知識蒸馴の推進。
• 学生の最後の畳み込み層の後に最小限で吸収可能な特徴変換を挿入して教師情報を保持する。
• 特徴拡張と分割を通じて多対一の知識伝達を可能にし、推論時パラメータを追加せずに転送経路を増やす。
• CIFAR-100とImageNetで最先端のKD性能を示し、ロジットベースKDおよびコントラストKDとの互換性を示す。

提案手法

• 学生ネットワークの最後の畳み込み層の後に、2層の線形特徴変換（FT）を挿入する。
• 最初のFT層は1x1畳み込みを用いてチャネルをN×C_tへ拡張し、F_seを生成する。二番目のFT層は別の1x1畳み込みでC_sへ縮約し、F_scを生成する。
• F_seをN個の非重複セグメントF_se^iに分割し、それぞれC_tチャネルを持つ。F_tに対してN parallel DISTILLATIONルートを実行する；L_normを最小化 = (1/N) Σ_i ||F_se^i − F_t||_2^2。
• FTを活性化なしの線形のままにし、推論時にはW_fc ← W_fc (W_sc W_se + I) によって次の全結合層へ統合する。
• 総訓練損失: L_total = L_ce + α L_norm（αはCIFAR-100でデフォルト10、ImageNetで8）; さらに性能を高めるためにL_kdおよび/またはL_crdを追加することも可能。
• アブレーションにはトレーニングを安定化させるための残差接続（線形）を含む。Nは通常8に設定。FTモジュールはアーキテクチャ変更を最小化するため最後の畳み込み層の後に配置。

実験結果

リサーチクエスチョン

• RQ1 多対一表現マッチングは、二段階KDにおける従来のOne-to-One特徴蒸留を上回ることができるか。
• RQ2 拡張係数NとL_normの重みαは、性能と訓練安定性にどう影響するか？
• RQ3 NORMアプローチは同種・異種の教師-学生ペアの両方に一般化するか、ロジットベースKDおよびコントラストKDとどう相互作用するか。
• RQ4 NORMの推論時の影響はどの程度か、FTを追加パラメータなしで分類器へ吸収できるか。

主な発見

• 同種の教師-学生ペアでのCIFAR-100において、NORMはベースラインより平均2.88%のTop-1精度向上を示す。
• 異種の教師-学生ペアでは平均5.81%、最大6.92%のTop-1精度向上。
• ImageNetでは、ResNet18とResNet34教師で72.14%トップ1（元70.13%）の絶対向上2.01%。
• ImageNetでは、MobileNetとResNet50教師で74.26%トップ1（元69.63%）の絶対向上4.63%。
• NORMはImageNetで主流のKD手法と比較してしばしば最善または競合的な結果を達成し、Vanilla KDや対比KDと組み合わせるとさらに改善できる（例: NORM+KD、NORM+CRD）。