Hybrid 架构：Mamba 与 Attention 的融合

纯 SSM 为什么不够

在上一篇文章中我们看到，SSM/Mamba 通过固定大小的状态向量 $x \in \mathbb{R}^N$ 实现了线性复杂度和常数推理缓存。但这个优势本身就是它的根本局限：N 维状态向量无法精确存储 M 个 token 的完整信息（当 M >> N 时信息必然丢失）。

这个限制在 copying task 中暴露得最明显。给定输入 “A B C D | ? ? ? ?”，模型需要将前半段精确复制到后半段。Transformer 的 Attention 矩阵可以直接从输出位置连线到源 token，精确复制任意长度的序列。而 SSM 必须将所有源 token 压缩进固定大小的状态——越早的 token 信息衰减越多，序列越长准确率越低。

Jelassi et al. (2024) 在 “Repeat After Me” 论文中严格证明：Transformer 仅需两层即可 copy 指数长度字符串，而 SSM 做不到。这不是工程问题，而是架构的根本理论限制。

NVIDIA 在 8B 规模实验中发现，纯 Mamba-2 模型在 12 个标准 benchmark 上比纯 Transformer 平均低 2.65 分。但有趣的是，8B Mamba-2-Hybrid（混合架构）反而比纯 Transformer 平均高 2.65 分。这个结论引出了当前的共识：Hybrid 架构是严格更优解——纯 Attention 太贵，纯 SSM 太弱，混合才是最佳方案。

三种融合范式

如何将 SSM 和 Attention 组合在一起？目前有三种主流的融合范式，各有优劣：

交替式 (Interleaved)：Mamba 层和 Attention 层按固定比例交替堆叠。例如每 8 层中 7 层用 Mamba、1 层用 Attention。实现简单，可灵活控制 SSM:Attention 比例。大部分层用 Mamba 节省 KV cache，少量 Attention 层保证 in-context learning 能力。代表模型是 AI21 的 Jamba。

并行式 (Parallel)：同一层内 SSM heads 和 Attention heads 并行运行，输出相加融合。每层同时获得精确检索（来自 Attention）和高效摘要（来自 SSM）。实现较复杂，需要协调两种 head 的维度和融合方式。代表模型是 NVIDIA 的 Hymba。

共享式 (Shared)：少量 Attention 层被多个位置的 Mamba 层复用（参数共享）。极致参数效率——仅 2 个 Attention 层即可补足 SSM 的 copying/ICL 缺陷，通过 LoRA adapter 在不同调用位置做特化。代表模型是 Zyphra 的 Zamba2。

三种范式的选择取决于目标场景：

长上下文 + 大模型 → 交替式（高 SSM 比例压缩 KV cache）
小模型 + 端侧部署 → 共享式（参数效率最高）
强 ICL 需求 + 精确检索 → 并行式（每层都有 Attention）

通用规则：SSM 比例越高 → 长序列效率越好但 ICL 越弱；Attention 比例越高 → 反之。

Jamba：大规模交替式 Hybrid

Jamba (AI21 Labs, 2024) 是第一个成功在大规模（52B 参数）部署 Hybrid 架构的生产模型。

架构设计：80 层，每 8 层为一组：7 层 Mamba + 1 层 Attention。部分层集成 MoE（16 experts, top-2 routing），使得总参数 52B 但每个 token 仅激活 12B 参数。

设计动机：在 256K context 下，纯 Transformer 的 KV cache 会占用大量显存。Jamba 的 KV cache 仅在 Attention 层生成（占总层数的 1/8），Mamba 层使用固定大小状态，使得 KV cache 大小仅为同等纯 Transformer 的 1/8。

关键数据：

256K context window，支持极长文本
52B 总参数 / 12B active 参数 → 单卡 80GB GPU 可部署
在同参数规模下：长上下文任务显著优于纯 Transformer，短上下文任务持平或略优

Jamba 验证了一个重要结论：大部分 Transformer 层可以被 Mamba 替换而不损失质量，少量 Attention 层足以维持 ICL 能力。这种 7:1 的比例成为后续交替式 Hybrid 设计的参考基线。

Zamba2：参数高效的共享式 Hybrid

