🤖 machine learning

Kronecker Factorization Improves Efficiency and Interpretability of Sparse Autoencoders

本論文は、潜在空間をヘッドへと分解し、構成的な共活性化を強制するために微分可能なAND型相互作用を採用することで、性能を損なうことなく解釈性を向上させ、特徴量の相関を捉え、計算コストを削減するスパースオートエンコーダ・アーキテクチャであるKronSAEを導入する。

原著者： Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

公開日 2026-09-01

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

大規模言語モデルは、テキストを数値のストリームへと変換することで人間の言語を処理する、膨大で複雑なシステムです。これらの数値は「アクティベーション（活性化）」として知られ、モデルの内部レイヤーを通じて流れ、単語の意味や文章の論理を運びます。これらの機械がどのように思考しているのかを理解しようとする研究者にとって、これらの隠れたストリームはブラックボックスです。これを開くために、科学者たちは「スパース・オートエンコーダ」と呼ばれるツールを使用します。このツールを、密で乱雑な数値の流れを取り込み、それを単純で明確な概念の長いリストへと分解する翻訳機だと考えてください。理想的には、このリストの各項目は、「数学」、「法的契約」、あるいは「青色」といった、特定の概念に遭遇したときにモデルが活性化させる、単一の明確なアイデアを表しています。このプロセスは「分解」と呼ばれ、研究者が機械の内部で回転している個々の歯車を見えるようにしてくれます。

しかし、標準的な翻訳機には限界があります。それらは、人間の言語が深く構造化されているという事実を無視し、あらゆる概念をリスト上の独立した平坦な項目として扱います。言葉やアイデアはしばしば予測可能なパターンで共に現れます。例えば、「地理」という概念は、「地図」や「方向」としばしば共起します。モデルがこれらのパターンを学習する際、標準的なツールは関係性を捉えることに苦労し、システムにその繋がりを暗黙的に学習させるか、最悪の場合、異なるアイデアを一つの混乱した特徴へと融合させてしまいます。これにより、結果として得られる概念のリストは解釈が難しくなり、計算効率も低下します。

T-Techの研究チームは、言語の自然な構造を最初から尊重する「KronSAE」と呼ばれる、新しいタイプの翻訳機の構築方法を提案しました。この設計は、内部の辞書をグループごとに整理します。各グループ内では、システムは既存の2つのより単純なコンポーネントを組み合わせることで、複雑な特徴を構築します。これは、シェフがベースとなる食材に特定のスパイスを組み合わせることで特定の料理を作るようなものです。その料理は、ベースとスパイスの両方が存在して初めて成立します。この新しいアーキテクチャでは、2つの基礎となる「親」信号が同時に活性化している場合にのみ、一つの特徴が活性化します。これにより、システムが単純な部分の組み合わせとして特徴を学習することを促す、組み込みのルールが作成されます。これは、言語が実際にどのように機能しているかを反映しています。

研究者たちは、Qwen2.5とGemma-2という2つの人気のある言語モデルを用いて、膨大な教育用ウェブページのデータセットでこのアプローチをテストしました。彼らは、この構造化されたアプローチが単に人間の言語を模倣しただけでなく、モデルの内部表現をより明確にしたことを発見しました。新システムを標準的な手法と比較した際、新しい設計はより特異的であり、互いに混同される可能性が低い特徴を生み出すことが分かりました。標準的な手法では、単一の特徴が多くのことを行おうとし、関連する複数の概念の意味を吸収して曖昧になってしまうことがよくあります。新しいシステムは、特徴を構築する際に強制的に特定の組み合わせから作ることで、この「吸収」を減少させました。その結果得られた特徴はより鋭くなり、より狭く精密なアイデアを記述するようになり、研究者にとって理解し、ラベル付けしやすくなりました。

明快さに加え、この新しい設計は大幅な効率向上をもたらしました。システムは単純なものを組み合わせることで複雑な特徴を構築するため、あらゆる可能性をゼロから計算する必要がありません。研究者たちは、元のデータを再構成する際の精度を維持しながら、エンコーダを実行するために必要な計算能力を40%以上削減できることを見出しました。これは、システムがより少ないリソースを使用して、同等の量の情報を学習できることを意味します。実験では、この大規模な計算コストの削減を行っても、パフォーマンスの低下は1%未満であり、これほど大きな速度と効率の向上に対しては無視できるほどのトレードオフでした。

