0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@GeneLab_999

Jambaってなんだ？ ── KVキャッシュを128GBから4GBに減らした設計

0
Last updated at Posted at 2026-08-25

Mamba-1やSSDの記事では、選択的スキャンとチャンク分解を自分の手で測って理解した。だが実務でMambaがどう使われているのかは、まだ実装を通して見ていなかった。この記事で考えてほしいのは、Attentionを完全に捨てるのではなく、どれだけ残せば十分なのかという配分の判断だ。手元の環境はコンシューマー向けハイエンド構成の一例（VRAM 32GB級のGPU）で、52Bパラメータのモデルはロードすらできない。だからこそ、KVキャッシュのサイズをPythonの計算式の上で再現し、AI21 Labsが公開した数値と照合するところから始めた。

この記事の対象読者

  • Transformer・Self-Attentionの仕組みをひととおり理解している人
  • Mamba-1やSSMがなぜ速いのかは知っていて、次は「実務でどう使われているか」を知りたい人
  • 長い文脈を扱おうとしてKVキャッシュでVRAMが尽きた経験がある人

この記事で得られること

Attention層とMamba層をどんな比率で混ぜ、KVキャッシュの節約とモデル品質のどちらを優先するかを、自分のユースケースに当てはめて判断できるようになります。

  • Jamba-v0.1の実コンフィグ値から、KVキャッシュのサイズを自分で計算できるようになります
  • Attention:Mamba比率が1:7に決まった根拠（実測結果）を説明できるようになります
  • MoEによって「activeパラメータ」と「totalパラメータ」が分かれる理由を説明できるようになります
  • 2026年時点のJamba系モデルのラインナップと、ライセンスの違いを把握できます

この記事で扱わないこと

  • Mamba自体の内部計算（選択的スキャンの実装等）の再解説（Mamba-1、SSMを参照してください）
  • Jambaのファインチューニング手順
  • 独自GPUクラスタでの事前学習の再現

1. TransformerとMambaが、別々に諦めていたもの

このセクションで分かること：TransformerのKVキャッシュ問題と、Mambaの文脈内学習（ICL）の弱点という、Jambaが埋めようとした2つの穴。

TransformerのSelf-Attentionは、それまでの全トークンを毎回参照する。参照するために、過去のK（キー）とV（バリュー）のベクトルを保存し続ける必要がある。この記憶領域がKVキャッシュで、文脈が伸びるほど線形に膨らむ。

一方のMambaは、状態を1つのベクトルに圧縮しながら系列を読み進めるSSMの一種で、過去をすべて保存し直さない。だから記憶領域は一定で済むが、代わりに苦手なことがある。

この記事では、Attention層を照合担当、Mamba層を書記と呼ぶことにする。書記は会話を聞きながら要点だけを1冊のノートに書き足していく。ノートは常に一定の厚さで済むが、書記は過去の発言を一字一句覚えているわけではない。照合担当は逆に、過去の発言をすべて原本のまま保管していて、必要なら何ページでも遡って確認できる。正確だが、原本の保管コストは会話が伸びるほど線形に増える。

AI21 Labsの論文は、書記だけに任せたときに何が起きるかを実際に測っている。IMDBのレビュー分類タスクでは、正解は「Positive」か「Negative」の2択のはずなのに、Mamba単体のモデルは「Very Good」「Funny」「3/10」のような答えを返すことがあった。原文の表現をそのまま引用せず、自分の言葉で要約してしまう。これがICL（文脈内学習）の弱さで、プロンプト中の例を見てその場で回答形式を真似る能力がAttentionより劣ることを示している。

モデル IMDB QuAC NarrativeQA
照合担当のみ（Attention） 84.1 27.9 45.8
書記のみ（Mamba） 48.8 20.2 27.7
書記＋照合担当（ハイブリッド） 90.9 26.6 43.7

（1.3Bパラメータモデル、250Bトークン学習。出典：Lieber et al. 2024, Table 6）

書記だけでは原本の引用ができず、照合担当だけでは保管コストが線形に増える。この対立を埋めるのがJambaのハイブリッド構造だ。

2. Jambaブロックの中身 ── 5つの自由度とKVキャッシュ4GBの根拠

前節で、書記だけでは正確な引用ができず、照合担当だけでは保管コストが線形に増えるという、両者の対立点を見た。では実際にAI21 Labsは、この2人をどんな配分でチームに編成したのか。このセクションで分かること：Jambaブロックを構成する5つの自由度と、それらが実際にどんな値に設定されているか。

Jambaは「Jambaブロック」という単位を繰り返して積み上げる構造を取る。1ブロックはl層で構成され、そのうちa層が照合担当（Attention）、m層が書記（Mamba）。さらにe層おきにMLPの一部が専門家パネル（MoE）に置き換わる。

