Sparse Logit Sampling: Accelerating Knowledge Distillation in LLMs

要約

知識の蒸留は、教師の出力ロジットを事前に計算してキャッシュすることができる場合、大規模な言語モデルの知識を蒸留するための費用対効果の高い手法になる可能性があります。
ただし、これをトレーニング前に正常に適用することは、ほとんど説明されていません。
この作業では、直感的であるが、生徒に教師の確率分布の偏った推定値を提供し、最適ではないパフォーマンスとキャリブレーションをもたらすなど、キャッシングトップK確率などのまばらな知識蒸留のための素朴なアプローチが得られることを証明します。
重要なサンプリングベースの方法「ランダムサンプリング知識の蒸留」を提案します。これは、公平な推定を提供し、勾配を期待する勾配を保存し、大幅にスパースの高いロジットを保存する必要があります。
私たちの方法により、300mから3Bのモデルサイズの範囲にわたって、完全な蒸留と比較して競争力のあるパフォーマンスを維持しながら、エントロピークロスベースのトレーニングと比較して、わずかなオーバーヘッド（<10％）の学生モデルのより速いトレーニングが可能になります。

要約(オリジナル)

Knowledge distillation can be a cost-effective technique to distill knowledge in Large Language Models, if the teacher output logits can be pre-computed and cached. However, successfully applying this to pre-training remains largely unexplored. In this work, we prove that naive approaches for sparse knowledge distillation such as caching Top-K probabilities, while intuitive, provide biased estimates of teacher probability distribution to the student, resulting in suboptimal performance and calibration. We propose an importance-sampling-based method `Random Sampling Knowledge Distillation’, which provides unbiased estimates, preserves the gradient in expectation, and requires storing significantly sparser logits. Our method enables faster training of student models with marginal overhead (<10%) compared to cross-entropy based training, while maintaining competitive performance compared to full distillation, across a range of model sizes from 300M to 3B.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google