Mamba 的召回缩放定律：基于哈希的理论与机制可解释性研究

Mamba 做关联回忆（Associative Recall, AR）时内部用的是什么算法，以及嵌入维 D 与状态维 N 必须如何随词表大小 V 和上下文中的事实数 Nf 缩放，才能达到接近完美的回忆。

• 场景：Transformer 解码时 KV cache 使内存随序列长度线性增长。Mamba、RWKV 等线性 RNN 用固定大小的循环状态换取常数内存，但把整个上下文压进一个向量会带来信息损失。作者认为，改进循环架构的内存利用因此成为活跃方向。
• 现有不足：AR、Multi-Query Associative Recall（MQAR）等合成基准与语言建模能力高度相关，已被用来诊断多种架构的回忆能力。作者指出，最接近的机制分析（Bick et al.）不研究 AR 也不给理论；Jelassi et al. 分析复制但不直接考察 AR，也不看模型实际学到的机制；Huang et al. 用 JL 引理给 Mamba 回忆界，但不做机制可解释性，也不考察实践中学到的机制。
• 核心观察：作者通过逐步去掉组件，把 Mamba 的回忆电路解释为隐式计算保持相似度的线性哈希函数；再借助 Johnson–Lindenstrauss（JL）引理，把容量写成 D、N、V、Nf 之间的缩放关系。
• 本文提出：Recall Scaling Laws。给定 V 与 Nf，该框架用于：（1）预测达到完美回忆所需的 D 与 N；（2）给定模型维度预测回忆成功概率；（3）分析多层模型和多头 SSM 模式。作者还给出从非压缩精确回忆到高概率回忆的四级保证阶梯，并称最后一级在信息论意义上基本最优。

关联回忆与合成基准

• Ba et al. [8]：提出 AR，用于评估回忆与检索。Arora et al. [9]：提出 MQAR，并称 AR 与语言建模性能强相关；Zoology 用该任务衡量并改进高效语言模型的回忆。
• 后续使用：一批工作用 AR 指导架构选择 [17, 15, 7, 24]、改进初始化 [16]、分析优化动态 [25]，并对循环 LLM 的局限给出理论与经验结论 [4, 5]。Wen et al. [4] 把上下文检索视为 RNN 相对 Transformer 的关键瓶颈。
• 最接近的三篇：Bick et al. [13] 用机制可解释性分析上下文检索的 gather-and-aggregate，但不研究 AR、不给理论。Jelassi et al. [26] 比较 Transformer 与 SSM 的复制能力并给理论，但不直接考察 AR，也不用机制可解释性看模型实际学到什么。Huang et al. [27] 用 JL 引理推导 Mamba 输入选择性对逼近能力、记忆和 AR 容量的影响，但不采用机制可解释性，也不调查实践中学到的机制。作者称本文给出更紧的回忆缩放律，同时写明状态大小与嵌入大小如何随词表缩放，并扩展到多层和多头；详细对比见附录 J。

Mamba 与选择性状态空间

• Gu and Dao [1]：Selective SSM（S6，即 Mamba），每步使用依赖输入的离散矩阵，并用 GPU 上的并行扫描实现。Dao and Gu [43]：Mamba-2，把 Transformer 与 SSM 通过结构化状态空间对偶联系起来，并引入多头 SSM 模式。
• 其他线性架构：RWKV [2]、Griffin 等线性 RNN [3]、Hyena [17]、H3 [15]、xLSTM [44]、DeltaNet [45]。作者把它们列为固定状态或线性递归的相关架构，计划在后续工作中分析其中一部分。

机制可解释性

• Elhage et al. [18]、Olah et al. [19]：把可观察行为连到具体组件，寻找网络中的紧凑电路。代表性例子包括 induction heads [14]、算术机制 [38–40]、任务向量 [41] 和 progress vectors [42]。
• 哈希与 JL：Andoni and Indyk [20]、Datar et al. [21]、Salakhutdinov and Hinton [22] 讨论保持相似度的哈希；Johnson and Lindenstrauss [23] 给出把高维点嵌入低维仍近似保持距离的引理。作者用这些工具分析回忆容量。