Jamba-v0.1（2024年3月公開のベースモデル）では、これらの自由度は次の値を取る。

記号 読み たとえでの役割 型 Jamba-v0.1での値
l エル 1ブロックの人数（層数） 整数 8
a:m エイ・タイ・エム 照合担当a人に対する書記m人の配分 比率 1:7
e イー 専門家パネルを何層おきに招集するか 整数 2
n エヌ 専門家パネルの総人数 整数 16
K ケイ 1トークンごとに実際に相談する専門家の人数 整数 2

実際の層の並びは次のようになる（HuggingFace transformersのconfiguration_jamba.pyのデフォルト値、attn_layer_period=8, attn_layer_offset=4、expert_layer_period=2, expert_layer_offset=1から再構成）。

注目してほしいのは、専門家パネル（MoE）は主に書記（Mamba）側のMLPに組み込まれ、照合担当（層4）のMLPは素の状態のままという点だ。設定上のオフセットが重ならないためにこうなる。

KVキャッシュはなぜ4GBで済むのか

KVキャッシュのサイズは、照合担当の人数、つまりAttention層の数に比例する。

\text{KVキャッシュ（バイト）} = 2 \times L_{attn} \times H_{kv} \times D_{head} \times S \times B \times \text{bytes}

この式が言っているのは、KVキャッシュの大きさはAttention層の数 $L_{attn}$ に比例して増えるということだ。書記（Mamba）はKとVを保存しないので、この式には一切登場しない。$H_{kv}$はKV用のヘッド数、$D_{head}$は1ヘッドあたりの次元、$S$は文脈長、$B$はバッチサイズを表す。

この式に、Jamba-v0.1の実コンフィグ値（hidden_size=4096, num_attention_heads=32, num_key_value_heads=8）と、比較対象のLlama-2・Mixtralの値を入れて計算する。

def kv_cache_bytes(num_attn_layers, num_kv_heads, head_dim, seq_len, batch=1, dtype_bytes=2):
    """KVキャッシュの理論サイズ（バイト）を計算する。
    2 = KとVの2本分。Attention層の数だけKVキャッシュが必要になる。
    """
    return 2 * num_attn_layers * num_kv_heads * head_dim * seq_len * batch * dtype_bytes

def to_gb(b):
    return b / (1024 ** 3)

hidden_size, num_attention_heads = 4096, 32
head_dim = hidden_size // num_attention_heads  # 128
seq_len = 262144  # 256K context

models = {
    "Llama-2-7B（MHA、GQAなし）": dict(num_attn_layers=32, num_kv_heads=32),
    "Mixtral-8x7B（GQA、全層Attention）": dict(num_attn_layers=32, num_kv_heads=8),
    "Jamba-v0.1（GQA、a:m=1:7）": dict(num_attn_layers=4, num_kv_heads=8),
}

for name, cfg in models.items():
    b = kv_cache_bytes(cfg["num_attn_layers"], cfg["num_kv_heads"], head_dim, seq_len)
    print(f"{name}: Attn層数={cfg['num_attn_layers']}, KVキャッシュ={to_gb(b):.1f} GB")

実行結果は次のとおり。

モデル Attn層数 KVキャッシュ（256K文脈、16bit）
Llama-2-7B（MHA、GQAなし） 32 128.0 GB
Mixtral-8x7B（GQA、全層Attention） 32 32.0 GB
Jamba-v0.1（GQA、a:m=1:7） 4 4.0 GB

この3つの数値は、AI21 Labsの論文Table 1（Llama-2 128GB / Mixtral 32GB / Jamba 4GB）と完全に一致する。つまりKVキャッシュが128分の1になった理由は特別な圧縮技術ではなく、単純にAttention層の数を32から4に減らしたことにある。

MoEはどれだけ効くのか

もう1つの自由度がMoEだった。Jambaは12B活性パラメータ／52B総パラメータという構成を取るが、この差はMoEによって生まれる。総パラメータのうち実際に計算に使われる分だけが「active」で、残りは「毎回は呼ばれない専門家」として温存される。

構成 HellaSwag WinoGrande NQ
Jamba（MoEなし） 36.6 62.5 15.4
Jamba＋MoE 38.1 66.0 18.9

（7Bパラメータモデル、50Bトークン学習。出典：Lieber et al. 2024, Table 7）

MoEを足すと、計算量（active parameters）をほぼ変えずに品質が底上げされる。専門家パネルの総人数nを増やすほど総パラメータは増えるが、1トークンごとに相談する人数Kを変えなければ、計算コストは変わらない。

Mamba層は大規模化すると内部の活性化値が暴走しやすく、学習が途中でスパイクすることがある。JambaではRMSNormを内部に追加することで安定させている（論文6.4節）。

