10月13日、amazon.scienceが「Can we change an AI model's architecture without retraining it from scratch?」と題した記事を公開した。この記事で紹介されている「Priming」手法の核心は、情報を非可逆圧縮するSSMが、完全な精度でトークンを保持するSWAより長文脈タスクで上回るという反直感的な発見と、学習済みTransformerのトークン予算をわずか0.5%未満使うだけでハイブリッドアーキテクチャへの変換が可能という点にある。スループットは最大2倍に改善しながら、元モデルと1%以内の精度差に収まる。
再学習なしにアーキテクチャを変える「Priming」
LLMの推論コストを左右するボトルネックのひとつが、Transformerのアテンション機構が抱えるKVキャッシュの問題だ。アテンション機構はこれまで処理したすべてのトークンのKey・Valueベクトルをキャッシュし続けるため、コンテキスト長が伸びるほどメモリ使用量が線形に増大する。
これを解消するアプローチとして、アテンション層とSSM（State Space Model）層を組み合わせたハイブリッドアーキテクチャが注目されている。SSMは過去の全情報を固定サイズの状態に圧縮するため、メモリフットプリントが一定に保たれる。ただし、こうしたハイブリッド構成を新たに試すには従来、数兆トークン・数十万GPU時間規模の事前学習が必要だった。
AWSの研究チームが発表したarXiv論文「Priming: Hybrid state space models from pre-trained Transformers」は、この問題を根本から回避する手法を提案する。元のTransformerの事前学習トークン予算の0.5%未満、すなわちトークン予算ベースで99%以上を削減した形で、既存の学習済みTransformerをハイブリッドモデルに変換できる。金銭・時間コストの削減はこのトークン予算削減に由来するが、具体的なGPU時間や費用の数値は元記事には記載されていない点に注意されたい。
Primingの仕組み：3ステップで変換する
Primingが成立する理論的根拠は、TransformerのアテンションもSSMも「過去のトークンが現在のトークンにどれだけ寄与するか」を決めるミキシング行列の積として統一的に表現できる点にある。この構造的な類似性を利用し、アテンション層のQuery・Key・Value射影行列をSSM層の入力射影として再利用する。
変換は以下の3段階で進む：
- 層の選択：各アテンション層をSliding Window Attention（SWA）で置き換えたときの精度劣化を指標に、ローカルな操作しか行っていない層を特定する
- 重み転送とアライメント：特定した層のアテンション重みでSSMを初期化し、短い蒸留プロセスで挙動を合わせる
- エンドツーエンド適応：次トークン予測でハイブリッドモデル全体をファインチューニングし、全層を協調適応させる
この手法はQwen、Llama、Mistralといった異なるTransformerファミリーや、DenseモデルとMoEモデルの両方で機能し、8Bから32B以上のスケールで検証済みだ。
同等の精度で、スループット最大2倍
研究の核となる成果が、32BパラメータのGated KalmaNet（GKA）ベースハイブリッドモデルの性能だ。
元のQwen3-32Bと同じ推論レシピで学習したGKA-Primed-HQwen3-32B-Reasonerを、数学コンテストベンチマーク「AIME 2025」で比較した結果：
- コンテキスト長128,000トークン時、スループット最大2.06倍
- 全問題セットで正答到達時間が1.2〜1.4倍速
- 最難問題では最大1.6倍速
- 思考トレースの長さはほぼ同等（全問題の中央値約14,000トークン）
スループット向上の主因は、アテンション層の半分をSSM層に置き換えることでKVキャッシュのサイズが約半分になり、同一ハードウェア上での並列実行数が増えることだ。トークンあたりのコストが下がるだけでなく、推論の質を維持したまま壁時計時間での正答到達が早くなっている点が重要である。
SSMの種類で性能差あり：GKA > GDN > Mamba-2
研究では3種類のSSMを比較している：
- Mamba-2：全状態に一様なスカラー減衰を適用するシンプルな実装。効率的だが識別能力は低い
- Gated DeltaNet（GDN）：デルタルールにより特定のキーに紐付いた値を選択的に上書きできる。Mamba-2より表現力が高い
- Gated KalmaNet（GKA）：消去方向が現在のキーだけでなく過去の全キーの統計に依存する。3者の中で最も表現力が高い
理論的な表現力のヒエラルキー（GKA > GDN > Mamba-2）は実験結果とも一致し、長文脈推論タスクでは差が広がり、短文脈タスクでは縮まるという傾向が確認された。32BスケールのハイブリッドGKAは、同じデータで学習した元Transformerと1%以内の精度差に収まっている。
「速度vs精度」をサービング時に調整できるGKA固有の特性
GKAが持つMamba-2やGDNにはない特性として、推論時の計算量を再学習なしに調整できる点がある。GKAは内部でカルマンフィルタに類似した反復的な回帰問題を解く構造になっており、各ステップで状態推定を繰り返すことで過去の文脈を精緻に圧縮する。このイテレーション数をサービング時に変更できるため、モデルを再学習することなく速度と精度のバランスを動的に操作できる。
論文では、通常30イテレーションで動作するGKAを10イテレーションに落とした場合、AIME 2025での精度が約1%低下する代わりにデコードが10%高速化することを示している。単一のモデルで速度と精度のトレードオフを動的に操作できるため、複数のモデルバリアントを維持せずに異なるレイテンシ要件に対応できる。
「SWAより圧縮の方が強い」という反直感的な結果
もうひとつ興味深い知見が、SSM層とSliding Window Attention（SWA）との比較だ。SWAは直近のトークンを完全な精度で保持するのに対し、SSMは過去を非可逆的に圧縮する。情報損失のないSWAが有利に思えるが、実験ではすべてのSSMベースハイブリッドがSWAハイブリッドを上回った。SWAのメモリ予算をSSMの2〜4倍に増やしても、長文脈ベンチマークではSSMが優勢だった。この結果は、固定サイズへの圧縮という一見不利な制約が、長文脈における情報の選択的保持という点で有利に働く可能性を示唆している。
モデルとコードの公開
学習コード、最適化カーネル、vLLMサービングプラグインとともに、8Bおよび32BスケールのPrimedハイブリッドモデル群がApache 2.0ライセンスで公開されている。モデルはHugging Faceのコレクションから入手できる。
詳細はCan we change an AI model's architecture without retraining it from scratch?を参照していただきたい。