基线方法与基准：实验基线是按 Thm. F.2 固定权重构造的设计模型，以及 Thm. 4.2 的理论预测；被测对象是简化线性 Mamba、完整非线性 Mamba、简化多层 Mamba 和简化多头 Mamba-2。基准为 AR（L = 2Nf + 1 的单查询关联回忆）和 MQAR（Arora et al. [9]：词表一分为二，上下文含 Nf 个不重复键值对，查询段含 Nf 个重复键，其余为从全词表随机采样的 padding）。

作者把本文定位为：在机制上抽出 Mamba 实际学到的哈希式回忆电路，并由此得到同时约束 D 与 N 的缩放律，且覆盖多层与多头；这与只给理论界、或不研究 AR 的先前工作相区别。

整体思路是先构造一个能解 MQAR 的最小线性 Mamba 电路，再用不变算子和隐状态反投影核验训练模型是否学到等价电路，最后把该电路写成两次 JL 压缩，得到 Recall Scaling Laws。

任务与简化模型

• MQAR：词表大小 V 为偶数，键、值各占一半。序列长 L，上下文是 Nf 个不重复键值对（2Nf 个 token），其余 L − 2Nf 个 token 为查询段：其中 Nf 个查询各重复上下文中的一个键，其余为从全词表随机采样的 padding。对每个查询 qi，模型应输出匹配键所对应的值。忽略位置的输出记为 *，不进入损失。AR 是 L = 2Nf + 1 的单查询情形。
• 完整 Mamba 块：嵌入维 D、状态维 N、层数 Λ。块内含 S6、Conv1D 与逐元素运算；扩展维 Din = expand · D，卷积核长 Dconv。式 (1)–(3) 给出门控、离散化与逐通道递推 ht。
• 简化模型（Alg. 1）：去掉门控、离散化、非线性、偏置、归一化和残差。令 A 为单位阵后，SSM 变成 ht = ht−1 + x̂t Bt⊤，ŷt = ht Ct。默认 expand = 2、Dconv = 2。

两条回忆电路

• 非压缩精确电路（Thm. 3.1）：D = N = V、expand = 2、Dconv = 2 时，单层简化 Mamba 的回忆概率为 1。Pin 复制输入，Conv1D 做复制与移位，使 SSM 输入成为相邻 token 对；SB、SC 分别取出前一个 token 作为键、当前 token 作为查询。因 one-hot 正交且键唯一，查询与键的内积变成 Kronecker delta，输出坍缩为对应值。作者称该电路因维度随 V 线性增长而不现实。
• 压缩电路（Thm. 3.2）：用 E ∈ R^{D×V} 与 F ∈ R^{N×D} 做两次线性压缩，Ẽ = F E。在 V、Nf 以及 L = 4Nf 下，若 ND = O(Nf log V)、expand = 2、Dconv = 2，单层简化模型以高概率解决任务。匹配项给出目标值加 O(εv + εk) 误差；Nf − 1 个非匹配项各留 O(εv εk) 残差。作者称最坏情况下残差相加，需 ND = O(Nf log V)（式中带根号的版本见正文阶梯第 2 级）；均值情况下残差部分抵消，ND = O(Nf log V) 即以高概率足够。定量条件见 Lem. 4.1 与 Thm. 4.2。
• Fig. 2 的数据流：嵌入 → Pin 与 Conv1D 形成相邻对 → SB、SC 抽出键与查询并再压缩 → 隐状态按外积累加存储 → 键查询匹配读出 → Pout 取价值的一半 → 用 E 做值的解压缩。

训练模型中的等价电路

