1つのモデルを深さごとに使える「Telescopic Language Model(TLM)」で訓練コストを減らす研究
この論文は、同じ言語モデルを複数の計算予算で使いたいという問題を扱います。通常は予算ごとに別の訓練や圧縮が必要です。研究者たちはTelescopic Language Model(TLM)と呼ぶ方法を提案し、同じモデルが層の深さを変えても有効に動くよう訓練しました。結果として、1回の訓練で20段階(prefix)すべてが実用的な性能を示し、品質と計算コストの全体的なトレードオフ(質−予算曲線下の面積)を43–44%削減できたと報告しています。
やり方は次の通りです。もとのモデルはTransformerという一般的な構造です。訓練では毎ステップ、モデルの「途中まで」(層を途中で切ったもの)をランダムに選びます。その切ったモデルに対して「次の単語を当てる」という通常の目的で学習させます。同時に、フルのモデルでも通常の学習を行います。この「確率的な接頭辞(prefix)監督」と、必ずフルモデルのパスも回す仕組み(論文中で“full anchor”と呼ばれる)により、訓練後の単一モデルは任意の深さで有効に動作します。学習時は各ステップで二回の順伝播・逆伝播を行いますが、推論(実際に使うとき)は特別な変更を必要としません。
実験は「200Mプロキシスイート」と呼ぶ条件で行われました。ここでは同じデータ流(FineWeb‑Eduの200億トークン)を使い、20段階の層ごとに性能を評価しています。TLMは20段階すべてで困難度を表す指標「パープレキシティ」(予測の当てやすさを示す値。小さいほど良い)と、その指標に敏感な下流タスクで妥当な性能を出しました。比較対象の「固定出口」方式(Matryoshka Language Model Suitesなど)は、監督を数点の深さだけに行うため、他の深さでは性能が極端に低く(パープレキシティが10^2〜10^5に達することがある)なりがちでした。TLMはフル容量では固定出口方式と同等の性能を保ちながら、全体のコストを単回の訓練で約12%低くできたと報告しています。
この手法が重要な理由は二つあります。一つは、同じモデルで多様な計算予算に対応できる点です。これにより予算ごとに別の訓練や圧縮を行う必要が減り、運用コストを下げられます。二つ目は、品質をどの深さに集中させるかが訓練時の「接頭辞のサンプリング密度」で調整可能になる点です。つまり、特定の深さに性能を集中させたいならそのように学習させられますし、連続的に使いたければ均等に学習させられます。これはモデルの設計(アーキテクチャ)ではなく訓練目標が弾力性(elasticity)を決める、という示唆につながります。
重要な注意点もあります。今回の結果は論文が示す「プロキシ」条件での評価に基づきます。これは実用的な最大規模の訓練とは異なるスケールでの検証を意味します。したがって、大規模モデルや別のデータセットにそのまま当てはまるかどうかは追加の検証が必要です。また、接頭辞のサンプリングを一部の深さに偏らせれば固定出口法と同等の性能を得られますが、その代わりに“連続性”は失われます。論文は訓練目標が鍵だと結論づけていますが、スケールや応用範囲については今後の研究で確かめる必要があります。