一文要約: Mambaは「Transformerは死んだ」ではありません — 二次的なAttentionボトルネックを言語モデリングの品質を失わずに回避できるという、初めての真剣な主張です。

第32章の概要: ポストTransformerアーキテクチャ — Mambaの選択的状態空間モデル、線形Attentionと再帰を持つRWKV、そしてAttentionとSSM層を混合して検索スタイルの忠実度と線形時間の長コンテキスト推論の両立を目指すハイブリッドスタック

32.1 Transformerのボトルネック

32.1.1 Attentionの複雑度問題

この本のほとんどはAttentionの研究に費やしてきました。数式はエレガントです:

Attention(Q,K,V)=softmax ⁣(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V

この数式の中に問題が潜んでいます。QKTQK^T の計算:

  • Qの形状は [N, d]
  • Kの形状は [N, d]
  • QKTQK^T の形状は [N, N]

すべてのトークンが他のすべてのトークンにAttentionします。計算もメモリも O(N²) でスケールします。

短いシーケンスでは問題ありません。長いシーケンスでは禁止的になります。

32.1.2 数値で見る

シーケンス長NAttention行列メモリ(FP16)
1,024100万要素2 MB
4,0961,600万要素32 MB
32,76810億要素2 GB
131,072 (128K)170億要素34 GB
1,048,576 (1M)1兆要素2 TB

4Kから128Kトークンへ(32倍長いシーケンス)は、Attention行列だけで 1,024倍のメモリ を必要とします。これがGPT-4が8Kコンテキストでローンチした理由であり、128Kへの拡張が設定フラグを切り替える以上の重要なエンジニアリングを必要とした理由です。

32.1.3 KVキャッシュが問題を複合する

第22章でKVキャッシュを扱いました: 生成中の再計算を避けるために過去のキーとバリューをキャッシュする技術。KVキャッシュは効率的な推論に必要ですが、シーケンス長に比例して増加します:

KVキャッシュサイズ = 2 × num_layers × num_heads × head_dim × seq_len × bytes_per_element

Llama 3-70Bの場合:

  • 4Kコンテキスト: 約2.5 GB
  • 128Kコンテキスト: 約80 GB
  • 1Mコンテキスト: 約625 GB(複数のサーバーが必要)

これはハードなトレードオフを作り出します: 長いコンテキストは同じハードウェア上の同時ユーザー数を減らすことを意味します。コンテキスト長の拡張は、Attention計算自体をより効率的にする技術があっても「無料」ではありません。

32.1.4 設計の問い

以下を満たすアーキテクチャはあるか:

  • 長距離依存関係のキャプチャ(Attentionが得意なこと)
  • O(N)時間での処理(O(N²)ではなく)
  • 推論時のO(1)状態(O(N) KVキャッシュではなく)
  • 言語モデリングでTransformerに匹敵する品質

State Space Models、特にMambaは今のところ最も信頼できる答えです。


32.2 State Space Models

32.2.1 RNNの文脈

Transformer以前、シーケンスモデリングにはRNNが使われていました。その核心的なアイデア: 履歴全体を固定サイズの隠れ状態に圧縮する。

h_t = f(h_{t-1}, x_t)    # 状態を更新
y_t = g(h_t)              # 出力を生成

RNNの利点: O(N)計算、O(1)推論メモリ(隠れ状態だけ)。

RNNの欠点: 順次計算(並列化できない)、勾配の消失/爆発(長距離依存関係のキャプチャが難しい)、すべての関連コンテキストを保持できない固定サイズのボトルネック。

Transformerは品質問題を解決しましたが、O(N²)スケーリング問題を再導入しました。SSMはRNNより良い品質でO(N)スケーリングを得ようとします。

32.2.2 状態空間方程式

SSMは制御理論と信号処理に起源があります。連続時間の定式化:

h(t)=Ah(t)+Bx(t)h'(t) = Ah(t) + Bx(t)
y(t)=Ch(t)+Dx(t)y(t) = Ch(t) + Dx(t)

ここで:

  • x(t)x(t) — 入力信号
  • h(t)h(t) — 隠れ状態(「メモリ」)
  • y(t)y(t) — 出力信号
  • A、B、C、D — 学習されるシステムパラメータ

パラメータには解釈可能な役割があります:

  • A: 隠れ状態が自律的にどのように進化するか(「忘却曲線」)
  • B: 入力がどのように隠れ状態に書き込まれるか(「新しい情報の重要性」)
  • C: 隠れ状態から出力をどのように読み取るか(「何を報告するか」)
  • D: 入力から出力への直接パス(ゼロに設定されることが多い)

32.2.3 離散化

コンピュータは連続信号ではなく離散シーケンスを扱います。ステップサイズ Δ\Delta で方程式を離散化する必要があります:

hk=Aˉhk1+Bˉxkh_k = \bar{A}h_{k-1} + \bar{B}x_k
yk=Chky_k = Ch_k

離散パラメータはゼロ次ホールド離散化から来ます:

Aˉ=eΔA\bar{A} = e^{\Delta A}
Bˉ=(eΔAI)A1B\bar{B} = (e^{\Delta A} - I)A^{-1}B

ステップサイズ Δ\Delta は各トークンが「表す」時間の量を制御します。Δ\Delta が大きいと、モデルは状態をより積極的に更新し、過去をより早く忘れます。

32.2.4 2つの計算モード

SSMには重要な構造的特性があります: 2つの等価な方法で計算できます。

再帰形式(推論用):

h = zeros(batch_size, d_state)
for k in range(seq_len):
    h = A_bar @ h + B_bar @ x[k]    # 状態更新
    y[k] = C @ h                     # 出力を読み取る

コスト: O(N)時間、O(1)メモリ(hだけ)。

畳み込み形式(学習用):

# 畳み込みカーネルを事前計算
K = [C @ B_bar,
     C @ A_bar @ B_bar,
     C @ A_bar^2 @ B_bar,
     ...]   # 長さN

# 畳み込みとして適用
y = conv1d(x, K)

コスト: FFT経由でO(N log N)、シーケンス全体で並列化可能。

SSMの優位性: 並列畳み込み形式で学習し、O(1)メモリの再帰形式で推論する。これがRNNにはできない方法で学習効率と推論効率を切り離します。

32.2.5 S4: 構造化状態空間モデル

2021年、StanfordのAlbert Guらが S4(Structured State Space Sequence Model)を提案しました。主要な貢献: 隠れ状態が長距離履歴を効率的に記憶できる、A(HiPPOと呼ばれる)の特定の構造化初期化を使用すること。

S4はコンテキスト長の制限によりTransformerが苦手とするタスクで、SSMが非常に長いシーケンス(数万要素)を扱えることを示しました。

しかしS4には致命的な制限がありました: パラメータが入力に依存しない。A、B、CマトリクスはどのトークンProcessingされているかにかかわらず同じです。モデルは一部のトークンにより多くのAttentionを向け、他には少なくするという決定を下せません。

これはクエリが動的に何にAttentionするかを選択するAttentionとは異なるメカニズムです。固定パラメータのSSMはすべての位置に固定の「Attention」を持ち、学習された忘却曲線によってのみ重み付けされます。


32.3 Mamba: 選択的状態空間

32.3.1 論文

2023年12月。Carnegie Mellon UniversityのAlbert GuとPrincetonのTri Daoが「Mamba: 選択的状態空間による線形時間シーケンスモデリング」を発表しました。主張: 言語モデリングタスクでTransformerの品質に匹敵するSSMアーキテクチャ。

これはその方向での最初の信頼できる結果でした。Mamba以前のSSMは特定のドメイン(音声、ゲノミクス、長シーケンス分類)では優れていましたが、言語でのTransformerには及びませんでした。

32.3.2 コアイノベーション: 入力依存の選択

Mambaはモデルの選択パラメータを入力の関数にします:

Δk=softplus(Linear(xk))\Delta_k = \text{softplus}(\text{Linear}(x_k))
Bk=Linear(xk)B_k = \text{Linear}(x_k)
Ck=Linear(xk)C_k = \text{Linear}(x_k)

ステップサイズ Δ\Delta は現在のトークンが何であるかに依存するようになりました。B(新しい情報の書き込み方)とC(何を読み出すか)も同様です。

AマトリクスはプログラムによるDesign選択として入力に依存しないままです。

入力依存の Δ\Delta が意味すること:

  • 大きな Δk\Delta_k: 状態が積極的に更新されます。モデルは効果的に古い状態を「忘れ」、現在のトークンに注目します。高解像度の観察のように振る舞います。
  • 小さな Δk\Delta_k: 状態はほとんど変わりません。モデルはメモリを保存し、現在のトークンを低重要性として扱います。多くのトークンを平滑化するように振る舞います。

モデルは重要なトークンに大きな Δ\Delta を割り当て、充填語に小さな Δ\Delta を割り当てることを学習します。これが選択です: モデルが何を記憶するかを決める。

32.3.3 選択とAttentionの比較

選択(Mamba)とAttention(Transformer)はどちらも、どの情報を伝えるかを決めるメカニズムです。メカニズムのレベルで異なります:

Attention:

各位置で、このトークンをすべての前のトークンと明示的に比較する。
バリューの加重和を取得する。
重みは(Q, K)の類似度に依存する。
コスト: O()  すべての位置対他のすべての位置。

選択的SSM(Mamba):

各位置で、隠れ状態をどれだけ更新するかを決める。
状態はすべての過去のトークンの圧縮サマリーを保持する。
決定(Δ、B、C)は現在のトークンに依存する。
コスト: O(N)  位置ごとに1つの状態更新。

MambaはAttentionのように過去のトークンに直接アクセスできません。圧縮された状態にアクセスします。これは表現力のトレードオフです — Mambaは位置ごとの表現力は低いですが、計算効率はずっと高い。

32.3.4 なぜ選択が言語に重要か

考えてみてください: 「エージェントはプルリクエストをマージしました、が3日前に提出したものを。」

「彼」を解決するには、多くのトークン前の「エージェント」に接続する必要があります。Attentionでは、これは明示的です: 「彼」は高いウェイトで「エージェント」にAttentionします。

非選択的なSSMでは、「エージェント」はその位置で状態に混入され、「彼」に到達するころには、間のすべてのトークンによって希釈されている可能性があります。モデルには「このエンティティは重要です、保存してください」というメカニズムがありません。

Mambaの選択メカニズムにより、モデルは効果的に「『エージェント』を見たとき、状態に目立つように書き込む(大きなB)。『the』やカンマを見たとき、最小限に更新する(小さなΔ)」と言えます。「彼」が解決を必要とするころには、「エージェント」はまだ状態でアクセス可能です。

32.3.5 ハードウェア対応アルゴリズム

パラメータを入力依存にすると、畳み込み計算が壊れます。カーネルが固定でないため、事前計算してFFTを適用できません。

Mambaの解決策: カスタムのハードウェア対応アルゴリズム。

重要な観察:

  1. 選択的スキャンはパラレルプレフィックスサムを使ってブロック内で並列化できる
  2. DRAM(遅いオフチップメモリ)ではなくSRAM(高速オンチップメモリ)にとどまれば、帯域幅コストは管理可能
  3. スキャン、A/B/Cの計算、出力射影を単一のCUDAカーネルに融合することで、中間のDRAM読み書きが不要になる

結果:

  • 学習スループット: 同等のシーケンス長でFlash Attention 2より約3倍速い
  • 推論: 位置ごとに一定のメモリでO(N)

32.3.6 Mambaブロックの構造

入力 x
  
  ├─────────────────────────────────────────┐
                                           
   線形射影(拡張)                         線形射影(拡張)
  Conv1d  SiLU活性化                     SiLU(ゲート)
  
  選択的SSM
  
  × (ゲートと要素ごとの乗算)
  
   線形射影(縮小)
  出力

主要な設計選択:

  • Attentionなし: 完全に削除
  • 拡張-縮小パターン: FFNのアップ射影/ダウン射影と類似
  • ゲーティング: 並列ブランチが乗算ゲートとして作用し、非線形性を提供

32.3.7 複雑度のまとめ

次元TransformerMamba
学習計算量O(N²d)O(Nd)
学習メモリO(N²) + O(N) KVO(N)
推論ステップごとO(N)(KVキャッシュ使用)O(1)
推論メモリO(N) KVキャッシュO(1) 状態

推論行が重要です。Transformerは生成するトークンごとにKVキャッシュが線形に増加します。Mambaの再帰状態はシーケンス長にかかわらず一定サイズです。


32.4 Mamba-2と状態空間の双対性

32.4.1 Mamba-1の限界

MambaはTransformerと同等の言語パープレキシティに達しましたが、分析によりAttentionが依然として優位を保つ領域が明らかになりました:

  • In-context learning: Transformerはfew-shot examplesをより効果的に使う
  • 正確な情報検索: 「第3段落で言及された正確な値は?」は状態ベースモデルには難しい
  • 複雑な多段階推論: AttentionでのPositionへの明示的アクセスが多段階推論を助ける

32.4.2 状態空間の双対性

2024年5月、DaoとGuがMamba-2論文を発表しました: 「TransformerはSSM: 構造化状態空間の双対性を通じた汎化モデルと効率的アルゴリズム」。

理論的な貢献: SSMとAttentionは同じ一般的フレームワークの2つのインスタンスです

両方とも構造化マトリクス乗算として書けます:

  • 密なマトリクス(構造なし)= フルAttention
  • 半分離可能マトリクス(特定の低ランク構造)= SSM

この「状態空間の双対性」(SSD)は2つのファミリーを数学的に結びつけます。それらは異なるパラダイムからの代替手段ではなく、シーケンス対シーケンス変換マトリクスにどの構造制約が課されるかで異なる、同じスペクトル上の点であることを示します。

実際的な結論: SSDフレームワークは中間点を示唆します。完全な密なAttentionでも完全な構造化SSMでもなく、その中間の何か。

32.4.3 Mamba-2の改善

理論を超えて、Mamba-2は実際的な恩恵をもたらします:

  • Mamba-1より学習スループットが 2〜8倍速い
  • より大きな状態次元: 16から256へ(長距離メモリが向上)
  • より良いテンソルコア利用率: 新しいアルゴリズムがGPUハードウェア構造にきれいに対応
タスクMamba-1Mamba-2Transformer
言語パープレキシティベースわずかに低いわずかに低い
長シーケンス(16K+)速い2〜8倍速い遅い
学習スループット高いより高い中程度

32.5 ハイブリッドアーキテクチャ: Jamba

32.5.1 動機

すべてのタスクで純粋なTransformerも純粋なMambaも明確に最良というわけではありません:

タスクタイプTransformerの優位性Mambaの優位性
短シーケンス、複雑な推論強い中程度
長シーケンス、スループットメモリボトルネック線形スケーリング
正確な情報検索強い弱い
トークンレベルルーティング強い強い
推論速度(長コンテキスト)遅い速い

明らかな問い: 組み合わせることはできないか?

32.5.2 Jambaアーキテクチャ

2024年3月、AI21 LabsがJambaをリリースしました — Transformer層、Mamba層、MoEを組み合わせた最初のプロダクションスケールモデル。

層構造(繰り返しパターン):

4層のグループ:
  層1: Mamba + MLP
  層2: Mamba + MLP
  層3: Mamba + MLP
  層4: Attention + MLP    4つのMamba層ごとに1つのAttention

一部の層のMLPはMoEで置き換えられます。

設計ロジック:

  • ほとんどの層は効率的なMambaを使用: O(N)計算、シーケンス処理の大部分を担う
  • 時折のAttention層: 必要なときに正確なグローバル情報検索を提供
  • MoE: 推論計算に比例せずに容量を増加

32.5.3 Jambaの設定

パラメータ備考
総パラメータ数52B
アクティブパラメータ数12BMoEがトークンごとの計算を削減
総層数32
Attention層数84層グループごとに1つ
Mamba層数24全層の4分の3
MoEエキスパート数16MoE層ごと
アクティブエキスパート数2Top-2ルーティング
コンテキストウィンドウ256K

1:3のAttention対Mambaの比率は恣意的ではありません。実証的に、これは高品質な検索タスクに十分なAttentionを提供しながら、ほとんどの計算を速いMamba層に保ちます。

32.5.4 なぜハイブリッドが純粋なアーキテクチャを上回れるか

長文書の処理:

[最初の8Kトークン]  [中間の240Kトークン]  [最後の8Kトークン]

Mamba層: 240Kの中間セクションを効率的に圧縮。
          状態ベースの圧縮: O()の爆発なし。

Attention層: 重要なセクション間の正確な検索。
              「最初に述べられた締め切りは?」

In-context learning:

[Few-shotの例]  [クエリ]

Mamba層: 例を圧縮された状態に処理。
Attention層: クエリ構造を例の構造に直接整合させる。

この組み合わせは、どちらの純粋なアーキテクチャも単独ではうまく処理できないタスクを処理します。

32.5.5 Jambaの性能

同クラスの比較:

モデルパラメータ数コンテキストVRAMスループット
Llama 2 70B70B4K140 GBベースライン
Mixtral 8x7B47B32K94 GB1.5×
Jamba52B(12Bアクティブ)256K100 GB

Jambaは4Kを処理するLlama 70Bよりも低いVRAMで256Kコンテキストを処理します。Mamba層が長距離圧縮の重労働を担っています。


32.6 その他の代替アーキテクチャ

32.6.1 RWKV: 再帰を通じた線形Attention

RWKV(Receptance Weighted Key Value)は、再帰と互換性のある方法でAttentionを再定式化するアーキテクチャです。

コアの数式(簡略化):

wkv_t = sum(exp(k_i + w*(t-1-i)) * v_i for i in range(t))
y_t   = sigmoid(r_t) * wkv_t

RWKVは次の2つの方法で計算できます:

  • Transformerとして: 学習中に位置をまたいで並列
  • RNNとして: 推論中に順次、一定状態

状態は時間減衰されます: 位置iからの情報は ew(t1i)e^{w \cdot (t-1-i)} で重み付けされるため、古いトークンの寄与は少なくなります。これはMambaで Δ\Delta が忘却を制御する方法に似ていますが、入力依存ではなく固定の減衰です。

RWKVには強力なオープンソースコミュニティがあります(RWKV-4〜RWKV-6)。多くの言語ベンチマークでMambaと競合し、小型バリアントを通じてローカルデプロイをサポートします。

32.6.2 RetNet: リテンション

RetNet(Retentive Network、Microsoft Research、2023年)は「リテンション」メカニズムを使います:

Retention(Q,K,V)=(QKTD)V\text{Retention}(Q, K, V) = (QK^T \odot D)V

ここでDは減衰マトリクスです: i ≥ jの場合 Dij=γijD_{ij} = \gamma^{i-j}、それ以外は0。以前の位置からの情報は幾何学的に減衰されます。

RetNetは3つの計算モードを提供します:

  1. 並列: Transformerのように、学習に使用
  2. 再帰: RNNのように、推論に使用
  3. チャンク: ハイブリッド、非常に長いシーケンスに有用

減衰構造は固定(γ\gamma で制御)であり、Mambaの入力依存の選択より制約があります。言語ベンチマークでは、RetNetは小型のTransformerと競合しますが、大型のものには及びません。

32.6.3 Hyena: 暗黙的な長畳み込み

Hyena(Stanford、2023年)はAttentionを暗黙的な長畳み込みで置き換えます:

入力 x
  
xからの小さなネットワーク経由でN個のフィルター関数 f1, f2, ..., fn を生成
  
出力 y = (((x  f1) · g1)  f2) · g2 ...

フィルターは小さなMLPによって暗黙的に生成されるため、N×Nマトリクスを保存せずに任意に長くできます。コスト: FFT経由でO(N log N)。

Hyenaは長シーケンス分類タスク(PathX、16Kシーケンス)で強い結果を示しました。標準的な言語モデリングベンチマークでは、小型のTransformerと競合しますがMambaレベルの言語品質には達しません。

32.6.4 アーキテクチャの比較

アーキテクチャコアメカニズム計算量言語品質長シーケンス
Transformer密なAttentionO(N²)最良限定的
Mamba選択的SSMO(N)Transformerに近い優秀
RWKV線形再帰O(N)良い優秀
RetNet減衰リテンションO(N)良い優秀
Hyena暗黙的畳み込みO(N log N)中程度優秀
Jambaハイブリッド(Attn + Mamba)O(N)強い優秀

32.7 コード例

32.7.1 最小限のSSM

import torch
import torch.nn as nn

class SimpleSSM(nn.Module):
    """概念理解のための最小限のSSM。"""

    def __init__(self, d_model: int, d_state: int = 16):
        super().__init__()
        self.d_state = d_state
        self.A      = nn.Parameter(torch.randn(d_state, d_state) * 0.01)
        self.B      = nn.Parameter(torch.randn(d_state, d_model) * 0.01)
        self.C      = nn.Parameter(torch.randn(d_model, d_state) * 0.01)
        self.delta  = nn.Parameter(torch.ones(1) * 0.1)

    def discretize(self):
        """ゼロ次ホールド離散化。"""
        A_bar = torch.matrix_exp(self.delta * self.A)
        B_bar = self.delta * self.B   # 簡略化
        return A_bar, B_bar

    def forward(self, x):
        """x: (batch, seq_len, d_model) → (batch, seq_len, d_model)"""
        batch, seq_len, _ = x.shape
        A_bar, B_bar = self.discretize()

        h = torch.zeros(batch, self.d_state, device=x.device)
        outputs = []
        for t in range(seq_len):
            h = h @ A_bar.T + x[:, t, :] @ B_bar.T   # 状態更新
            y = h @ self.C.T                            # 出力を読み取る
            outputs.append(y)

        return torch.stack(outputs, dim=1)

32.7.2 選択的SSM(Mambaスタイル)

class SelectiveSSM(nn.Module):
    """Mambaスタイルの選択的SSM。重要点: B、C、Δが入力に依存する。"""

    def __init__(self, d_model: int, d_state: int = 16, d_conv: int = 4):
        super().__init__()
        self.d_state = d_state
        self.in_proj = nn.Linear(d_model, d_model * 2)
        self.conv    = nn.Conv1d(d_model, d_model, kernel_size=d_conv,
                                 padding=d_conv - 1, groups=d_model)

        # 入力依存パラメータ  Mambaのイノベーション
        self.B_proj     = nn.Linear(d_model, d_state)
        self.C_proj     = nn.Linear(d_model, d_state)
        self.delta_proj = nn.Linear(d_model, d_model)

        # Aは入力に依存しないが特殊な初期化を使う
        A = torch.arange(1, d_state + 1).float()
        self.A_log = nn.Parameter(torch.log(A))

        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch, seq_len, _ = x.shape

        xz   = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)

        x = x.transpose(1, 2)
        x = self.conv(x)[:, :, :seq_len]
        x = x.transpose(1, 2)
        x = torch.nn.functional.silu(x)

        # 入力依存パラメータを計算
        B     = self.B_proj(x)                                        # (B, L, d_state)
        C     = self.C_proj(x)                                        # (B, L, d_state)
        delta = torch.nn.functional.softplus(self.delta_proj(x))      # (B, L, d_model)
        A     = -torch.exp(self.A_log)                                 # (d_state,)

        y = self._selective_scan(x, A, B, C, delta)
        y = y * torch.nn.functional.silu(z)
        return self.out_proj(y)

    def _selective_scan(self, x, A, B, C, delta):
        batch, seq_len, d_model = x.shape
        h = torch.zeros(batch, d_model, self.d_state, device=x.device)
        outputs = []
        for t in range(seq_len):
            delta_t = delta[:, t, :].unsqueeze(-1)           # (B, D, 1)
            A_bar   = torch.exp(delta_t * A)                  # (B, D, d_state)
            B_bar   = delta_t * B[:, t, :].unsqueeze(1)       # (B, D, d_state)
            h       = h * A_bar + x[:, t, :].unsqueeze(-1) * B_bar
            y       = (h * C[:, t, :].unsqueeze(1)).sum(-1)   # (B, D)
            outputs.append(y)
        return torch.stack(outputs, dim=1)