• 不变算子：解不唯一，正交变换 Q 下 Pin′ = Pin Q、E′ = Q⊤ E 给出等价电路。作者定义不依赖于这种自由度的有效算子 Êin、Πk,in、Πq,in、Πv,out，以及 Gvv = Πv,out Πv,in、Gkq = Πk,in⊤ Πq,in。输出可写成 yt = Στ Gvv ξτ ξτ⊤ Gkq ξt，ξt 为相邻 token 对。若学到的解等价于 Thm. 3.2，则 Gvv 应近似只保留值块，Gkq 应近似只连接查询与已存储的键。
• 隐状态即哈希表：理想电路把事实存成 V×V 外积表 Ht = Σn vn kn⊤。压缩后 Ht′ 是 D×N 的有损表，本身看起来不可解释；Alg. 2 把它投回词表空间得到 Ht″。作者称 Fig. 3 中 Ht″ ≈ Ht，因此学到的是“哈希、存储、匹配、检索”的接口。JL 投影被作者视为保持相似度的线性哈希。
• Conv1D：不能直接观察权重，因为它们与其他线性变换混在一起。作者用 Gvv、Gkq 检验复制与移位：期望 xt Gvv ξτ ≈ δt,τ，以及 ξt Gkq⊤ ξτ ≈ δt,τ−1（Fig. 15）。Tab. 8 的消融称，只含 SSM 与长度为 2 的 Conv1D 的简化块接近完美回忆，去掉 Conv1D 后完全失败。附录 F.3 把 Gvv、Gkq 非零块中比单位阵更复杂的图案解释为键值选择性。

缩放律

• 最坏情况（Lem. 4.1）：D = 4 log V / εv²，N = 4 log V / εk²。当 0 < εv, εk < 1 且 εv + εk + 2 Nf εv εk < 1/2 时，单层简化模型可完美解决 AR；若进一步 εv + εk + L εv εk < 1/2，则可完美解决 MQAR。
• 概率形式（Thm. 4.2）：在 Nf ≫ N, D 且 V ≫ 1 时，AR 与 MQAR 的成功概率分别用标准正态 CDF Φ 写出。AR 的自变量是 sqrt(ND / Nf) − sqrt(2 log V)；MQAR 还对查询位置做平均，分母里出现 Nf 与位置 t。证明见附录 F.4。
• 统一形式（Rem. 4.3）：precall ≈ Φ( sqrt(ND) / sqrt(a Nf + N) − sqrt(2 log V) )。AR 取 a = 1，MQAR 取 a ≈ 5/4。Rem. 4.4：要让 AR 概率 ≥ 1 − δ，需要 ND = Ω(Nf log(V/δ))。Lem. 4.5：任意有限状态循环模型，若隐状态有 S 个标量、每个 b 比特精度，并以大于 1 − δ 的概率解决该 AR（或 L = 4Nf 的 MQAR），则 bS = Ω(Nf log V)。对固定精度的 Mamba，这读作 ND = Ω(Nf log V)（Rem. F.8）。
• 引言中的四级阶梯：N = D = V 的非压缩模型可对每个上下文精确回忆（Thm. 3.1）；哈希模型在最坏情况下以更小维度、ND = O(Nf log V)（带根号的版本）达到精确回忆（Lem. 4.1）；均值情况下 ND = O(Nf log V) 即高概率足够（Thm. 4.2、Rem. 4.3）；高概率回忆也需要 ND = Ω(Nf log V)（Rem. 4.4，以及一般信息论下界 Lem. 4.5）。作者概括为：所需状态内存对事实数线性、对词表大小对数。

多层与多头

• 多层（Thm. 4.6）：把上式中的 N 换成 ΛN，即有效状态 Neff = ΛN。AR 仍 a = 1，MQAR 仍 a ≈ 5/4。证明见附录 G.2。
• 多头（Thm. 4.7）：固定 D、N，M > 1 的单层简化 Mamba-2 在 MQA、MKA、MVA、MHA 四种模式下满足 pMQA = pMKA ≤ psingle = pMVA ≤ pMHA。psingle 是同 D、N 的单头概率。作者称只有同时增大键与查询维度的 MHA 能提高回忆，代价是更多可训练参数；MQA 与 MKA 因跨头共享值压缩而更弱。

