Jamba — Hybrid SSM-Transformer

상태 공간 모델(State Space Model; SSM)과 트랜스포머(Transformer)는 서로 다른 것을 원합니다. 트랜스포머는 어텐션(attention)을 통해 품질을 사들이지만 그 비용은 시퀀스 길이의 제곱(quadratic)에 비례합니다. SSM은 순환식(recurrence)을 통해 선형 시간 추론(linear-time inference)과 상수 메모리(constant memory)를 손에 넣지만 품질이 다소 뒤처집니다. AI21의 Jamba(2024년 3월)와 Jamba 1.5(2024년 8월)는 둘을 같은 모델 안에 함께 담았습니다. Mamba 레이어(layer) 7개마다 트랜스포머 레이어 1개, 한 칸씩 건너뛰며 적용되는 전문가 혼합(Mixture of Experts; MoE), 그리고 단일 80GB GPU에 들어가는 256k 토큰 길이의 문맥 창(context window)이 그 핵심입니다. Mamba-3(ICLR 2026)는 복소수(complex-valued) 상태 공간과 다중 입력 다중 출력(Multi-Input Multi-Output; MIMO) 사영(projection)을 도입해 SSM 쪽을 한층 더 조였습니다. 이 강의에서는 두 아키텍처(architecture)를 처음부터 끝까지(end-to-end) 읽어 가며, 순수 SSM이나 순수 트랜스포머의 장문 문맥(long-context) 시도가 버티지 못한 3년 동안 어째서 하이브리드(hybrid) 조리법만이 살아남았는지 설명합니다.

유형: Learn 언어: Python (표준 라이브러리, 레이어 비율 계산기) 선수 지식: Phase 10 · 14 (오픈 모델 아키텍처), Phase 10 · 17 (네이티브 희소 어텐션(Native Sparse Attention)) 예상 시간: 약 60분

학습 목표

• Jamba 블록(block)을 구성하는 세 가지 기본 요소(primitive)인 트랜스포머 레이어, Mamba 레이어, MoE와 1:7 비율 및 한 칸 건너뛰기(every-other) 인터리빙(interleaving) 조리법을 설명합니다.
• SSM의 순환식이 상위 수준(high level)에서 어떤 모양을 가지는지, 그리고 그것이 어떻게 상수 메모리 추론을 가능하게 하는지 말로 풀어냅니다.
• 256k 문맥에서 Jamba 모델의 키-값 캐시(KV cache) 크기(footprint)를 계산하고, 같은 조건에서 순수 트랜스포머 모델이 필요로 하는 양과 비교합니다.
• Mamba-3가 도입한 세 가지 혁신, 즉 지수-사다리꼴 이산화(exponential-trapezoidal discretization), 복소수 상태 갱신(complex-valued state update), MIMO를 이름으로 가리키고 각각이 겨냥하는 문제를 짚어냅니다.

문제

어텐션은 시퀀스 길이(sequence length)에 대해 제곱 복잡도를 가집니다. 상태 공간 모델은 선형입니다. 이 차이는 문맥이 길어질수록 누적됩니다. 256k 토큰에서 트랜스포머의 어텐션 맵(attention map)은 헤드(head)마다 약 650억(65B) 개의 항목(entry)을 가집니다. 반면 SSM의 순환 상태(recurrent state)는 시퀀스 길이와 무관하게 고정 크기입니다.

순수 SSM 모델(Mamba, Mamba-2)은 작은 규모에서는 트랜스포머의 펄플렉시티(perplexity)에 근접하지만, 상태 추적(state-tracking) 과제에서 뒤처지고 일부 인-컨텍스트 검색(in-context retrieval) 부류에서 실패합니다. 직관은 이렇습니다. SSM은 과거 이력(history)을 고정 크기 상태로 압축해 담는데, 이력이 길어지면 정보가 새어 나갑니다. 반면 어텐션은 모든 것을 정확하게 기억하지만 그 대가로 제곱 비용을 지불합니다.