32.7.3 スケーリング比較

def compare_scaling():
    """異なるシーケンス長でのウォールクロック時間を比較。"""
    import time
    d_model = 512
    batch   = 4

    seq_lengths = [512, 1024, 2048, 4096, 8192]
    print(f"{'seq_len':>8}  {'Attention':>10}  {'SSM':>10}  {'比率':>8}")
    print("-" * 50)

    for seq_len in seq_lengths:
        x = torch.randn(batch, seq_len, d_model).cuda()

        attn = nn.MultiheadAttention(d_model, 8).cuda()
        torch.cuda.synchronize()
        t0 = time.time()
        for _ in range(10):
            attn(x, x, x)
        torch.cuda.synchronize()
        attn_time = (time.time() - t0) / 10

        ssm = SelectiveSSM(d_model).cuda()
        torch.cuda.synchronize()
        t0 = time.time()
        for _ in range(10):
            ssm(x)
        torch.cuda.synchronize()
        ssm_time = (time.time() - t0) / 10

        ratio = attn_time / ssm_time
        print(f"{seq_len:>8}  {attn_time:>10.4f}s  {ssm_time:>10.4f}s  {ratio:>7.2f}×")

# 典型的な出力:
#  seq_len   Attention         SSM     比率
# --------------------------------------------------
#      512    0.0023s    0.0018s    1.28×
#     1024    0.0045s    0.0032s    1.41×
#     2048    0.0156s    0.0061s    2.56×
#     4096    0.0589s    0.0118s    4.99×
#     8192    0.2234s    0.0233s    9.58×

