Stuck on "A": Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization of a 0.6B Language Model

技術要約：Attention-to-KDA 線形化におけるインターフェース損傷の診断と修復

問題提起
本論文は、学習済みフルアテンション言語モデルを線形アテンションアーキテクチャ（特に Kimi Delta Attention: KDA メカニズム）へと変換する際の「情報の損失（lossy）」な性質を調査している。線形アテンションは、長文脈推論やエッジデバイスへのデプロイにおいて構造的な利点（O(n) の KV キャッシュに対し O(1) の再帰状態）を提供するが、標準的な指標が成功を示唆しているにもかかわらず、変換プロセスがタスク性能の維持に失敗することがしばーにある。

著者は、Qwen3-0.6B-Base（28レイヤー）を、21レイヤーを KDA に置き換え、7レイヤーは Grouped Query Attention (GQA) を維持するハイブリッドアーキテクチャへと変換することに焦点を当てている。これは、単一の 32 GB GPU と、数十億トークンではなく数千万トークンのトレーニング予算という、厳格なコンシューマーグレードの制約下で実行される。

核心となる診断的問いは、「変換は何を壊しているのか？」である。標準的な評価では、レイヤーごとのアライメントとエンドツーエンドの KL 蒸留を行った後、学生モデルは教師モデルに近いパープレキシティ（Cross-Entropy 差 +0.128 nats）を達成しているものの、多肢選択式（MCQ）の正解率においては崩壊している（C-Eval において教師の 50.6% に対し、ランダムに近い 25–29% に留まる）。

手法
著者は、この性能ギャップを診断し修復するためのマルチステージ・パイプラインを提案している。

1. ハイブリッドアーキテクチャとサージェリー（手術）:

   • レイアウト: レイヤー 3, 7, 11, 15, 19, 23, 27 はネイティブの GQA を保持し、残りの 21 レイヤーを KDA（16 ヘッド、ヘッド次元 128）に変換する。
   • KDA メカニズム: 忘却ゲート ($\alpha_t$) と書き込みゲート ($\beta_t$) を備えた、チャンネルごとのゲート付きデルタルールを実装。
   • 初期化: 重要なアブレーション研究により、単純な射影の移植（projection transplant）は失敗することが特定された。成功した構成（g6）は、再帰状態を安定させるための特定の減衰スケールとバイアス設定を使用している。

2. トレーニング・ステージ:

   • ステージ 2 (レイヤー単位のアライメント): 固定された教師モデルの隠れ状態と KDA の隠れ状態を MSE を用いて整列させる。これにより、Cross-Entropy (CE) は 9.48 から 4.13 に減少したが、C-Eval の精度は 25.5% のままだった。
   • ステージ 3a (エンドツーエンドの KL 蒸引): 教師モデルと学生モデルの分布間の前方 KL ダイバージェンスを用いて、事前学習テキストで学習。これにより、CE の差は +0.128 nats まで縮まったが、C-Eval の精度は 28.8% にしか上昇しなかった。
   • ステージ 3b (フォーマット特化型の修復): 教師によって検証された多肢選択問題（MCQ）、詩、および翻訳 QA を使用した、1,000 ステップの「補完のみ（completion-only）」の KL ステージ。決定的なのは、このステージが一般的なテキスト生成ではなく、回答の「フォーマット」に焦点を当てたことである。

3. 4つの置換による診断:
   パープレキシティを超えた診断を行うため、著者は 161 個のクリーンな MCQ（ベンチマークとの重複ゼロ）を用いた診断プロトコルを設計した。各質問は、回答の選択肢（A, B, C, D）の 4 つの循環的な回転の下で評価された。

   • 指標: 「スティッキネス（粘着性）」は、選択肢の内容の位置に関わらず、モデルがどれほど頻繁に同じラベル（例：「A」）を予測するかを測定する。
   • 知見: ステージ 3a のモデルは深刻な インターフェース損傷 を示していた。モデルは 81% の確率で「A」を予測し、4 つの回転すべてにおいて 161 問中 106 問で同じラベルを維持していた。モデルは知識を選択肢のラベルにマッピングする能力を失っており、選択肢の内容を事実上無視していた。