명백한 해법은 둘 다 쓰는 것입니다. 정확한 회상(exact recall)이 중요한 자리에는 트랜스포머 레이어를 두고, 나머지에는 SSM 레이어를 둡니다. 그리고 그 비율(ratio)을 조정합니다. Jamba는 이 하이브리드 조리법을 운영 등급(production-grade) 규모에서 처음 출시한 모델입니다. 총 52B 파라미터(parameter), 활성(active) 12B, 문맥 256k, 단일 80GB GPU라는 구체적인 수치를 함께 내놓았습니다. Jamba 1.5는 이 계열(family)을 총 398B / 활성 94B까지 확장합니다. Mamba-3(ICLR 2026)는 하이브리드를 다시 만들 때 기준이 될 만한, 현재 가장 강력한 순수 SSM 기준선(baseline)입니다.

이 강의는 세 논문을 모두 읽고 "올바른 비율을 고르는" 사고 모형(mental model)을 만드는 것을 목표로 합니다.

개념

한 페이지로 보는 SSM

상태 공간 모델(SSM)은 고정 크기 상태 h를 이용해 시퀀스 x_1, ..., x_N을 처리합니다.

h_t = A h_{t-1} + B x_t
y_t = C h_t

각 시점(step)에서 상태는 선형 동역학(linear dynamics) A로 진화하고, 입력 B x_t를 받아들이며, 출력 C h_t를 내놓습니다. A, B, C는 학습 대상이 될 수 있습니다. 중요한 성질은 y_t를 계산할 때 직전 상태 h_{t-1}과 현재 입력 x_t만 있으면 된다는 것입니다. 더 이전 시점의 x는 직접 필요하지 않습니다. 메모리는 상수이며, 추론(inference)은 토큰당 O(1)입니다.

모델링 품질을 좌우하는 비결은 A의 구조입니다. S4(Gu 2021)는 학습 과정에서 긴 합성곱(long convolution)으로 효율적으로 평가할 수 있도록 매우 구조화된 행렬(matrix)을 사용했습니다. Mamba(Gu, Dao 2023)는 고정된 A, B, C를 데이터 의존적(data-dependent)인 값으로 교체했습니다. 이것이 바로 "선택적(selective)"이라는 표현의 정체입니다. Mamba-2(2024)는 구조를 한층 더 단순화했고, Mamba-3(2026)는 특정 지점에 복잡성을 다시 들여놓습니다.

핵심 성질은 디코더 LLM에서 SSM 레이어가 어텐션 레이어의 즉시 교체용(drop-in replacement) 부품이 될 수 있다는 것입니다. 다만 점점 커지는 KV 캐시 대신, 레이어별로 고정 크기를 가지는 상태가 들어섭니다.

Jamba 블록

Jamba 블록은 두 숫자에 따라 레이어를 번갈아 배치합니다.

• l: 어텐션 대 Mamba 비율(attention-to-Mamba ratio). Jamba는 l = 8을 씁니다. 즉 Mamba 레이어 7개마다 트랜스포머 레이어 1개가 들어가며, 7개의 Mamba와 1개의 어텐션이 그룹 하나에 8개 레이어를 이룹니다.
• e: MoE 적용 빈도(frequency). Jamba는 e = 2를 씁니다. 즉 한 칸씩 건너뛰며 MoE를 적용합니다.

한 블록 안의 레이어 배열은 다음과 같습니다.

M  M  M  M  M  M  M  A    (Mamba 7개 + Attention 1개)
|  M  |  M  |  M  |  M    (|로 표시된 자리에 MoE 적용)

각 Jamba 블록은 8개의 레이어로 구성됩니다. 블록 4개 깊이, 즉 총 32개 레이어가 쌓이면 Mamba 레이어 28개와 어텐션 레이어 4개를 얻고, 그중 16개 레이어는 MoE를 사용합니다.

왜 1:7 비율인가

AI21은 절제(ablation) 실험을 수행했습니다. 어텐션 대 Mamba 비율이 자신들의 장문 문맥 평가(long-context eval)에서 파라미터당 펄플렉시티(perplexity-per-parameter)와 인-컨텍스트 회상(in-context recall)을 동시에 가장 좋게 만드는 지점이 어디인지 찾는 작업이었습니다.