また、本研究では、このシステムが概念間の隠れた関係をどの程度学習できるかについても調査しました。特徴間の関係があらかじめ分かっている合成データを用いた制御実験において、新しいシステムは標準的な手法よりもはるかに上手くこれらのパターンを回収することに成功しました。新しいシステムは、関連する概念を内部構造内の同じグループにまとめることを学習しましたが、標準的な手法はそれらを散漫に配置してしまいました。現実世界のデータを見たとき、研究者たちは、テキスト中で自然に共起する特徴が、新しいシステムにおいては実際に同じ内部グループにマッピングされていることを観察しました。このことは、アーキテクチャが言語の統計的な規則性を捉え、アイデアが実際にどのように接続されているかを反映した形でモデルの知識を整理していることを示唆しています。

研究者たちはまた、特徴自体の性質についても調査しました。彼らは、システム内のより単純な「親」コンポーネントは、多くの場合、複数の異なるアイデアを表現できる広範で抽象的なものであることを見出しました。しかし、これらの親が組み合わされると、結果としての特徴は非常に特異なものになります。例えば、「方向」に関連する広範なコンポーネントが、「医学用語」に関連する別のコンポーネントと組み合わさると、「医学的な指示」に関する特徴が作成されます。この階層構造により、システムは膨大な数の様々な概念を作るために、基本コンポーネントを再利用することが可能になりました。この組成的なアプローチは、特徴の解釈可能性を高めるだけでなく、知識を保存し検索するためのより効率的な方法も提供しました。

結局のところ、この研究は、私たちが扱うデータと同じくらい、解釈ツールをどのように設計するかも重要であることを示唆しています。言語における概念の組み合わせ方を模倣する単純な構造的バイアスを導入することで、研究者たちは、より効率的で透明性の高いシステムを作り上げました。これらの知見は、モデルがトレーニング中に偶然これらの関係を発見することを期待するのではなく、アーキテクチャ自体にそれらを組み込むことができることを示しています。このアプローチは、人工知能を理解するための新しい道を提示しており、複雑なシステムの理解を向上させるための構造化された学習への探求を促すべく、コードを公開してさらなる発展を呼びかけています。

技術要約: KronSAE

問題提起

疎な自己符号化器（Sparse Autoencoders: SAEs）は、高密度な言語モデルの活性化を、疎で解釈可能な特徴量へと分解するための標準的なツールである。しかし、標準的なSAEアーキテクチャは、通常、潜在的な辞書を独立した座標の平坦な集合として扱う。この「フラットな潜在空間」というデフォルト設定は、自然言語に見られる階層性、構成性、および体系的な共活性化といった、より高次の構造的関係を明示的にエンコードできていない。その結果、これらの関係性は、学習中に暗黙的に出現するか、あるいは事後解析を通じて回収される必要がある。さらに、標準的なエンコーダは、 $F$ 個の独立した潜在変数への高密度な射影を必要とするため、計算コストが $O(Fd)$ となり、辞書サイズが大きくなるにつれて高価になる。

手法: KronSAE

著者らは、潜在空間に階層的かつ構成的な事前分布を導入するエンコーダ側の修正案として、KronSAEを提案する。これは、入力を直接 $F$ 個の独立した特徴量へ射影するのではなく、潜在空間を $h$ 個の独立した「ヘッド」へと分解するものである。

アーキテクチャ

分解 (Factorization): 各ヘッド $k$ について、入力活性化 $x \in \mathbb{R}^d$ は、2つの低次元の「前潜在変数（pre-latents）」へと射影される：
- ベースベクトル $p^k \in \mathbb{R}^m$ （ $u^k = P^k x$ を介して）
- 拡張ベクトル $q^k \in \mathbb{R}^n$ （ $v^k = Q^k x$ を介して）
- ここで、 $m, n \ll d$ であり、総辞書サイズは $F = h \times m \times n$ である。
構成的相互作用 (mAND): ヘッドのポスト潜在変数（post-latent）は、微分可能な相互作用カーネルである mAND を用いて、ベースと拡張の前潜在変数のペアによる構成によって形成される：
$z^k_{i,j} = \text{mAND}(u^k_i, v^k_j) = \begin{cases} \sqrt{u^k_i v^k_j} & \text{if } u^k_i > 0 \text{ and } v^k_j > 0 \\ 0 & \text{otherwise} \end{cases}$
この演算は論理積（AND）を近似しており、ポスト潜在変数が活性化するためには、両方の親となる前潜在変数が活性化している必要があることを保証する。得られる行列 $z^k \in \mathbb{R}^{m \times n}$ は、フラット化されてポスト潜在ベクトルを形成する。
疎性 (Sparsity): すべてのヘッドからのフラット化されたベクトルが連結され、標準的な TopK 選択が適用され、最終的な疎な潜在ベクトル $f \in \mathbb{R}^F$ が生成される。

