技術概要：大規模言語モデルの継続的学習のための SAE-FD

問題提起

大規模言語モデル（LLM）は、複数のタスクに逐次的に適応する際、破滅的忘却に直面する課題を抱えています。既存の継続的学習（CL）手法は、通常、密な表現空間（重み、勾配、または出力空間）で動作します。これらのアプローチの根本的な限界は、特徴の重なり（feature superposition）にあり、これは密な活性化空間内の重なり合う次元に複数の概念が符号化される現象です。この絡み合いにより、以前に学習した知識を選択的に保護しながら、新たなタスクの学習を誤って制約してしまうことなく行うことが困難となり、結果として最適ではない可塑性と安定性のトレードオフが生じます。

手法：SAE-FD

著者は、事前学習されたスパースオートエンコーダ（SAE）のスパース特徴空間にモデル表現を固定する手法、**SAE-FD（スパースオートエンコーダ特徴蒸留）**を提案します。密な活性化をスパースで過剰基底（over-complete basis）に分解することで、SAE-FD は重なり合った表現の解離を図り、より標的を絞った正則化を可能にします。

この手法は、3 段階のパイプラインとして動作します。

1. SAE 学習（ステージ 1）

継続的学習のシーケンス開始前に、ベースモデルからの多様なテキスト活性化を用いてゲート付き SAEを学習します。

• 入力: 分類、QA、知識ベンチマークなど多様なソースからの最終層 MLP 活性化（$h \in \mathbb{R}^d$）。
• アーキテクチャ: ゲート付き SAE は、エンコーディングをゲーティング経路（$g$）と大きさ経路（$m$）に分離し、ReLU 前の特徴（$f_{pre}$）と ReLU 後の特徴（$f$）を生成します。
• 目的: 再構成損失（$\|\hat{h} - h\|_2^2$）と L1 スパース性ペナルティ（$\lambda_{\ell1}\|f\|_1$）を最小化します。
• 結果: 密な活性化をスパース特徴（$D \gg d$）に分解する凍結された SAE が得られ、概念を効果的に解離します。

2. タスク学習とアンカー捕捉（ステージ 2）

シーケンス内の各タスク $t$ に対して:

• ファインチューニング: 標準的な言語モデル化目的を用いて、LoRA を使用してモデルをファインチューニングします。
• アンカー捕捉: タスクの学習後、モデルは少量のアンカーサンプル（$N_a$）に対してフォワードパスを実行します。トークンごとの MLP 出力（$H_i$）を抽出して保存します。
• バッファ: これらのアンカーはタスクを超えて蓄積されます。その後の学習中、バッファからバッチが均一にサンプリングされ、蒸留損失を計算することで、コストをすべての過去のタスクにわたって均等化します。

3. スパース認識型特徴蒸留（ステージ 3）

新たなタスク（$t+1$）の学習中、モデルは凍結された SAEを通じて符号化された保存されたアンカーと比較することで正則化されます。蒸留損失（$L_{FD}$）は、2 つの相補的な項から構成されます。

• コサイン方向損失（$L_{cos}$）: 現在の ReLU 前特徴ベクトルとアンカー特徴ベクトルの間の角度のずれに対してペナルティを課します。これにより、表現のグローバルな方向が維持されます。
• アクティブ大きさ損失（$L_{mag}$）: アンカーにおいてアクティブであった（ReLU 後 $> 0$）特徴に対してのみ、平均二乗誤差（MSE）を適用します。これにより、以前のタスクに寄与する特徴の大きさを制約しつつ、非アクティブな特徴は新たなタスクのために採用されるために自由に保たれます。

適応的 $\lambda$ メカニズム

タスク学習と正則化のバランスを取るため、著者は総損失 $L = L_{task} + \lambda \cdot L_{FD}$ における重み付け係数 $\lambda$ に対する適応メカニズムを導入します。