シーケンス長が増えるにつれてSSMの優位性が増します。8Kトークンで、SSMは約10倍速い。128Kトークンでは、優位性は桁で測られます。


32.8 アーキテクチャ選択ガイド

32.8.1 Transformer

選ぶ場面:

  • 4K以下のシーケンス(Attentionのオーバーヘッドは管理可能)
  • 正確な情報検索が必要なタスク(Attentionはこれを直接処理する)
  • In-context learningが中心(TransformerはFew-shotの例をより良く使う)
  • エコシステムの成熟度が重要(ライブラリ、ハードウェア最適化、事前学習済みウェイト)

32.8.2 Mamba / SSM

選ぶ場面:

  • 8K以上のシーケンス(二次コストが支配的になり始める)
  • レイテンシよりスループットが重要
  • KVキャッシュが非実用的なメモリ制約環境
  • 信号処理、音声、ゲノミクス(SSMはこれらのドメインに強い事前分布を持つ)

32.8.3 ハイブリッド(Jambaスタイル)

選ぶ場面:

  • 32K〜1M+トークンのコンテキスト長
  • 強い言語品質と長距離効率の両方が必要
  • 同じモデルで検索タスクと逐次処理の両方
  • より新しく、まだ実戦テストが少ないスタックを採用する意欲がある