• 어텐션이 너무 많은 경우(1:1): 품질은 올라가지만 메모리와 속도가 나빠집니다.
• 어텐션이 너무 적은 경우(1:15): 메모리는 훌륭하지만 인-컨텍스트 검색이 실패합니다.
• 가장 좋은 지점(sweet spot): 1:7 또는 1:8.

직관은 이렇습니다. 트랜스포머 레이어는 정확한 회상과 상태 추적을 담당하고, Mamba 레이어는 값싸게 대량의 처리(cheap bulk processing)를 담당합니다.

위치 인코딩(Positional encoding)

Mamba 레이어는 순환식 자체가 위치를 의식(position-aware)하도록 만들어져 있습니다. 원래의 Mamba 기반 하이브리드에서 어텐션 레이어는 회전 위치 인코딩(Rotary Position Embedding; RoPE)을 쓰지 않았습니다. SSM 레이어가 위치 정보를 제공하기 때문입니다. Jamba 1.5는 더 긴 문맥으로의 일반화를 위해 어텐션 레이어에 RoPE를 추가했습니다. 이는 경험적인 장문 문맥 평가에 기반한 사후(post-hoc) 보정입니다.

메모리 예산(memory budget)

Jamba-1의 형태(32 레이어: 28 Mamba + 4 Attention, 은닉 차원(hidden) 4096, 어텐션 헤드 32)를 살펴봅니다.

• KV 캐시(어텐션 레이어 한정): 256k에서 BF16 기준 2 * 4 * 32 * 128 * 256k * 2 = 8.4 GB. 어텐션 레이어 4개만 기여합니다.
• SSM 상태: 토큰 접두(prefix)당 28 * hidden * state_size처럼 보이지만, 시퀀스 길이에 따라 커지지 않고 레이어별로 고정 크기를 유지합니다. 전형적인 Mamba 상태가 특징(feature)당 16이고 은닉 차원이 4096이라면 28 * 4096 * 16 * 2 = 3.7 MB로 전체가 끝납니다.

같은 은닉 차원과 32 레이어를 갖는 순수 트랜스포머가 32개 헤드의 풀(full) 멀티 헤드 어텐션(Multi-Head Attention; MHA)을 쓴다면 256k BF16에서 2 * 32 * 32 * 128 * 256k * 2 = 128 GB가 됩니다. KV 캐시가 8배 줄어드는 셈입니다. 2024년 대부분의 모델이 채택한 그룹 쿼리 어텐션(Grouped-Query Attention; GQA) 8 기준선과 비교해도 2 * 32 * 8 * 128 * 256k * 2 = 32 GB이며, Jamba의 1:7 하이브리드는 16 GB로 여전히 2배 더 작습니다.

이것이 AI21이 말하는 "단일 80GB GPU에서 256k 문맥"의 구체적인 의미입니다. 풀 MHA 순수 트랜스포머의 KV 캐시는 들어가지 않습니다. GQA 기준선조차 가중치(weights)와 활성값(activations)을 위한 공간을 거의 남기지 않습니다. Jamba는 들어갑니다.

Mamba-3: 2026년의 순수 SSM 기준선

Mamba-3(ICLR 2026, arXiv:2603.15569)는 순수 SSM 쪽에 세 가지 혁신을 도입합니다.

1. 지수-사다리꼴 이산화(Exponential-trapezoidal discretization). Mamba-2가 쓰던 오일러 방법(Euler-method) 이산화를 더 표현력 있는 순환식으로 교체합니다. x_t에 바깥쪽으로 합성곱을 적용하는 대신, 핵심 순환식 안에서 상태-입력에 합성곱과 비슷한 연산(convolution-like operation)을 적용합니다.

2. 복소수 상태 갱신(Complex-valued state update). 이전 세대의 Mamba들은 상태 행렬을 복소수(S4)에서 실수 대각(real diagonal, Mamba), 그리고 스칼라 항등(scaled identity, Mamba-2)으로까지 단순화해 왔습니다. Mamba-3는 다시 복소수 값을 도입합니다. 이는 상태에 데이터 의존적인 회전 임베딩(rotary embedding)을 곱하는 것과 동등합니다. 이전의 실수 단순화가 잃어버렸던 상태 추적 능력을 되살리는 효과를 냅니다.

