[論文レビュー] Self-distillation with Batch Knowledge Ensembling Improves ImageNet Classification

Yixiao Ge, Xiao Zhang+7 more|arXiv (Cornell University)|Apr 27, 2021

Advanced Neural Network Applications参考文献 58被引用数 16

ひとこと要約

本稿では、ミニバッチ内に存在する他のサンプルからの知識を、サンプル間の類似度に基づいてオンラインでアンサンブルする自己蒸留手法であるBatch Knowledge Ensembling（BAKE）を提案する。この手法は、追加のパラメータが一切不要で、計算コストもわずか1.5%増加に抑えられながら、ImageNet分類性能を向上させ、Swin-Tのトップ-1精度を+0.7%向上させる。

ABSTRACT

The recent studies of knowledge distillation have discovered that ensembling the "dark knowledge" from multiple teachers or students contributes to creating better soft targets for training, but at the cost of significantly more computations and/or parameters. In this work, we present BAtch Knowledge Ensembling (BAKE) to produce refined soft targets for anchor images by propagating and ensembling the knowledge of the other samples in the same mini-batch. Specifically, for each sample of interest, the propagation of knowledge is weighted in accordance with the inter-sample affinities, which are estimated on-the-fly with the current network. The propagated knowledge can then be ensembled to form a better soft target for distillation. In this way, our BAKE framework achieves online knowledge ensembling across multiple samples with only a single network. It requires minimal computational and memory overhead compared to existing knowledge ensembling methods. Extensive experiments demonstrate that the lightweight yet effective BAKE consistently boosts the classification performance of various architectures on multiple datasets, e.g., a significant +0.7% gain of Swin-T on ImageNet with only +1.5% computational overhead and zero additional parameters. BAKE does not only improve the vanilla baselines, but also surpasses the single-network state-of-the-arts on all the benchmarks.

研究の動機と目的

• 教師信号としてのone-hotラベルの限界を是正し、教師信号の質を向上させることによる分類性能の向上を図ること。
• 知識蒸留の計算コストとメモリコストを低減しながら、性能を維持または向上させること。
• 1つのネットワークを用いて、ミニバッチ内のサンプル間でオンラインかつ軽量な知識アンサンブルを実現すること。
• サンプル間の知識伝達に基づく洗練されたソフトターゲットを用いることで、一般化性能とロバストネスを向上させること。

提案手法

• BAKEは、同じミニバッチ内の他のサンプルの予測を、それらのサンプル間の類似度に基づいて重み付けして即時に知識をアンサンブルする。
• サンプル間の類似度は、現在のネットワークの特徴量を用いて推定され、類似度の高いサンプルがより大きな寄与を示す。
• 知識伝達とアンサンブルは反復的に適用され、収束まで繰り返されるが、高コストを避けるために効率的な推論により近似される。
• 本手法は1つのネットワークで洗練されたソフトターゲットを生成し、自己蒸留を実現するため、複数の教師ネットワークや追加のブランチを必要としない。
• 本フレームワークはResNetやSwin Transformerを含むさまざまなアーキテクチャと互換性があり、sync BatchNormを用いた分散学習もサポートする。

実験結果

リサーチクエスチョン

• RQ1 パラメータの追加や計算コストの増加なしに、ミニバッチ内のサンプル間で知識アンサンブルを実施することで、蒸留性能を向上させられるか？
• RQ2 即時のサンプル間知識伝達が、自己蒸留におけるソフトターゲットの品質にどのように影響を与えるか？
• RQ3 SOTAの単一ネットワーク手法と比較して、BAKEはImageNetおよび下流タスクでどの程度の精度向上を達成できるか？
• RQ4 ラベルノイズや分布シフトにさらされたテスト時ロバストネスベンチマークにおいて、BAKEはどの程度のロバストネスを示すか？

主な発見

• BAKEは、トレーニング時の計算コストを3.7%増加させるのみで、ImageNetにおけるResNet-50の精度を+1.2%（76.8% → 78.0%）向上させた。
• BAKEは、計算コストをわずか+1.5%増加させ、追加パラメータを一切追加しない状態で、ImageNetにおけるSwin-Tの精度を+0.7%（81.3% → 82.0%）向上させた。
• BAKEは、CIFAR-100、TinyImageNet、CUB-200-2011、MIT67、Stanford Dogsを含む、すべての評価ベンチマークで単一ネットワークSOTA手法を上回った。
• 本手法は、COCO検出およびセグメンテーションタスクにおける転移学習性能を向上させ、学習済み表現の一般化能力が向上したことを示した。
• BAKEは、ImageNet-R、ImageNet-A、ImageNet-Pなどの摂動を加えたデータセット上でもテスト時ロバストネスを向上させ、モデルのロバストネスが向上したことを示した。
• アブレーションスタディの結果、近似的な推論による無限反復知識アンサンブルが、単一反復アンサンブルよりもよりロバストな性能を達成することが確認され、sync BatchNormは、クラス別サンプリングを伴う分散学習において不可欠であることがわかった。