実務でロードする場合は、公式のHuggingFaceのtransformersライブラリを使うのが早い。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 8bit量子化でロードする例（mambaモジュールは量子化をスキップする）
quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_skip_modules=["mamba"],
)
model = AutoModelForCausalLM.from_pretrained(
    "ai21labs/Jamba-v0.1",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    quantization_config=quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")

3. なぜ1:7なのか ── Attentionは「8層に1層」で足りるという実測

前節で、Jamba-v0.1が実際にa:m=1:7という比率を採用していると確認した。だがこの数字はどこから来たのか。もっと減らせないのか、あるいはもっと増やす必要はないのか、を確かめないと、自分のユースケースで比率を変える判断ができない。このセクションで分かること：1:3と1:7を比較した実測結果と、比率ごとのKVキャッシュ節約幅を見積もる方法。

AI21 Labsは1.3Bパラメータのモデルで、a:mの比率を1:3と1:7で比較する実験をしている。

構成 HellaSwag WinoGrande NQ
照合担当のみ（Attention） 36.4 62.4 14.5
書記のみ（Mamba） 36.1 62.6 14.5
Jamba（1:3、MoEなし） 37.2 65.1 16.5
Jamba（1:7、MoEなし） 37.2 65.1 16.0

（1.3Bパラメータモデル、250Bトークン学習。出典：Lieber et al. 2024, Table 4）

1:3から1:7へ、照合担当をさらに半分近くまで減らしても性能がほぼ変わらないという表を最初に見たとき、てっきり誤植を疑った。だが論文の結論は明快で、「両者にほぼ性能差はないが、1:7のほうが計算効率が良い」ため、より大きなモデルでは1:7を採用したと書かれている。7Bパラメータでの比較でも、書記＋照合担当のハイブリッドは書記のみ・照合担当のみの両方を上回っている。

つまり、Attentionは全体の8分の1（12.5%）残すだけで、書記の弱点を解消できる（原本を正確に引用できない問題）というのがこの実測の核心だ。これは前節のIMDB分類タスクの結果とも整合していて、論文では「4層しかAttentionがないにもかかわらず」needle-in-a-haystack評価で高い精度を保っている、とも報告されている。

この関係を、自分のVRAM予算から比率を逆算する形で確かめてみる。

def kv_cache_bytes(num_attn_layers, num_kv_heads, head_dim, seq_len, batch=1, dtype_bytes=2):
    return 2 * num_attn_layers * num_kv_heads * head_dim * seq_len * batch * dtype_bytes

def to_gb(b):
    return b / (1024 ** 3)

hidden_size, num_attention_heads, num_key_value_heads = 4096, 32, 8
head_dim = hidden_size // num_attention_heads
seq_len = 262144
total_layers = 32  # Jamba-v0.1と同じ総層数に固定し、比率だけ変える

for m in [0, 1, 3, 7, 15, 31]:
    if m == 0:
        attn_layers, label = total_layers, "全層Attention"
    else:
        attn_layers, label = max(1, round(total_layers / (m + 1))), f"1:{m}"
    kv = kv_cache_bytes(attn_layers, num_key_value_heads, head_dim, seq_len)
    print(f"{label}: Attn層数={attn_layers}, KVキャッシュ={to_gb(kv):.1f} GB")
比率 Attn層数 KVキャッシュ（256K文脈）
全層Attention（純粋なTransformer相当） 32 32.0 GB
1:1 16 16.0 GB
1:3 8 8.0 GB
1:7（Jamba-v0.1の採用値） 4 4.0 GB
1:15 2 2.0 GB
1:31 1 1.0 GB

論文の実測（Table 4）が示すのは、1:3と1:7のあいだに性能差がほとんどないという事実だけで、1:15や1:31が同様に安全かどうかまでは検証されていない。したがって「品質を落とさずKVキャッシュを削れる範囲」は、少なくとも1:7までは実測の裏付けがあり、それより先は自分のタスクで検証が必要な領域、というのが現時点で言える境界線になる。

ここまでのまとめ

TransformerのAttentionは正確だがKVキャッシュが線形に膨らみ、MambaはKVキャッシュを持たない代わりに文脈内学習が弱い。Jambaはこの2つを、1ブロック8層のうち1層だけを照合担当（Attention）にするという配分で組み合わせている。この1:7という比率は、1:3との比較実験でほぼ性能差がないことを確かめたうえで、計算効率を優先して選ばれた。KVキャッシュはAttention層の数にだけ比例するので、この比率を変えることが、そのままVRAM予算のコントロールになる。

4. 2026年時点のJambaと、ライセンスの変更

ここまでは2024年3月に公開されたJamba-v0.1というベースモデルの設計判断を見てきた。だが実務でJambaの採用を検討するなら、現在提供されているモデルとライセンス条項を知らないと判断を誤る。このセクションで分かること：現行モデルのラインナップと、ライセンスがApache 2.0から変わった影響。

2026年1月8日時点の公式ドキュメントによると、現在提供されているJamba系モデルは次の3種類。

モデル パラメータ 最大文脈長 バージョン スナップショット
Jamba Large 398B（94B active） 256K 1.7 2025-07
Jamba Mini 52B（12B active） 256K 2 2026-01
Jamba 3B 3B 256K 2 2026-01

Jamba 3Bはオンデバイス用途を想定した小型モデルで、API経由ではなく重み配布での提供になっている。ラインナップはJamba-v0.1（2024）→1.5→1.6→1.7／Mini 2（2026-01）と継続的に更新されており、256Kコンテキストという特徴は一貫して維持されている。

ライセンスには重要な違いがある。最初のJamba-v0.1（ベースモデル）はApache 2.0で、事実上の無制限利用が可能だった。だがJamba 1.5以降（1.5／1.6／1.7／Mini 2／3B）はJamba Open Model Licenseという別のライセンスに切り替わっている。このライセンスは研究・商用利用ともに許可する permissive なものだが、利用者（またはその関連会社）の年間売上が5,000万米ドルを超える場合、商用目的での利用に制限がかかる条項がある。企業での採用を検討する場合は、この閾値を必ず確認したほうがいい。

トラブルシューティング

症状 原因 対処
Fast Mamba kernels are not available のようなエラー mamba-ssm / causal-conv1d が未インストール pip install mamba-ssm causal-conv1d を実行するか、use_mamba_kernels=Falseで低速フォールバックを使う
80GBクラスのGPUでもロード時にOOMになる bf16のままフル精度でロードしている BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["mamba"])で8bit量子化する
vLLMでの推論がエラーになる vLLMのバージョンが古い Jamba 1.5系はvLLM 0.5.4以上が必須（モデルにより異なるので公式カードを確認）
ライセンスが商用利用不可だと思い込む v0.1（Apache 2.0）と1.5以降（Jamba Open Model License）を混同している 年商5,000万ドルを超える場合のみ制限がかかる条項であることを確認する
位置エンコーディングを追加しようとしてしまう Transformerの感覚でRoPEが必要だと思い込む JambaはMamba層が暗黙の位置情報を与えるため、明示的な位置エンコーディングなしで動く設計（論文6.5節）