3. 다중 입력 다중 출력 사영(Multi-Input Multi-Output; MIMO). 특징별 스칼라 사영 대신 행렬값(matrix-valued) 사영을 사용합니다. 디코드 지연(decode latency)을 늘리지 않으면서도 모델링 능력과 추론 시점의 하드웨어 활용도(hardware utilization)를 함께 끌어올립니다.

1.5B 파라미터 규모에서 Mamba-3는 Gated DeltaNet 대비 다운스트림 평균 정확도를 0.6 포인트 끌어올립니다. MIMO 변형은 거기에 1.2 포인트를 더해 총 1.8 포인트의 개선을 냅니다. 같은 상태 크기에서 Mamba-3는 Mamba-2를 절반의 상태로 따라잡습니다.

Mamba-3는 아직 운영 환경의 하이브리드로 대규모 출시되지는 않았습니다. 그러나 다음 Jamba 급(class) 모델의 SSM 쪽 후보로는 가장 자연스러운 카드입니다.

언제 하이브리드를 선택할 것인가

하이브리드가 이기는 경우는 다음과 같습니다.

• 문맥이 충분히 길어서 순수 트랜스포머의 KV 캐시가 부담스러워지는 경우(64k 이상).
• 과제가 짧은 거리(SSM에 유리)와 긴 거리(트랜스포머가 필요)의 구조를 함께 다루는 경우.
• 단일 GPU 메모리 예산으로 배포해야 하는데, 트랜스포머 KV 캐시만으로 이미 자리가 모자라는 경우.

하이브리드가 지는 경우도 있습니다.

• 문맥이 짧은 경우(16k 미만). SSM에 따르는 오버헤드(overhead)가 낭비되며, 이때는 순수 트랜스포머로 충분합니다.
• 모든 위치에서 모든 위치로의 어텐션(everywhere-to-everywhere attention)이 필요한 과제(깊은 추론(deep reasoning), 여러 문서 간 교차 참조). 하이브리드에서 어텐션 레이어가 드문드문 배치되어 있다는 점이 불리하게 작용합니다.
• 조 단위(trillion-parameter) 최전선 모델로 규모를 키우는 경우. 현재의 능력 경쟁에서는 순수 트랜스포머 + 다중 잠재 어텐션(Multi-head Latent Attention; MLA) + MoE 조합(DeepSeek-V3 방식)이 앞서고 있습니다.

경쟁 지형도

2026년의 풍경은 이렇습니다. 순수 트랜스포머 기반 MoE가 최전선을 지배하지만, 256k 이상의 문맥이라는 틈새는 하이브리드가 차지합니다. Mamba-3의 상태 추적 개선은 다음 세대 하이브리드의 비율을 더 낮은 쪽으로, 다시 말해 SSM은 더 많이, 어텐션은 더 적게 가는 방향으로 끌어내릴 가능성이 있습니다.

사용해보기

code/main.py는 하이브리드 아키텍처를 위한 메모리 계산기입니다. SSM-트랜스포머 비율과 은닉 차원/레이어 수 설정이 주어지면 다음을 계산해 줍니다.

• 목표 문맥 길이에서의 KV 캐시.
• SSM 상태 메모리.
• 다양한 모델 형태(shape)에 대해 문맥 길이 N에서의 총 메모리.

계산기는 다음 구성을 지원합니다.

• 순수 트랜스포머 기준선(KV 캐시가 N에 따라 커집니다).
• Jamba 방식의 1:7 하이브리드.
• 순수 SSM(KV 캐시가 아예 없습니다).

공개된 형태에 대한 숫자는 Jamba-1과 Jamba-1.5 논문에서 그대로 가져왔고, 가상의 변형(hypothetical variant)에 대해서는 같은 식으로 외삽(extrapolate)했습니다.

실제 배포에서 통합(integration) 시 고려할 점은 다음과 같습니다.