32.8.4 フィールドのトレンド

近い将来(2025〜2026年):

  • Transformerは新しいプロジェクトのデフォルトとして残る
  • SSM層が長コンテキストのハイブリッドモデルに登場する
  • ハードウェアベンダーがMambaのスキャン操作を最適化する

中期(2026〜2028年):

  • ハイブリッドアーキテクチャがフロンティアシステムの新しい標準になる可能性がある
  • 長コンテキストが一般的なプロダクションサービングにSSMが採用される

長期:

  • アーキテクチャ探索がより良いシーケンスモデリングのプリミティブを発見するかもしれない
  • SSD理論はAttentionとSSMの両方が特殊ケースとして出現する統一フレームワークを示唆する

32.8.5 学術 vs 産業の視点

学術界 はSSMを理論的な対象として興奮しています。AttentionとSSMを構造化マトリクス理論で結びつけるSSDフレームワークは真の数学的洞察と見なされています。

産業界 は慎重です。Transformerスタックには何年もの最適化があります: Flash Attention、KVキャッシュ圧縮、量子化ツール、Speculative Decoding、TransformerをProduction環境で速くするすべてのインフラ。Mambaは独自のスタックが必要で、そのスタックはまだ成熟途上です。

現実的なパス: TransformerコンポーネントとSSMコンポーネントを混合するハイブリッドアーキテクチャにより、SSMのツールが追いついてくる間に両方の恩恵を受けられます。


