技術要約: 有理型スパースオートエンコーダ (Rational Sparse Autoencoder: RSAE)

問題提起

スパースオートエンコーダ（SAE）は、大規模言語モデル（LLM）のメカニスティックな解釈可能性のための標準的なツールであり、内部活性化を過剰な特徴量辞書のスパースな線形結合へと分解するものである。しかし、既存のSAEファミリーは、固定された非平滑なエンコーダ非線形性、具体的にはReLU（$\ell_1$正則化を伴う）、TopK、およびJumpReLUに依存している。これらの固定された活性化プリミティブは、特定の病理的性質をもたらす：

• ReLU SAE: 活性化された特徴量の大きさの縮小（magnitude shrinkage）と、「死んだ」潜在変数の持続的な集団に苦しむ。
• TopK SAE: ソフトなペナルティをハードな基数制約に置き換えることで、不活性な特徴量への勾配の流れを遮断し、補助的な復活損失（revival loss）を必要とする。
• JumpReLU: 学習可能な閾値を導入するが、そのインジケーター・ゲートを通じたバックプロパゲーションのために連続緩和サロゲートを必要とする。

これらのハードコードされたメカニズムは、観測された事前活性化の幾何学的構造に適応するモデルの能力を制約し、再構成の忠実度とスパース性の間のトレードオフを歪める可能性がある。

手法

著者らは、固定されたエンコーダ活性化を学習可能な有理関数に置き換えた**有理型スパースオートエンコーダ（RSAE）**を提案する。RSAEは、標準的なSAEのスケルトン（アフィン事前活性化層、それに続くスパース活性化ブロック、および線形デコーダ）を保持しつつ、要素ごとに適用される学習可能な有理関数として活性化関数 $\phi(\cdot)$ を変更する。

アーキテクチャ

RSAEの活性化は次のように定義される：
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
ここで、$h$ は事前活性化ベクトルであり、$C_{in}$ と $C_{out}$ は学習可能なスケーリングパラメータである。また、$r_{(a,b)}(t) = P(t)/Q(t)$ は、学習可能な係数 $(a, b)$ を持つ有理関数である。スケーリングパラメータは、事前活性化の分布を有理関数の設計区間 $[-1, 1]$ にマッピングし、出力を期待される特徴量の大きさにマッピングする。

理論的基礎

本論文は、有理関数と現在使用されている区分的アフィン・ゲートとの間の表現上の非対称性を確立している：

1. 上界（有理近似）: Zolotarevの符号関数を用い、学習可能な有理活性化が、ターゲットとなる精度 $\epsilon$ に対して多項式対数サイズの規模で、ReLU、JumpReLU、および供給された閾値付きTopKゲートを近似できることを証明している（具体的には、マージン $\delta$ を持つ不連続ゲートに対して $O(\log(1/\epsilon)\log(1/\delta))$）。
2. 下界（区分的アフィン限界）: 逆に、特定の $O(1)$ パラメータを持つ有理ターゲットマップを、スカラー出力の単層区分的アフィンエンコーダ（ReLU/JumpReLU/TopK）を用いて近似する場合、$\Omega(\epsilon^{-1/2})$ 個の活性化座標が必要であることを示している。これは、有理型活性化がSAEエンコーダにとってより効率的な表現クラスを提供することを示唆している。

トレーニング・パイプライン

RSAEは、学習済みの「教師（teacher）」SAEをアップグレードするために設計された、2段階のパイプラインによって構築される。

1. 初期化:
   • レメズ交換（Remez Exchange）: $L_\infty$ 誤差を最小化するために、緩和されたレメズ交換アルゴリズムを用いて、合成的な高密度グリッド上で教師の活性化プリミティブ（ReLU、JumpReLU、またはTopK）に有理関数を適合させる。
   • キャリブレーション: 有理係数とスケーリングパラメータ（$C_{in}, C_{out}$）を、教師の実際の事前活性化分布に対する二乗誤差を最小化することで適応させる。
   • 重みの継承: エンコーダおよびデコーダの重みをそのままコピーする。
2. ファインチューニング: すべてのパラメータ（重み、バイアス、および有理係数）の凍結を解除し、標準的な $\ell_1$ 正則化再構成目的関数の下で共同最適化する。

主な貢献

1. RSAEアーキテクチャ: 浅いSAEエンコーダに特化した近似理論に基づいた、学習可能な有理活性化を持つスパースオートエンコーダの導入。
2. 理論的表現力のギャップ: 有理型活性化が固定されたSAEゲートのコンパクトな近似を提供する一方で、区分的アフィンエンコーダは単純な有理ターゲットを近似するために大幅に多くの活性化座標を必要とするという証明。
3. 2ステップ・トレーニングアルゴリズム: 実用的な初期化およびファインチューニング手順により、無視できるほどのオーバーヘッドで、既存の学習済みSAEをRSAEへと「ドロップイン」置換することを可能にする。
4. 実証的検証: 3つのオープンウェイト言語モデル（GPT-2 small、Pythia-160m、Gemma-2-2B）およびすべてのベースライン活性化ファミリーに対する包括的な評価。

結果

著者らは、残差ストリームの活性化に対して、学習済みベースラインと比較してRSAEを評価している。結果は、ファインチューニング後の厳格な改善を示している：

• 再構成指標: RSAEは、テストされたすべての構成において、教師モデルと比較して一貫して低い平均二乗誤差（MSE）、より低い $\ell_0$（スパース性）、および高い「生存」特徴量の割合を達成している。
• ダウンストリーム指標: RSAEはダウンストリームの挙動を改善し、残差ストリームを遮断した際のクロスエントロピーの劣化（$\Delta CE$）の減少と、損失回復率（$LR$）の向上を示している。
• 解釈可能性: Pythia-160mにおけるスパース・プロービング実験は、RSAEが特徴レベルの解釈可能性を保持していることを示している。8つのバイナリ分類タスクに対するフル辞書プローブの精度は、教師モデルと同等であり、大幅な低下は見られない。
• 効率性: このアップグレードは、オートエンコーダあたりわずかなスカラーパラメータしか追加しない。初期化と2,000ステップのファインチューニングを含む全パイプラインは、単一のコンシューマ向けGPU（NVIDIA RTX 5090）上で数分で実行可能である。
• 堅牢性: 改善は、異なるホストモデル、ベースラインの活性化ファミリー、および幅広いベースラインスパース性レベル（$\ell_0 \in [51, 683]$）にわたって一貫している。

意義と主張

本論文は、RSAEが、解釈可能性を犠牲にすることなく、一致したスパース性においてモデルの忠実度を厳密に向上させる、学習済みSAEのための「ドロップイン」修正であることを主張している。固定された非線形性を学習可能な有理関数に置き換えることで、RSAEはホストモデルにおける事前活性化の特定の幾何学的構造に適応し、ハードコードされた活性化プリミティブの限界を克服する。

著者らは、これらの改善が無視できるほどの計算コスト（単一GPUでの数分）と最小限のパラメータオーバーヘッドで達成されることを強調している。彼らはRSAEを、既存のSAEを強化し、メカニスティックな解釈可能性に必要な単意味的な特徴を維持しながら、再構成の質とダウンストリームのタスク性能の両方を向上させることができる、ユニバーサルなアップグレード戦略として位置づけている。本論文は、評価が3つのモデルと3つの活性化ファミリーをカバーしているものの、このアプローチは汎用的であり、他の最先端モデルやアーキテクチャのバリアント（例：Gated, BatchTopK, Matryoshka SAEs）にも拡張可能であると述べている。