实验设置

• 模型：单层简化线性 Mamba、单层完整非线性 Mamba、简化多层 Mamba，以及固定总通道维 D 的简化多头 Mamba-2（MQA、MKA、MVA、MHA）。设计模型使用 Thm. F.2 的固定权重，不训练。
• 任务：AR 与 MQAR。Fig. 1 与 Fig. 7 的三组体制为 (i) V = 1024, L = 64, Nf = 16；(ii) V = 512, L = 128, Nf = 32；(iii) V = 512, L = 64, Nf = 16。更大规模为线性模型 D = 150、N = 75，V 从 1000 到 50000（对数间隔），L 从 1000 到 4000，Nf = L/4（Fig. 11、Fig. 12）。
• 指标：回忆准确率；理论对比用 precall ≈ Φ(x − b)，b ≈ sqrt(2 log V)。网格报告 best-of-seeds 准确率。附录 K.5 说明，作者关心的是是否存在回忆解，而不是优化找到它的概率。
• 种子与收敛：附录 K 的表给出每个网格点 3 或 5 个种子；Fig. 7 体制 (iii) 的完整模型用 5 个种子。Fig. 20 把达到准确率 ≥ 0.1 的种子计为收敛。早停准确率在该表中为 1.0 或 0.999。
• 实现：附录 K 给出数据集、模型、训练超参、梯度裁剪与算力。MQAR 实现跟随 Arora et al. 的 Zoology。论文写明主要实验代码放在公开仓库，正文首页给出地址。

主结果

Fig. 1（及更全的 Fig. 7）把回忆准确率画成 D–N 平面上的像素颜色，从左到右是设计模型、Thm. 4.2 理论、训练后的线性模型、训练后的完整模型。正文没有列出每个格子的准确率，因此下表只收论文写明的数值关系，不填未写出的格子。

设置	论文写出的结果	出处
单层，D–N 网格	理论与训练线性模型对齐；完整模型与线性模型接近；四列都呈 D–N 反向权衡	Fig. 1，Sec. 5.1
线性 AR	a = 1，曲线为 Φ(x − b)	Fig. 4(a)，Rem. 4.3
线性 MQAR	a ≈ 5/4；Fig. 9 另一处写作 a = 5/4	Fig. 4(b)，Fig. 9
完整非线性 MQAR	a_full ≈ (1/2) a_linear，启发式确定	Fig. 4(c)
多层 MQAR	(N, Λ) = (12,1)、(6,2)、(4,3)、(3,4) 时 pMQAR ≈ 0.38	Fig. 10
大规模线性模型	D = 150, N = 75；按 mV 拟合的斜率 1.18–1.56，理论斜率为 1	Fig. 11、Fig. 12

• 作者的解读：列 (c) 与列 (b) 对齐，说明理论预测贴近经验结果；列 (d) 与列 (c) 接近，说明去掉门控等组件后仍保留核心回忆行为。D 与 N 呈 Rem. 4.4 所预测的反向权衡。
• 缩放曲线：失败到接近完美回忆的过渡可收成一维变量 x = 1/σ，σ² 为输出 yt 的理论方差。MQAR 因查询更多而方差更大（a 更大），同样的 D、N 下更难达到高回忆。AR 与 MQAR 在重缩放后服从同一条律。Fig. 9 称 a = 1、a = 5/4、a = 5/8 这三列在三个体制上一致；其中 5/8 对应完整模型，与 Fig. 4 中 a_full ≈ (1/2) a_linear 且 a_linear ≈ 5/4 的写法相符。
• 作者对完整模型：两者共享哈希式键值压缩瓶颈，N 与 D 是主要容量约束。作者把“完整模型为何符合带启发式 a 的同一形式”留作后续工作，理论目前只对简化线性模型推导。

机制核验