用語集

  • KVキャッシュ｜Attention層が過去の全トークンを参照するために保持するK・Vベクトルの記憶領域｜たとえでは「照合担当が積み上げる原本の山」
  • ICL（文脈内学習）｜追加の学習なしに、プロンプト中の例だけを見て応答の形式を真似る能力｜たとえでは「その場で他人の答え方を真似る」こと
  • SSM（状態空間モデル）｜状態を1つのベクトルに圧縮しながら系列を処理するモデルの総称｜たとえでは「書記が取る要約メモ」
  • MoE（専門家混合）｜複数の専門家ネットワークのうち一部だけを毎回選んで使う仕組み｜たとえでは「毎回違う専門家パネルを招集する」こと
  • GQA（グループ化クエリAttention）｜複数のクエリヘッドでKVヘッドを共有し、KVキャッシュを削減する手法
  • RMSNorm｜活性化値の大きさを揃える正規化手法。Jambaでは大規模学習時のMamba層の暴走を防ぐために使われている
  • active／totalパラメータ｜MoEにより生まれる「実際に計算に使うパラメータ数（active）」と「重みとして保持している総数（total）」の乖離

学習ロードマップ

  1. SSM → Mamba-1 → Mamba-2 → SSD：Mamba側の内部計算を追う
  2. Transformer → Self-Attention：Attention側の内部計算を追う
  3. GPU → VRAM：KVキャッシュ以外のメモリ消費要因を理解する
  4. HuggingFace：実際にモデルをロードして手を動かす

まとめ

Jambaを調べ始める前は、AttentionとMambaを混ぜるという発想を「両方のいいとこ取り」という曖昧な言葉でしか理解していなかった。だが実際の設計判断は、KVキャッシュという1つの指標に対して「照合担当を何人残すか」という配分問題に還元されていた。1:7という比率は思いつきではなく、1:3との比較で性能差がほぼないことを確かめたうえでの選択だったし、KVキャッシュが128GBから4GBに減った理由も、特別な圧縮ではなくAttention層を32から4に減らしただけという単純な事実だった。次に長文脈のRAGパイプラインでVRAM不足に直面したときは、Attentionをゼロにする以外にも、この配分という選択肢があることを思い出したい。

参考文献

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?