🤖 machine learning

From Markov to Laplace: How Mamba In-Context Learns Markov Chains

本文证明了单层 Mamba 模型能够高效地在上下文学习中实现马尔可夫链的最优拉普拉斯平滑估计量，从而在理论上首次建立了 Mamba 基于卷积的架构与贝叶斯/极小极大最优统计估计之间的正式联系。

原作者： Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, Michael Gastpar, Ashok Vardhan Makkuva

发布于 2026-06-23

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Marco Bondaschi, Nived Rajaraman, Xiuying Wei, Kannan Ramchandran, Razvan Pascanu, Caglar Gulcehre, Michael Gastpar, Ashok Vardhan Makkuva

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大局观：一种新型的 AI 大脑

想象一下当前人工智能领域的超级明星——Transformer（大多数聊天机器人背后的引擎），它就像一位才华横溢但有些笨重的图书管理员。它可以阅读整本书并瞬间找到联系，但如果书太长，它就会变得疲惫且缓慢，因为它试图同时记住每一个单词。

迎来 Mamba，一种更新型的 AI 模型。它就像一个敏捷、快速奔跑的信使。它处理信息的速度更快，占用的内存更少，是那个笨重图书管理员的一个极具前景的替代方案。但科学家们并不完全理解这个信使为什么如此聪明。他们知道它效果很好，但并不清楚其中的秘诀。

这篇论文就像是一个侦探故事，旨在弄清楚 Mamba 究竟是如何解决一种被称为**马尔可夫链（Markov Chains）**的特定类型谜题的。

谜题：预测下一步

为了测试 Mamba，研究人员给了它一个名为“下一标记预测（Next Token Prediction）”的游戏。

设定： 想象一个事件序列，比如一串彩色珠子（红、蓝、红、红、蓝……）。
规则： 下一个珠子的颜色取决于它之前的珠子颜色。这就是一个“马尔可夫链”。
挑战： 模型看到一串随机的珠子，必须猜出下一个是什么。至关重要的一点是，游戏的规则（例如红色紧跟蓝色出现的概率）会随着每一串新的珠子而改变。模型必须通过观察当前的字符串来“即时”掌握规则。这被称为上下文学习（In-Context Learning, ICL）。

发现：Mamba 是完美的统计学家

研究人员发现了一些令人惊讶的事实。即使是一个单层的 Mamba（一个非常简单的模型版本），也学会了成为这个游戏的完美统计学家。

在统计学领域，当你不完全了解规则时，有一种预测下一个珠子的“黄金标准”方法。它被称为拉普拉斯平滑（Laplacian Smoothing）（或“加 $\beta$ ”估计量）。

类比： 想象你正在猜测一副扑克牌中的下一张牌。如果你已经看到了 10 张 A 和 0 张 K，你可能会猜下一张是 A。但一个聪明的统计学家会想：“等等，我还没见过 K，也许只是运气不好。”所以，他们会在计数中加入一个微小的“幽灵”K 来避免过于武断。这可以防止他们在尚未观察到某种情况时就断言“概率为零！”。

论文的观点： Mamba 不仅仅是在瞎猜；它在数学上学会了完美地执行这种“幽灵计数”。它计算出之前模式出现的次数，并自动加上那一点点“平滑”处理，就像最优统计公式所要求的那样。

核心要素：卷积“手电筒”

研究人员问道：Mamba 是如何做到这一点的？ 是因为它复杂的门控机制吗？还是因为它的非线性激活函数？

他们通过移除 Mamba 的部分组件并观察哪些部分失效来进行实验。

发现： 最重要的部分是卷积（Convolution）。
类类比： 把卷积想象成 Mamba 用来照向近期历史的一个手电筒。
- 为了预测下一个珠子，Mamba 需要知道：“在过去的几步中，‘红色’紧跟‘蓝色’出现了多少次？”
- 卷积就像一个在历史记录上滑动的窗口，能够瞬间统计这些模式。
- 研究人员发现，如果你拿走了这个手电筒（卷积），Mamba 就会变“盲”并无法完成任务。如果你只保留手电筒（并移除复杂的门控机制），Mamba 仍然能完美地解决这个谜题。