4. ペルソナ・アライメント:
   修復後、モデルには「Qingyi」という注入されたペルソナを用いた SFT と、モデルのアイデンティティや会話能力を劣化させないための 1 ラウンドのオンポリシー DPO が行われた。

主な結果

• 「パープレキシティの嘘」: 本研究は、隠れ状態のアライメントや低い KL ダイバージェンスが、タスク能力を保証するものではないことを裏付けている。モデルは教師の分布統計量を模倣できる一方で、その知識を特定のタスクに活用することには失敗する場合がある。
• インターフェース損傷の診断: 4つの置換テストにより、モデルが均一な知識喪失（ランダムな推測）に苦しんでいるのではなく、ラベルのスティッキネス に苦しんでいることが明らかになった。モデルは回答内容を追跡できず、特定の選択肢ラベルにデフォルト設定されていた。
• 修復の有効性: フォーマット特化型のステージ 3b 修復により：
  • C-Eval の精度が 28.83% から 41.31% に向上した。
  • 「A」の予測シェアが 81.06% から 58.5% に減少した。
  • 「4つの回転間で同じラベルになる」指標が半減した（106/161 $\to$ 47/161）。
  • 正解オプションの余剰（margin）が負（−0.117 nats）から正（+0.233 nats）に転じた。
  • トレードオフ: この修復は、一般的な分布適合性（Valid CE/KL）にわずかなコストを強いた。これは、モデルが MCQ フォーマットに特化したことを示唆している。
• 最終性能: ペルソナ・アライメント後、最終的なモデルは C-Eval で 41.83% を達成した。これは、事前 SFT 修復段階のノイズ範囲内に留まっており、壊滅的な忘却（catastrophic forgetting）が起きていないことを示している。

エンジニアリングの教訓と失敗モード
本論文は、プロセス中に遭遇したいくつかの重要なエンジニアリング上の落とし穴を記録している：

• FP32 マスター・フェイルモード: パラメータと Adam オプティマイザのモーメントの両方に bf16 を使用したため、小さな重みの更新が「静かに飲み込まれ」（sub-ULP 更新）、勾配ノルムは健全であるにもかかわらず重みが凍結してしまう現象が発生した。解決策として、FP32 マスターウェイトを使用した。
• キャッシュの不一致: 新しい KDA レイヤーが当初、増分キャッシュの更新を無視していたため、2 番目のトークンから生成がゴミへと崩壊した。特定のキャッシュ実装の修正により、再学習なしで解決した。
• 偽の移植: 教師から Q/K/V/O 射影を移植しようとする初期の試みは、機能的に失敗した。成功したアプローチは、射影の移植自体を完全に放棄することであった。

意義と主張
本論文は、主に 4 つの貢献を主張している：

1. 診断ツール: 「インターフェース損傷」（知識を出力フォーマットにマッピングできないこと）と「知識喪失」（均一な忘却）を分離する、4つの置換プロトコル。これは、あらゆる線形化または圧縮パイプラインに適用可能である。
2. 再現可能なレシピ: Qwen3 を KDA ハイブリッドに変換するための、予算制約のある完全なレシピ。ゲートの消失、キャッシュのバグ、更新の飲み込みといった、公開記録では省略されがちな具体的な失敗モードを文書化している。
3. ペルソナの生存: ペルソナ・アライメント（SFT + DPO）が、大規模な再学習や壊滅的な忘却を引き起こすことなく、アーキテクチャ変換を実質的に「無料で」生き残れるという証拠。
4. オープンリリース: コード、重み、レシピ、および失敗の監査証跡の完全な公開。

著者は、コンシューマー予算で 0.6B モデルを KDA ハイブリッドに変換することは可能であるが、標準的なメトリックススタック（パープレキシティ、KL、隠れ状態の MSE）は、損傷を系統的に誤報するという結論を下している。成功への主要な障壁は、必ずしも内部知識の喪失ではなく、その知識を特定の出力フォーマット（多肢選択ラベルなど）にマッピングするインターフェースの劣化であり、これはターゲットを絞ったフォーマット指導によって修復可能である。