ABKD: Pursuing a Proper Allocation of the Probability Mass in Knowledge Distillation via $α$-$β$-Divergence

要約

知識蒸留（KD）は、通常、フォワードカルバック – イブラー発散（FKLD）またはリバースKLD（RKLD）を使用して、出力分布間の発散を最小限に抑えることにより、大規模な教師モデルから小規模な学生モデルに知識を転送します。
1ホットのラベルと比較して、教師の分布によって提供されるより広範な監督情報のため、効果的なトレーニングパラダイムになりました。
KDのコアチャレンジは、2つのモード濃度効果のバランスをとることにあります。
勾配の更新中に確率がどのように再割り当てされるかを分析することにより、これら2つの効果がFKLDとRKLDで絡み合っているが、極端な形で絡み合っていることがわかります。
具体的には、両方ともFKLDが弱すぎるため、学生はターゲットクラスに集中できません。
対照的に、両方ともRKLDで強すぎるため、教師からのより広い分布情報を無視しながら、生徒はターゲットクラスを過度に強調します。
この不均衡に対処するために、$ \ alpha $ \ beta $ -divergenceを備えた一般的なフレームワークであるABKDを提案します。
私たちの理論的結果は、ABKDがFKLDとRKLDの間のスムーズな補間を提供し、これらの効果の間で効果的なトレードオフを達成することを示しています。
12の教師と学生の設定を備えた17の言語/ビジョンデータセットでの広範な実験は、その有効性を確認します。
このコードは、https：//github.com/ghwang-s/abkdで入手できます。

要約(オリジナル)

Knowledge Distillation (KD) transfers knowledge from a large teacher model to a smaller student model by minimizing the divergence between their output distributions, typically using forward Kullback-Leibler divergence (FKLD) or reverse KLD (RKLD). It has become an effective training paradigm due to the broader supervision information provided by the teacher distribution compared to one-hot labels. We identify that the core challenge in KD lies in balancing two mode-concentration effects: the \textbf{\textit{Hardness-Concentration}} effect, which refers to focusing on modes with large errors, and the \textbf{\textit{Confidence-Concentration}} effect, which refers to focusing on modes with high student confidence. Through an analysis of how probabilities are reassigned during gradient updates, we observe that these two effects are entangled in FKLD and RKLD, but in extreme forms. Specifically, both are too weak in FKLD, causing the student to fail to concentrate on the target class. In contrast, both are too strong in RKLD, causing the student to overly emphasize the target class while ignoring the broader distributional information from the teacher. To address this imbalance, we propose ABKD, a generic framework with $\alpha$-$\beta$-divergence. Our theoretical results show that ABKD offers a smooth interpolation between FKLD and RKLD, achieving an effective trade-off between these effects. Extensive experiments on 17 language/vision datasets with 12 teacher-student settings confirm its efficacy. The code is available at https://github.com/ghwang-s/abkd.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google