关键结论： 这个“手电筒”（卷积）才是英雄。它允许 Mamba 回望过去、统计模式出现的次数，并应用最优的统计平滑，而不需要一个深奥复杂的“大脑”。

局限性：手电筒需要多大？

论文还研究了谜题的难度变化。

如果游戏取决于最后 1 颗珠子（一阶），一个小手电筒就能胜任。
如果游戏取决于最后 5 颗珠子（五阶），手电筒就需要更宽，以便同时看清这 5 颗珠子。
定理： 论文证明，要处理一个依赖于前 $k$ 个步骤的游戏，模型的“记忆容量”（隐藏维度）需要随 $k$ 指数级增长。这就像尝试记住一个密码：密码越长，要在脑海中同时记住它的难度就呈指数级上升。

与 Transformer 的对比

论文将 Mamba 与 Transformer（沉重的图书管理员）进行了对比：

Transformer： 为了解决这个“计数”谜题，一个 Transformer 通常需要两层（两个协同工作的脑子）来构建一种被称为“归纳头（Induction Head）”的机制来进行模式计数。单层 Transformer 则会失败。
Mamba： 单层 Mamba 能立即解决问题，因为其内置的卷积机制天生擅长高效计数。

总结

这篇论文揭示了 Mamba 在从上下文中学习时的超能力来自于一个特定的架构特征：卷积。

Mamba 学会了像一个完美的统计学家一样行动，使用被称为拉普拉斯平滑的方法来预测序列中的下一项。
它通过使用一个**卷积“手电筒”**来统计过去的模式并应用正确的统计调整来实现这一点。
即使是在一个非常简单的单层模型中也能实现，而 Transformer 则需要更复杂的结构才能完成同样的工作。

作者得出结论，这是首次有人正式将 Mamba 与这些最优统计估计量联系起来，证明了 Mamba 不仅仅是一个快速的模型，它从根本上理解如何高效地计数和对数据进行平滑处理。

技术摘要：从马尔可夫到拉普拉斯：Mamba 如何进行上下文学习以掌握马尔可夫链

问题陈述

尽管基于 Transformer 的模型驱动了近年来的 AI 进展，但其计算复杂度（随序列长度呈二次方增长）促使人们对高效替代方案（如结构化状态空间模型，SSMs，特别是像 Mamba 和 Mamba-2 这样的选择性 SSMs）产生了浓厚兴趣。尽管这些模型在经验上取得了成功并提升了推理速度，但 Mamba 的基本学习能力仍未得到充分理解。具体而言，目前缺乏关于 Mamba 上下文学习（In-Context Learning, ICL）机制的理论理解，而针对 Transformer ICL 的研究（例如归纳头机制）已有大量成果。

本文旨在解决的核心问题是：我们能否系统地表征 Mamba 的 ICL 能力？ 作者通过研究 Mamba 在随机马尔可夫链上的表现来探讨这一问题，这一框架此前曾被用于揭示 Transformer 的基本机制。

研究方法

本文采用了 马尔可夫-ICL 框架，其中输入 token 采样自随机的 $k$ 阶马尔可夫链。

数据生成： 转移核从狄利克雷先验（Dirichlet prior）中独立采样，确保每个测试序列都遵循唯一的马尔可夫分布。模型必须在不进行显式参数更新的情况下，在上下文中估计下一个 token 的分布。
最优基准： 该任务的理论最优估计器是 拉普拉斯平滑（Laplacian smoothing）（即 add- $\beta$ ）估计器，它是贝叶斯最优且极小极大（minimax）最优的。该估计器需要统计特定上下文中的 token 出现次数，并应用加性平滑。
架构分析： 作者分析了 Mamba 架构（特别是 Mamba-2），重点关注其递归、选择性和卷积机制。他们引入了一个简化的变体 MambaZero，该变体仅保留了嵌入、基于卷积的输入选择性和线性层，去除了门控和非线性部分，以隔离核心学习机制。
理论表征： 本文提供了一个构造性证明，展示了 MambaZero 如何能够表示最优的拉普拉斯估计器。同时，作者也确立了循环架构表示此类估计器所需的隐藏维度下界。