互換性と効率性

互換性: KronSAE は、TopK、Matryoshka（入れ子状の辞書）、Switch SAE（混合エキスパート）、および JumpReLU のような異なる活性化関数を含む、標準的な SAE の目的関数およびバリアントと互換性がある。
計算上の利点: エンコーダを $F$ 個の潜在変数への高密度な射影から、 $h(m+n)$ 個の前潜在変数への射影に置き換えることで、KronSAE はエンコーダのパラメータ数と浮動小数点演算数（FLOPs）を削減する。エンコーダのコストは、 $O(d \cdot F)$ ではなく $O(d \cdot h(m+n))$ にスケールする。

主な貢献

エンコーダ階層の特定: 本論文は、特徴量の構成に内在する構造的な帰納バイアスに基づき、SAE におけるエンコーダ側の階層性の欠如を、設計上の欠落した軸として特定した。
KronSAE の設計: 再利用可能な前潜在因子からポスト潜在特徴を構築する、mAND 相互作用を備えたクロネッカー分解エンコーダの導入。
解釈性の向上: 「特徴吸収（一つの特徴が別の特徴の厳密な部分集合となる現象）」を減少させ、標準的な TopK SAE よりも特異的で解釈しやすい潜在変数をもたらすことを示す経験的証拠。
再構成・計算トレードオフ: 説明分散（EV）の低下を 1% 未満に抑えつつ、エンコーダの FLOPs を 40% 以上削減できることを実証。

実験結果

著者らは、FineWeb-Edu データセットを用いて、Qwen2.5-1.5B および Gemma-2-2B モデルの残差ストリームに対して SAE を訓練した。

再構成忠実度: 同等のトークン予算（2億5千万トークン）の下で、KronSAE は、様々な辞書サイズ（ $2^{14}$ から $2^{16}$ ）およびレイヤーにおいて、標準的な TopK SAE と同等の説明分散（EV）を達成した。
解釈性: 自動化されたパイプライン（検出およびファジングスコア）を用いた結果、KronSAE のバリアント（特に $m \ge 2$ の場合）は、一貫して TopK ベースラインを上回った。これらの特徴量は、より低いトークンエントロピーとより低いマルチトークン比率を示しており、より高い特異性を示唆している。
相関学習:
- 実データ: 標準的な TopK SAE における相関する特徴量は、同一の KronSAE ヘッド内の特徴へとマップされる傾向があり、これはアーキテクチャが関連する概念をうまくグループ化できていることを示唆している。
- 合成データ: ブロック構造化された共分散を用いた制御されたトイ実験において、KronSAE は TopK SAE よりも有意に良く（高い RV 係数）、グラウンドトゥルースの相関構造を回収した。
特徴吸収 (Feature Absorption): KronSAE の構成は、ベースラインと比較して平均吸収率および完全吸収スコアが低く、特徴の「包含（subsumption）」が減少していることを示した。
効率性: パレート図の分析によれば、KronSAE は再構成品質を維持しながら、エンコーダの FLOPs を 40% 以上削減できる。

意義と主張

本論文は、KronSAE が、表現学習において疎性と分解がいかに共に機能し得るかについての新しい視点を提供すると主張している。アーキテクチャ的な共活性化の事前分布（mAND 相互作用による）を課すことで、再構成の忠実度を損なうことなく、構造化された特徴の相互作用を学習しやすくしている。

著者らは、KronSAE を以下の手法として位置づけている：

解釈性の強化: 特徴吸収を減少させ、より特異で単一意味的（monosemantic）な特徴を作成する。
効率性の向上: エンコーダの計算コストを劇的に削減することで、大規模な辞書を持つ SAE の訓練と実行をより現実的なものにする。
言語構造との整合性: 特徴を独立した平坦な座標として扱うのではなく、言語の構成的な性質（例：ある概念がベースとなるカテゴリと特定の修飾語の交わりであること）を明示的にモデル化する。

著者らは、制限事項として、評価が現在中規模の Transformer モデルに限定されていること、および分解されたエンコーダの実装に伴うエンジニアリング上のオーバーヘッドが中程度ではあるがゼロではないことを挙げている。また、トイモデルは帰納バイアスを検証しているものの、現実世界の相関構造はそれとは異なる可能性があることも認めている。