• 测了什么：在训练后的线性 Mamba 上检查 Gvv、Gkq（Fig. 16）、隐状态反投影（Fig. 3、Fig. 13）以及 Conv1D 的复制移位（Fig. 15、Tab. 8）。
• 结果：作者称 Gkq 只关注查询 xt 与已存键 xτ−1，Gvv 只关注已存值 xτ，其余条目约等于 0；非零块比单位阵更复杂，附录 F.3 将其归于键值选择性。Fig. 3 中 Ht′ 看起来不可解释，投回词表后 Ht″ ≈ Ht。去掉长度为 2 的 Conv1D 后，简化块从接近完美回忆变为完全失败（Tab. 8；正文未给出表内逐格数字）。
• 作者的解读：训练解与 Thm. 3.2 的电路等价，隐状态是可解压的哈希表，Conv1D 的移位是该电路的必要部分。Fig. 13 在 D、N 为 (64, 32)、(64, 24)、(48, 32) 的设置下重复这一可视化。

多层与多头

• 多层（Fig. 5、Fig. 10）：准确率不由 N 与层数 Λ 分开决定，而由 Neff = ΛN 决定。左图的 N–Λ 网格呈反向关系；右图在 AR 与 MQAR 上都收到 Φ(x − b)。Fig. 10 给出一个可回查的点：上述四组 (N, Λ) 的 pMQAR 都约为 0.38。作者称这与 Thm. 4.6 一致。
• 多头（Fig. 6）：在 MQAR 上改变每头状态维 N 与头数 M，总通道维 D 固定。作者称与 Thm. 4.7 一致：MVA 与单头基线相当；MQA 与 MKA 因跨头共享值压缩而更弱；只有 MHA 通过增大有效状态提高准确率，同时参数更多。正文没有给出 Fig. 6 各点的准确率数字。
• 与语言建模的联系：作者称 Thm. 4.7 的趋势与完整 Mamba-2 的困惑度消融一致（Dao and Gu [43, Sec. 9.4, Table 5]）：MQA 与 MKA 的困惑度都差于 MVA。控制参数量时 MVA 更强，因此实践中被选用。作者认为本文分析可能抓住了完整 Mamba 在真实 NLP 任务中的同一权衡。该困惑度表不在本文中，本文未转写其数字。

其他消融与分析

• 种子：Fig. 19 显示 Fig. 7 体制 (iii) 的五个种子上，回忆相变边界一致，种子之间主要是孤立的优化失败。Fig. 20 左图为每格收敛种子比例，右图按该格 best-of-seeds 准确率分箱的每种子收敛率。正文未给出这些比例的具体数字。
• 大规模斜率：Fig. 12 按上下文长度 L 的线性拟合斜率为 1.18–1.56，理论斜率为 1。作者称理论缩放是近似关系，因此这一偏差可以理解。
• 重复键：附录 I.2 讨论键可重复的 MQAR。正文未给出该设置下的准确率数字。
• 电路极小性：附录 I.1 做回忆电路的极小性消融；与 Tab. 8 一起，作者用来支持“SSM + 长度 2 的 Conv1D”这一最小组合。
• 设计模型：Fig. 1(a)、Fig. 7(a) 用固定权重而非训练，作为次优但便于分析的基线；作者称其与理论列和训练列处于同一 D–N 权衡。
• 负面与缺口：完整模型的 a 是启发式拟合，不是从定理推出。Fig. 12 的经验斜率高于理论斜率 1。Gvv、Gkq 的非零块不是单位阵。理论不刻画门控分支。网格的逐格准确率大多只以颜色呈现，除 Fig. 10 的约 0.38 外，正文没有逐格数字。

作者指出的局限与后续方向