Zamba2 (Zyphra, 2024) 走了一条完全不同的路线：极致参数效率。

架构设计：Mamba2 backbone + 仅 2 个共享 Attention 层。这 2 个 Attention 层以 ABAB 模式在多个位置被复用——它们的 Q/K/V projection 权重是共享的，但每个调用位置附加独立的 LoRA adapter 做特化。这意味着不同位置的 Attention 计算共享大部分参数，但通过低秩调整仍能表现出位置特异性。

核心创新：

共享 Attention 参数 + LoRA 特化：兼得参数效率和层间差异
嵌入拼接 (Embedding Concatenation)：将原始 embedding 拼接到每个 Attention block 的输入，防止深层信息退化
Mamba2 backbone：利用 SSD 的 chunk-wise 算法加速训练

关键数据（2.7B 参数）：

推理效率等于 1-2B 参数的纯 Transformer
输出质量等于 3-4B 参数的纯 Transformer
对比 Phi-3 3.8B：2× 更快的 TTFT (Time to First Token)，27% 更少显存，1.29× 更低生成延迟

Zamba2 的启示是：Attention 的作用在于”补足”SSM 的缺陷，而非替代。仅 2 个共享 Attention 层就足够了——它们主要负责那些 SSM 做不好的精确检索任务。这使得 Zamba2 成为端侧部署的理想选择。

Hymba：并行融合 + Meta Tokens

Hymba (NVIDIA, 2024) 提出了最精细的融合方案：在每一层内同时运行 Attention heads 和 SSM heads。

架构设计：每层内部包含两组 heads：Attention heads 和 SSM heads。输入 token embedding 同时送入两组 heads，它们独立计算后输出相加。此外，Hymba 引入了 Meta Tokens——一组可学习的 token 前缀，拼接到输入序列前端。Meta tokens 存储全局关键信息（如语言特征、任务类型），减少 Attention 需要从实际 token 中检索的信息量。

1. 输入 + Meta Tokens

进一步优化：

Cross-layer KV sharing：相邻层共享 Attention 的 KV cache，减少存储开销
Partial sliding window attention：部分 Attention heads 使用局部窗口而非全局 Attention，进一步压缩计算

关键数据（1.5B 参数）：

超越 Llama-3.2-3B（+1.32% 平均分数），参数量仅为其一半
KV cache 比 Llama-3.2-3B 小 11.67×
Throughput 高 3.49×

Hymba 的并行式设计保证了每一层都能同时利用 Attention 的精确检索和 SSM 的高效摘要。代价是实现复杂度最高，但效果也最好——尤其适合需要强 ICL 能力的小模型。

模型对比

下面的表格汇总了主要 Hybrid 模型的关键指标。Throughput 和 KV Cache 大小相对于同规模纯 Transformer baseline：

几个关键观察：

没有统一最优比例：Jamba 用 7:1，Zamba2 用 ~6:1，Hymba 用 1:1 — 最优 SSM:Attention 比例因模型规模和目标任务而异
Hybrid 一致优于两个极端：纯 Transformer（太贵）和纯 Mamba（太弱）都不是最优解
KV cache 压缩是核心收益：所有 Hybrid 模型的 KV cache 都显著小于纯 Transformer

总结与展望

Hybrid 架构代表了当前序列建模的共识方向：

维度	纯 Attention	纯 SSM	Hybrid
训练效率	$O(N^2)$	$O(N)$	$O(N)$ ~ $O(N^2)$
推理 KV cache	$O(N)$	$O(1)$	大幅减少
ICL / Copying	精确	受限	精确（靠 Attention 层）
长序列支持	困难	原生	良好
工程复杂度	成熟	中等	较高

趋势与开放问题：

最优混合策略仍在探索：如何自动搜索最优 SSM:Attention 比例和放置位置？
SSM 能否突破 copying 限制？ 更大的状态维度或新的状态更新机制可能缩小差距
硬件协同设计：未来芯片是否会为 SSM 的 scan 操作提供专用硬件支持？
统一框架：Mamba-2 的 SSD 框架暗示 Attention 和 SSM 可能最终统一为同一个数学框架的不同特例