技術要約：RMSNorm Transformerにおける符号付き置換座標輸送

問題提起

現代の大規模言語モデル（LLM）のワークフローでは、ステアリングベクトル、スパースオートエンコーダ（SAE）の辞書、トップkニューロン集合、LoRAアダプタといった「座標インデックスを持つアーティファクト」を異なるモデルチェックポイント間で移動させる作業がますます重要になっています。しかし、「このニューロンは何であるか？」という問いは、固定されたゲージ（座標系）が存在しない限り、定義が不適切です。標準的なアライメント手法（Git Re-Basinや重みマッチングなど）は、置換のみの対称群（$S_d$）を想定していますが、RMSNormを利用するモデルに対してはこの仮定が正しくありません。

平均減算を含むLayerNormとは異なり、RMSNormは平均減算を行わないため、その残差ストリームは、置換に加えて各座標に対する独立した符号反転を許容します。したがって、RMSNormのネイティブな離散ゲージは、$S_d$ではなく、超八面体群 $B_d = S_d \ltimes \{\pm 1\}^d$（符号付き置換群）となります。この符号成分を無視すること（$S_d$のみのアライメント）は、座標インデックスを持つアーティファクトの不整合を引き起こし、転送時に重大な性能低下や完全な機能崩壊を招くという構造的な失敗モードをもたらします。

手法

1. ネイティブゲージの理論的特性化

本論文は、一般的なチャネルゲイン $\gamma$ を持つトランスフォーマーについて以下を確立しています：

• LayerNorm: $O(d)$ の最大座標保存部分群は $S_d \times C_2$（置換とグローバルな符号反転）である。
• RMSNorm: 最大座標保存部分群は、完全な符号付き置換群 $B_d$ である。
• 示唆: 連続的な直交回転（SliceGPTやQuaRotで使用されるもの）は関数保存には有効ですが、疎な座標の同一性を破壊します。疎なオブジェクト（特定のニューロンやSAE特徴量など）の同一性を保持するのは $B_d$ のみです。

2. 符号周辺化によるゲージ回復

ペアとなった活性化関数から正しい $B_d$ ゲージを回復するために、本論文は標準的なハンガリアン・マッチング・アルゴリズムの修正を提案しています：

• コスト関数: 符号付き相関 $\sum C_{i, \pi(i)}$ を最大化する代わりに、絶対相関 $\sum |C_{i, \pi(i)}|$ を最大化します。
• 符号の回復: 最適な置換 $\pi^*$ が絶対コスト行列を通じて見出された後、符号は $s_i = \text{sign}(C_{i, \pi^*(i)})$ として回復されます。
• 理論的保証: 非相関座標集団のケースにおいて、標準的な符号付き相関マッチングは、正しい負の相関を持つエントリーを能動的に回避してしまうため、約50%（正の符号の割合）の構造的な精度天井を持つことを本論文は証明しています。符号周辺化マッチングはこの障害を取り除き、100%の精度へと収束させます。

3. 軌跡合成による平行輸送

2つのエンドポイント間のみをマッチングするのではなく、本論文は学習軌跡に沿った平行輸送を導入しています：

• 隣接するチェックポイント間で、符号周辺化マッチングを用いてローカルな $B_d$ ゲージを推定します。
• これらのローカルなゲージを合成して、遠く離れたチェックポイント間（または共通のベースを共有する実行間）のグローバルな輸送マップを形成します。
• このアプローチは、直接的なエンドポイント・マッチングにおける「ホロノミー」の問題を回避し、隣接するチェックポイントの高い割り当てマージンを活用します。

4. ゲージ共変な輸送規則

本論文は、$B_d$ の下での様々な座標インデックス・アーティファクトの明示的な変換規則を導出しています（表1）：

• ベクトル（ステアリング、プローブ、SAEエンコーダ）: $v \mapsto vG$ として変換される。
• 行列（SAEデコーダ、LoRA WRITE）: $W \mapsto W G^\top$ として変換される。
• オプティマイザ状態（AdamW）: 第1モーメント $m_W$ は重みと共に変換され、第2モーメント $v_W$ は（二乗であるため）符号に依存せず、置換のみによって変換される。