32.9 章のまとめ

32.9.1 Transformerが持つ問題

O(N²) Attention: 計算もメモリもシーケンス長と共に二次的にスケールします。128Kコンテキストモデルは4Kコンテキストモデルの1,024倍のAttentionメモリが必要です。

KVキャッシュ: 効率的な生成を可能にするO(N)キャッシュは無限に増加し続けます。70Bモデルでは1Mトークン = 625 GB。

32.9.2 SSMが提供するもの

再帰状態: シーケンス長にかかわらずO(1)推論メモリ。隠れ状態はすべての過去のコンテキストを固定サイズのベクトルに圧縮します。

O(N)計算: ペアごとの比較ではなく、位置ごとに1つの状態更新。

学習の並列性: 学習中に等価な畳み込み計算、まだ並列化可能。

32.9.3 Mambaが加えるもの

選択: Δ\Delta、B、Cを入力依存にします。モデルはトークンごとに状態に何を書き込み、どのくらい積極的に更新するかを決定します。これがTransformerとの品質ギャップを縮める仕組みです。

ハードウェア対応アルゴリズム: 選択的スキャンを現代のGPU上で効率的に保つカスタムCUDAカーネル。

32.9.4 重要な数式

離散SSM:

hk=Aˉhk1+Bˉxk,yk=Chkh_k = \bar{A}h_{k-1} + \bar{B}x_k, \quad y_k = Ch_k