• 대부분의 운영용 추론 서버(vLLM, SGLang)는 Jamba와 Mamba를 지원합니다. 다만 구체적인 버전을 확인할 필요는 있습니다.
• 256k 문맥에서는 Jamba의 메모리 이점이 동시 요청 처리량(concurrent-request throughput)에서 드러납니다. 같은 VRAM 위에 트랜스포머 시퀀스보다 Jamba 시퀀스를 더 많이 올릴 수 있습니다.
• 단독 모델로서의 Mamba-3는 아직 운영 환경에 출시되지 않았습니다. 1.5B 규모의 연구용 미리보기(research preview) 단계입니다.

산출물 만들기

이 강의는 outputs/skill-hybrid-picker.md를 만들어 냅니다. 작업 부하 명세(workload specification)로 문맥 길이 분포(context length profile), 과제 구성(task mix), 메모리 예산(memory budget)이 주어지면, 순수 트랜스포머와 Jamba 방식의 하이브리드, 순수 SSM 가운데 무엇을 골라야 할지 추천하고, 메모리와 품질 사이의 절충(tradeoff)을 명시적으로 짚어 줍니다.

연습문제

1. code/main.py를 실행해, 같은 형태의 32 레이어 순수 트랜스포머(은닉 4096, 헤드 32개)와 Jamba-1 하이브리드가 256k 문맥에서 차지하는 KV 캐시를 계산해 봅니다. AI21 논문이 주장한 약 8배의 메모리 절감을 직접 확인합니다. (난이도: 쉬움)

2. 계산기를 수정해 1:3 하이브리드(Mamba 4개 : 어텐션 1개)와 1:15 하이브리드(Mamba 14개 : 어텐션 1개)를 모델링합니다. 비율별 KV 캐시를 그래프(plot)로 그려 봅니다. 어떤 비율에서 KV 캐시가 SSM 상태 메모리와 같아지나요? (난이도: 중간)

3. Jamba 논문(arXiv:2403.19887)의 3절(Section 3)을 읽고, Mamba-2가 더 빠른데도 AI21이 Mamba-1을 사용하는 이유를 설명합니다. 힌트(hint): 하이브리드 절제 실험(ablation) 절에 그 근거가 정리되어 있습니다. (난이도: 중간)

4. Jamba 1.5 Large(총 398B, 활성 94B)에서 한 칸씩 건너뛰며 적용되는 MoE의 파라미터 오버헤드(parameter overhead)를 계산합니다. 활성 비율(active ratio)을 DeepSeek-V3(37B/671B)와 비교하고, 왜 Jamba의 아키텍처가 활성 비율을 더 높은 쪽으로 끌어올리는지 설명합니다. (난이도: 어려움)

5. Mamba-3 논문(arXiv:2603.15569)의 3절을 읽고, 복소수 상태 갱신이 데이터 의존적인 회전 임베딩과 어째서 동등한지 세 문장으로 설명합니다. Phase 7 · Lesson 04의 RoPE 유도와 연결지어 답합니다. (난이도: 어려움)

핵심 용어

더 읽을거리

• Lieber et al. — Jamba: A Hybrid Transformer-Mamba Language Model (arXiv:2403.19887) — 원조 Jamba 논문입니다. 비율 절제 실험과 256k 문맥 주장을 담고 있습니다.
• AI21 — Jamba 1.5: Hybrid Transformer-Mamba at Scale (arXiv:2408.12570) — 규모를 키운 가족 라인업입니다. 398B/94B와 12B/52B 공개 릴리스(public release)를 다룹니다.
• Gu, Dao — Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752) — Jamba가 토대로 삼은 선택적 SSM 논문입니다.
• Dao, Gu — Mamba-2 (arXiv:2405.21060) — 단순화된 구조적 상태 공간(structured-state-space)의 후속작입니다.
• Lahoti et al. — Mamba-3 (arXiv:2603.15569, ICLR 2026) — 복소수 상태와 MIMO를 통해 2026년 순수 SSM의 최전선을 다룹니다.
• Gu et al. — Efficiently Modeling Long Sequences with Structured State Spaces (arXiv:2111.00396) — LLM용 SSM 계보의 출발점이 된 S4 논문입니다.