主な結果

1. 回復精度

• 符号周辺化 vs 符号付き: 7B/8B RMSNormモデルを用いた実験において、素朴な符号付き相関マッチングは置換の約49.5%と符号の約50%しか回復できませんでした。一方、符号周辺化マッチングは**100%**の回復を達成しました。
• 輸送 vs エンドポイント・マッチング: Qwen2.5-1.5Bのファインチューニング・ベンチマーク（1500ステップ）において、ベースの合成による輸送は**91.1%の座標を回復しましたが、直接のエンドポイント・マッチングは60.3%**の回復に留まりました。この改善はベースのみを参照したことによるものではなく（ベース経由のエンドポイントは11.4%）、多くの高マージンなローカルゲージを合成したことによるものです。

2. アーティファクト転送性能

$B_d$ に適した輸送は、複数のアーティファクト・クラスにおいて $S_d$ のみの輸送を大幅に上回りました：

• LoRAアダプタ: Qwen2.5-1.5B (SST-2) において、$S_d$ 輸送はアダプタの精度を未適応のベースレベルまで崩壊させましたが（0.523 vs 0.525）、$B_d$ 輸送は利得の**99.7%**を保持しました（0.944）。
• SAE再構成: TinyLlamaにおいて、$S_d$ 輸送は再構成NMSEが1.08（失敗）となったのに対し、$B_d$ 輸送は0.004（ほぼ完璧）を達成しました。
• ステアリングベクトル: Qwenのセンチメント・ステアリングは、$B_d$ の下で**95.8%の効果を維持しましたが、$S_d$ の下では17.2%**でした。拒絶ステアリングベクトルは、$S_d$ の下では符号が反転し、実質的に挙動を逆転させてしまいました。

3. ステートフルな学習とマージ

• オプティマイザ状態: $S_d$ のみの輸送を行ったAdamW状態（置換は正しいが符号が+1に強制されている）を用いて学習を再開すると、参照点から機能的に同一のチェックポイントから開始しているにもかかわらず、学習の軌跡が乖離しました。$B_cd$ 輸送は、数値的な許容範囲内で軌跡を保持しました。
• モデルマージ: 座標保存型マージにおいて、$S_d$ のみの整列は破壊的な符号の不一致を導入し、高い「ピーク・アバブ・コード（chord）」障壁を生み出しました。符号付き整列はこれらの障壁を排除しました（ゲージ・スクランブル・テストにおいて0.00に減少）。

意義と主張

本論文は、RMSNormベースのアーキテクチャ（Llama 2、Llama 3、Qwenなどの広く使用されているモデルを含む）において、座標の同一性は符号付き置換までしかウェルディファインド（定義可能）ではないと主張しています。

• 再現性: 座標インデックスによる解釈可能性の主張（例：「ニューロンXは拒絶ニューロンである」）は、特定の $B_d$ ゲージが固定され適用されない限り、実行間で再現可能ではありません。本論文は、そのような主張を監査するためのネガティブコントロールとして「ゲージ・スイープ」を提案しています。
• ツールの転送: $S_d$ のみの手法の失敗は単なる軽微なエラーではなく、構造的な崩壊です。ステアリングベクトル、SAE辞書、およびオプティマイザ状態を正常に移動するには、符号付き輸送が必要です。
• 限界: 本論文は、$B_d$ が座標の同一性を保持する一方で、密な融合や圧縮において座標の同一性が主要な対象ではない場合には、連続的な直交手法が依然として有用であることを述べています。さらに、FFNの隠れユニットは、残差ストリームのゲージとは異なる独自の置換対称性（$S_{d_{ffn}}$）を持っています。

要約すると、本研究はRMSNormモデルを正しく整列させるための理論的基礎と実践的なアルゴリズムを提供しており、符号成分のゲージを無視することが、座標ベースのワークフローを根本的に破綻させることを実証しています。