Mambaの選択:

Δk,Bk,Ck=f(xk)\Delta_k, B_k, C_k = f(x_k)

複雑度:

AttentionSSM
学習計算量O(N²d)O(Nd)
学習メモリO(N²)O(N)
推論メモリO(N) KVO(1) 状態

32.9.5 私の見解

Mambaは私が見た中で言語タスクでTransformerの支配に対して説得力のある主張をする最初のアーキテクチャです — ベンチマークで勝ることによってではなく(初期のMambaはほぼ同等でした)、線形時間の複雑度で同じ品質に到達できることを証明することによって。

SSDの洞察が最も理論的に興味深いです: AttentionとSSMは異なるパラダイムではなく、同じマトリクスフレームワーク上の異なる構造制約です。これはどのような特定のタスクに対しても正しいアーキテクチャは純粋なTransformerでも純粋なMambaでもなく、SSDスペクトル上のどこかの点であることを示唆します — これはまさにJambaのハイブリッドアプローチが実証的に発見していることです。

実践者へ: ハイブリッドアーキテクチャに注目してください。純粋なSSMのストーリーは理論的には魅力的ですがツールへの投資が不足しています。ハイブリッドのストーリーこそが最初にプロダクションの勝利が現れる場所です。


章のチェックリスト

この章を終えたら、次のことができるようになっているはずです:

  • AttentionのO(N²)コストがどこから来るか、長コンテキストでなぜ重要かを説明できる。
  • SSMの状態方程式と、A、B、Cがそれぞれ何を制御するかを説明できる。
  • SSMが2つの計算モード(再帰と畳み込み)を持つ理由と、それぞれの用途を説明できる。
  • Mambaの選択メカニズムが何をするか、入力依存のΔがどのように「選択」を実装するかを説明できる。
  • 状態空間の双対性の洞察と、TransformerとSSMの統一についての含意を説明できる。
  • Jambaのアーキテクチャ選択と、1:3のAttention対Mambaの比率が理にかなう理由を説明できる。
  • 与えられたシーケンス長とタスクに対して、Transformer、Mamba、ハイブリッドアーキテクチャを選択できる。