• 简化模型：Limitations 写明，分析依赖简化模型，因此有局限；特别是没有完整刻画完整 Mamba 架构的每个组件对回忆的贡献（Sec. 6）。
• 门控：同一节举例说，门控分支在回忆中的理论作用仍不清楚（Sec. 6）。
• 其他架构：作者称未来工作打算把分析扩展到 xLSTM [44]、RWKV [2]、DeltaNet [45]，以理解架构改动如何影响回忆能力（Sec. 6）。这是计划，不是已经完成的分析。
• 完整模型的瓶颈假设：Sec. 5.2 写明，完整非线性模型与简化模型共享哈希式键值压缩瓶颈这一假设，留待以后做详细调查。

实验覆盖范围

• 模型范围：实验包括按 Thm. F.2 固定权重的设计模型、训练后的单层线性模型、训练后的单层完整非线性模型、简化多层模型，以及简化多头 Mamba-2 的 MQA、MKA、MVA、MHA（Sec. 5，Fig. 1、Fig. 5、Fig. 6）。
• 任务与规模：主网格使用 Fig. 7 的三组 (V, L, Nf)；附录 A.2 另测线性模型在 V 从 1000 到 50000、L 从 1000 到 4000、Nf = L/4、D = 150、N = 75 上的 MQAR（Fig. 11、Fig. 12）。
• 报告方式：每个网格点用多个随机种子，报告 best-of-seeds 准确率；附录 K 的表写明种子数为 3 或 5，体制 (iii) 的完整模型为 5（Sec. 5.1，App. K.5，Fig. 19）。
• 机制实验：在训练后的线性模型上检查不变算子、隐状态反投影和 Conv1D 消融（Sec. 3.2，Fig. 3、Fig. 13、Fig. 15、Fig. 16，Tab. 8）。
• 论文未报告的内容：主 D–N 网格除 Fig. 10 的 pMQAR ≈ 0.38 外，正文未逐格列出准确率；Fig. 6 未给出各头模式的准确率数字；Tab. 8 的“接近完美”和“完全失败”没有配上表内数字；评审不使用单独的 LLM 评审模型，成功由任务准确率定义。论文未报告语言建模本身的复现数字，只引用 Dao and Gu [43] 的困惑度趋势。

作者从机制可解释性出发，反推 Mamba 解关联回忆的内部算法，并据此写出 Recall Scaling Laws。

• 问题：固定大小的循环状态把上下文压成一个向量。给定词表 V 和事实数 Nf，嵌入维 D 与状态维 N 要多大，单层、多层和多头 Mamba 才能以高概率做对 AR 与 MQAR。
• 电路：去掉门控、离散化、非线性和残差后，Conv1D 把相邻 token 配成对，SSM 用外积累加存储，两次线性投影分别压缩值与键查询。非压缩时 D = N = V 可精确回忆（Thm. 3.1）。压缩时，JL 投影充当保持相似度的哈希；训练模型的 Gvv、Gkq 以及投回词表的隐状态与该电路一致，去掉长度 2 的 Conv1D 则回忆失败。
• 定律：高概率下 ND = O(Nf log V)，同时 ND = Ω(Nf log V)。统一形式是 precall ≈ Φ(sqrt(ND) / sqrt(a Nf + N) − sqrt(2 log V))，AR 的 a = 1，MQAR 的 a ≈ 5/4。多层把 N 换成 ΛN。多头满足 pMQA = pMKA ≤ 单头 = pMVA ≤ pMHA。任意固定精度的有限状态循环模型也有 bS = Ω(Nf log V)。
• 实验：D–N 网格上，设计模型、理论、训练线性模型和完整模型呈同一反向权衡。完整模型用启发式 a_full ≈ (1/2) a_linear 拟合，表现好于线性模型，但缩放形式相同。多层 MQAR 中 Neff 相同的四组 (N, Λ) 准确率都约为 0.38。D = 150、N = 75 的线性模型在 V 至 50000、L 至 4000 上，经验斜率 1.18–1.56，理论斜率为 1。
• 作者的结论：理论预测与经验结果对齐，所需状态内存对事实数线性、对词表对数。作者同时说明，门控等完整组件的理论角色尚未刻画，并计划把同一分析扩展到 xLSTM、RWKV 和 DeltaNet。