技術的サマリー：Budgeted LoRA

問題定義

大規模言語モデル（LLM）はスケーリングを通じて強力な性能を達成しますが、これには多大な展開コストが伴います。知識蒸留（KD）は、大規模な教師モデルから小規模な学生モデルへ知識を転移させることで解決策を提供しますが、既存のパイプラインは推論時に構造的に効率的なモデルを生み出すことに失敗することが多いです。

主要な限界は、LoRA（Low-Rank Adaptation）のようなパラメータ効率型微調整手法の広範な使用にあります。LoRA は、密なバックボーンを固定し低ランクの更新を追加することで学習可能なパラメータ数を削減しますが、密なバックボーン自体は変更されません。その結果、LoRA ベースの蒸留は学習コストを削減する学生モデルをもたらしますが、密な行列乗算がそのまま残るため、支配的な推論コストは削減されません。現在の手法は、推論中の学生アーキテクチャ内で計算がどのように割り当てられるかを明示的に制御することなく、パラメータ数や学習効率の最適化に焦点を当てることが多いです。

手法：Budgeted LoRA

著者は、モデル圧縮を構造化された計算割り当て問題として再定義するBudgeted LoRAという蒸留フレームワークを提案します。学生アーキテクチャを固定するのではなく、この手法は学習の終了時点で保持される密な計算の最終的な割合を決定するグローバルな計算予算を導入します。

このアプローチは 3 つの段階で動作します：

1. 予算感知ゲート付き LoRA モジュール:
   学生モデルは、選択された線形射影（例：アテンションの$q/k/v/o$および MLP の$up/gate/down$）をゲート付きモジュールに置き換えます。これらのモジュールは以下を計算します：
   $$y = d \cdot Wx + \frac{\alpha}{r_{max}} B((Ax) \odot g)$$
   ここで：

   • $W$は固定された密な重みです。
   • $d \in [0, 1]$は、予算スケジューラによって制御される密パス保持係数です。
   • $A, B$は LoRA 因子であり、$g$は学習されたランクごとのゲートを表します。
   • この定式化により、モデルは異なるモジュール間で密パスと低ランクパスの両方のアクティブ容量を動的に変化させることができます。

2. 予算スケジューリングとコントローラ:
   グローバルな予算パラメータ $F \in [0, 1]$ は、学習終了時点で保持する密な計算の目標割合を指定します。コントローラは以下によってこの予算を強制します：

   • コサインスケジュールを使用して、目標となる密コスト軌道 $b(t)$ を定義します。
   • アクティブな密な射影の MAC 数に基づいて、保持される総密コスト $C(t)$ を計算します。
   • 目標予算を満たすために、モジュールレベルの保持係数 ($d$) を貪欲に低下させ、まず低コストのモジュールの除去を優先します。
   • これにより、学習中に密パスから低ランクパスへの計算の制御された転移が実現されます。

3. 学習後圧縮:
   学習後、モデルは実際の推論節約を実現するための統合ステップを経ます。最終的な保持係数 $d$ に基づき、各モジュールは 3 つの展開形態のいずれかにマッピングされます：

   • ケース 1（密パスの削除）: $d < \epsilon_{zero}$ の場合、密パスは完全に除去され、剪定された低ランク成分のみが残ります。
   • ケース 2（SVD 近似）: $\epsilon_{zero} \le d < \epsilon_{lr}$ の場合、残りの密な寄与は低ランク SVD によって近似され、LoRA 因子と融合されます。
   • ケース 3（密パスの統合）: $d \ge \epsilon_{lr}$ の場合、密パスは保持され、LoRA 更新と統合されて標準的な密線形層となります。

このパイプラインは、単一の予算ダイヤル $F$ でパラメータ化された学生モデルのファミリーを生成し、効率性と性能の間の直接的なトレードオフを可能にします。

主な貢献

• 新しい定式化: 本論文は、固定アーキテクチャの最適化ではなく、グローバルな予算制約下での構造化された計算割り当て問題として蒸留を扱うことを提案します。
• Budgeted LoRA フレームワーク: コントローラによって強制される密保持と学習された低ランク割り当てを組み合わせ、学習後圧縮を経て展開効率の高い学生モデルを生成する手法を導入します。
• 制御可能なトレードオフ: 単一の予算パラメータが、学習コストと推論コストにわたるパレートフロンティアを形成する、制御可能な品質 - 効率トレードオフをもたらすことを実証します。
• ICL 保持: Budgeted LoRA が、標準的な LoRA ベースの蒸留よりも、関数スタイルのコンテキスト内学習（ICL）プローブにおいて著しく高い性能を保持することを示します。

結果

著者は、0.13B モデルから初期化された 6 層の学生モデルを用いて Budgeted LoRA を評価し、0.13B、0.5B、1B のさまざまなスケールの教師モデルから蒸留を行いました。

• 品質 - 効率トレードオフ:

  • 中程度の予算（$F=0.4$）において、Budgeted LoRA は標準的な KD-LoRA と同等のパープレキシティを達成しつつ、圧縮モジュール推論で1.74 倍の高速化を実現し、学習計算量を0.59 倍に削減しました。
  • 積極的な予算（$F=0.0$）において、この手法は圧縮モジュール推論で4.05 倍の高速化と0.38 倍の学習計算量を実現し、パープレキシティの低下は中程度にとどまりました。
  • この手法は明確なパレートフロンティアを描き、標準的な LoRA と比較して制御されたパープレキシティ低下を伴う低コストの運用ポイントを提供します。

• コンテキスト内学習（ICL）保持:

  • 関数スタイルの ICL プローブ（Function Vectors スイートから適応）において、Budgeted LoRA は、特に大規模な教師モデルから蒸留する際に、標準的な KD-LoRA を大幅に上回りました。
  • 例えば、1B 教師モデルの場合、KD-LoRA のマクロ平均 ICL スコアは 2.7 でしたが、Budgeted LoRA（$F=0.0$）は 5.0、（$F=0.4$）は 5.6 に達しました。
  • これは、高次な振る舞いを保持するには、単にパープレキシティを一致させるだけでなく、予算化された密保持を通じて計算構造を維持することが重要であることを示唆しています。

意義と主張

本論文は、計算制約のある蒸留において、振る舞いを保持することはパープレキシティの一致や単なるパラメータの削除よりも、密な計算が低ランクパスへどのように転移するかを制御することであると主張します。

著者は、標準的な LoRA が密なバックボーンをそのまま残すため、推論節約を実現できないと指摘します。対照的に、Budgeted LoRA は密パスと低ランクパスを競合するリソースとして扱います。結果は以下を示唆しています：

1. 構造化された割り当ての重要性: 密容量の段階的かつ予算化された削減は、急激な削除と比較して、低ランクパスへの能力転移をより滑らかにします。
2. パープレキシティを超えて: 高次の振る舞い（ICL など）は、最終的なパープレキシティが標準的な LoRA と似ていても、学習プロセスが明示的に計算の構造を管理する際に、よりよく保持されます。
3. 展開効率: この手法は、学習時の予算化シグナルを学習後圧縮を通じて実際の推論時の節約に変換することに成功し、密から完全な低ランクまでのモデルスペクトラムを提供します。

この研究は、効率性を単なるパラメータ数の問題ではなく、学習および蒸留プロセス中の計算割り当ての問題として捉えるための一歩として位置づけられています。