核心贡献

1. Mamba 学习最优拉普拉斯估计器

作者证明，即使是 单层 Mamba 也能高效地学习所有马尔可夫阶数的拉插拉斯平滑上下文估计器。

经验证据： 实验表明，单层 Mamba 在二元和多状态马尔可夫链上的表现能达到最优估计器的水平，显著优于 1 层 Transformer（后者无法完成该任务），并能与 2 层 Transformer 持平。
令人惊讶的发现： 与通常需要多层来实现归纳头（计数机制）的 Transformer 不同，Mamba 仅用一层即可实现这种最优的计数能力。

2. 卷积的关键作用

通过消融实验，论文指出卷积是使 Mamba 具备 ICL 能力的最核心组件，其重要性甚至超过了门控或非线性激活。

消融结果： 从 Mamba 中移除卷积会导致其无法完成马尔可夫预测任务，即使增加层数也无法解决；相反，一个仅包含卷积的简化模型（MambaZero）能够成功学习该任务。
机制： 卷积窗口大小 $w$ 必须满足 $w \ge k + 1$ （其中 $k$ 为马尔可夫阶数），才能追踪必要的上下文以进行计数。卷积允许模型维持一个状态，从而有效地根据当前上下文统计转移 $i \to j$ 的发生次数。

3. 理论表征与极限

表示容量： 作者提供了一个构造性证明（定理 1），表明 MambaZero 可以精确表示一阶马尔可夫过程的拉普拉斯估计器。其机制涉及使用状态 $H_t$ 来累积转移计数，其中卷积确保状态更新依赖于正确的上下文历史。
基本极限： 定理 2 确立了循环架构的下界：为了用有限比特精度表示 $k$ 阶马尔可夫过程的最优估计器，隐藏维度 $d$ 与精度 $p$ 的乘积必须按 $\Omega(2^k)$ 比例缩放。这与 Transformer 不同，Transformer 可以使用 3 层结构，使隐藏维度随 $k$ 线性增长来表示这些估计器。
选择性： 在非马尔可夫的“切换”马尔可夫过程中，论文展示了 Mamba 如何利用其选择性机制（衰减因子 $a_t$ ）在遇到切换 token 时重置计数，展示了所学机制的适应性。

实验结果

马尔可夫任务： 单层 Mamba 在各种马尔可夫阶数和状态空间大小下，与最优拉普拉斯估计器的损失差距几乎为零。
与 Transformer 的对比： 虽然 2 层 Transformer 可以解决该任务，但 1 层 Transformer 不行。Mamba 用 1 层即可解决。
自然语言： 在 WikiText-103 数据集上的实验显示，从 Mamba-2 中移除卷积会导致困惑度（perplexity）上升 11%，而从 Transformer 中移除卷积仅导致困惑度上升 2%。这表明，相比于 Transformer，卷积在 Mamba 处理复杂语言任务时发挥着更关键的作用。
切换过程： Mamba 通过将状态衰减因子 $a_t$ 设置为 0，正确实现了在遇到“切换” token 时重置转移计数的策略，有效地模拟了非平稳环境下的最优贝叶斯策略。

意义与主张

本文声称提供了 Mamba 与最优统计估计器之间的首次正式联系。

理论洞察： 研究揭示了 Mamba 的 ICL 能力并非仅仅是一种涌现属性，而是其通过卷积和递归实现最优计数能力的结构化基础。
架构理解： 该工作阐明了对于 Mamba 而言，卷积是其学习上下文相关统计数据的首要驱动力，这将其与 Transformer 基于注意力的机制区分开来。
泛化性： 研究结果从合成马尔可夫数据扩展到了自然语言建模，表明所识别的机制（由卷积驱动的计数）是 Mamba 在更广泛语言任务中取得成功的根本原因。

作者对未来的研究方向保持谦逊，指出虽然他们表征了表示能力，但 Mamba 收敛至这些估计器的学习动力学仍是一个开放的研究领域。他们还强调，其重点在于理解 Mamba 的内在能力，而非直接进行与 Transformer 的架构对比，尽管这种对比提供了必要的背景信息。