パート9を終えて

これでパート9が完結です。Transformerスタックのコアメカニズムからフロンティアまでをたどりました: アライメント学習、スパースエキスパートルーティング、拡張推論、そして最初の信頼できるポストTransformerの代替手段まで。

RLHF、MoE、テスト時計算スケーリング、そしてなぜMambaのO(N)が重要かを — これらの章を見ずに — 説明できるなら、付録と原論文を読む準備ができています。

この分野は速く動きます。これらの章は2026年4月のスナップショットです。動かないもの: 新しい論文を理解可能にする第一原理的な思考です。パート1〜8で構築した数学と直感が耐久性のある部分です。パート9のすべてはその応用です。

旅を共にしてくれたことに感謝します。引き続き学んでいきましょう。

このページを引用する
Zhang, Wayland (2026). 第32章: ポストTransformerアーキテクチャ — Mambaとハイブリッドモデル. In Transformer アーキテクチャ:直感から実装まで. https://waylandz.com/llm-transformer-book-ja/chapter-32-post-transformer-architectures
@incollection{zhang2026transformer_ja_chapter_32_post_transformer_architectures,
  author = {Zhang, Wayland},
  title = {第32章: ポストTransformerアーキテクチャ — Mambaとハイブリッドモデル},
  booktitle = {Transformer アーキテクチャ:直感から実装まで},
  year = {2026},
  url = {https://waylandz.com/llm-transformer-book-ja/chapter-32-post-transformer-architectures}
}