• 固定された $\lambda$ の代わりに、システムは蒸留損失が総損失に占める固定された寄与率 $\rho$ を目標とします。
• $\lambda$ は、$\rho = \frac{\lambda L_{FD}}{\lambda L_{FD} + L_{task}}$ を満たすように指数移動平均（EMA）を通じて更新されます。
• 挙動: $\lambda$ は新たなタスクが開始された際（$L_{task}$ が大きく忘却のリスクが高い場合）に高く始まり、モデルが収束するにつれて緩和され、ファインチューニング中の過度な制約を防ぎます。

主要な貢献

1. SAE-FD フレームワーク: SAE のスパースで結合の解けた特徴空間で正則化を行う最初の継続的学習手法であり、密な空間に内在する特徴の重なり問題に対処します。
2. スパース認識型蒸留損失: コサイン方向の維持とアクティブな特徴に対する選択的な大きさの一致を組み合わせる新しい損失関数であり、SAE 表現のスパース性を活用して新たなタスクとの干渉を最小化します。
3. 適応的正則化: タスク損失と蒸留損失の相対的な大きさに基づいて正則化強度を動的に調整する適応的 $\lambda$ メカニズムにより、高リスク期間中の強力な保護を確保しつつ、収束を妨げません。

実験結果

著者は、3 つのモデルアーキテクチャ（LLaMA-2-7B、Vicuna-7B、Mistral-7B）にわたる 2 つのベンチマークで SAE-FD を評価しました。

• TRACE ベンチマーク（8 つの異種タスク）:
  • LLaMA-2-7B において、SAE-FD は**52.70% の平均精度（AA）を達成し、−0.46% の後方転送（BWT）**のみで済みました。
  • これは、最も強力な公開された正則化ベースのベースライン（GORP）を**+2.3% の AAで上回り、忘却を34%**削減（−0.46 vs −0.70 の BWT）しました。
  • BWT が−3% よりも悪い結果を示した他の LoRA ベースの手法（EWC、O-LoRA、TreeLoRA）を大幅に上回りました。
• 汎化性:
  • SAE-FD は、3 つのモデルファミリーすべてにおいて、一貫して最高の AA と BWT を達成しました。
  • Vicuna-7B において、O-LoRA を**+10.1% の AA**で上回り、フルファインチューニングのベースラインを含むすべての手法の中で最高の BWT（−0.69%）を達成しました。
• アブレーション研究:
  • 損失コンポーネント: コサイン方向損失とアクティブ大きさ損失の両方が必要であり、いずれか一方のみを使用すると AA が低下するか、忘却が増大しました。
  • 適応的 $\lambda$: 適応メカニズムを固定の $\lambda=1.0$ に置き換えると、BWT は−0.46% から−2.38% に悪化しました。
  • SAE 対 生 MSE: 密な活性化上の生 MSE 蒸留よりも、SAE 特徴空間での蒸留が優れており、結合の解けた空間の利点を確認しました。
• 効率性: 計算オーバーヘッドは SeqLoRA の約1.85 倍であり、ステップあたりのコストは過去のタスクの数に関係なく一定です。ストレージには、アンカー活性化のためにタスクあたり約 0.8 GB が必要です。

意義と主張

本論文は、SAE-FD が機械的解釈可能性（SAE 経由）と継続的学習を成功裡に結びつけたと主張しています。ニューラルネットワークの理解のために開発されたスパースで解離された表現を利用することで、この手法はそれらを制御するための効果的な帰納的バイアスを提供します。

著者は、スパース特徴空間で動作することで、知識のより標的を絞った維持が可能となり、密な空間正則化を悩ませる「特徴の重なり」問題が解決されると述べています。この手法は優れた可塑性と安定性のトレードオフを達成し、LLM が以前に学習したタスクで高い性能を維持しつつ、ゼロから再学習することなく進化の続くタスクに適応することを可能にします。

著者が指摘する限界:

• モデル固有の SAE の学習に、一度きりのトレーニングコストが必要である。
• アンカー活性化のストレージ要件はタスクの数に比例して増加する（ただし圧縮可能）。
• TRACE ベンチマークの実験は固定されたタスク順序で行われたが、複数モデルでの一貫性は堅牢性を示唆している。
• 実験設定の違い（例：フルファインチューニング対 LoRA、エポック数）により、いくつかのベースラインとの比較は慎重に解釈する必要がある。