The Bayesian Geometry of Transformer Attention

Transformer模型在上下文中似乎常常表现出贝叶斯推理的行为，但要严格验证这一点一直不可能：真实数据缺乏解析形式的后验分布，而大型模型又将推理与记忆混淆在一起。我们通过构建“贝叶斯风洞”——即后验分布具有闭式解且记忆被证明不可能发生的受控环境——来解决这一问题。在这些设定下，小型Transformer模型能够以10⁻³至10⁻⁴比特的精度复现贝叶斯后验，而容量相当的多层感知机（MLP）则相差数个数量级，从而揭示出明显的架构差异。 在两类任务——双射消去和隐马尔可夫模型（HMM）状态追踪——中，我们发现Transformer通过一种一致的几何机制实现贝叶斯推断：残差流作为信念的载体，前馈网络执行后验更新，而注意力机制则提供基于内容的可寻址路由。几何诊断揭示了正交的关键向量基底、查询与关键向量之间逐步对齐的现象，以及一个由后验熵参数化的低维值空间流形。在训练过程中，该流形逐渐展开，而注意力模式保持稳定，这种“框架—精度分离”现象正是近期梯度分析所预测的结果。 综上所述，这些结果表明，层级化注意力机制通过几何设计实现了贝叶斯推断，从而解释了为何注意力机制不可或缺，以及为何扁平架构会失败。“贝叶斯风洞”为从小型可验证系统出发、机械性地理解大语言模型中观察到的推理现象提供了坚实基础。

论文试图验证Transformer是否在上下文中执行贝叶斯推理，这是一个长期被推测但难以严格验证的问题。由于真实数据缺乏解析后验分布，且大模型容易混淆推理与记忆，此前无法进行严谨验证。该问题在机制解释层面具有新颖性，尤其是在区分架构能力方面。

提出“贝叶斯风洞”（Bayesian wind tunnels）——一种受控环境，其中真实后验可解析获得且记忆不可能发生。在此设定下，发现小型Transformer能以极高精度（10^{-3}–10^{-4}比特）复现贝叶斯后验，而容量匹配的MLP则失败；并揭示了Transformer通过残差流作为信念载体、FFN执行后验更新、注意力实现内容寻址路由的几何机制。这一将架构组件映射到贝叶斯计算角色的机制性解释是核心创新。

设计了两个可控任务：双射消除和隐马尔可夫模型（HMM）状态追踪，确保理论后验可解且无记忆可能；实验显示Transformer与贝叶斯后验高度一致，MLP差距显著；引入几何诊断方法，发现正交键基、查询-键对齐演进、由后验熵参数化的低维值流形；观察到训练中流形展开而注意力模式稳定的现象（帧-精度分离），支持最新梯度理论预测；代码和设置有望开源，为小模型到大模型的推理机制桥梁提供可复制基础。

1. Language Models as Knowledge Bases? (2019) 2. In-context Learning and Induction Heads in Transformer Language Models (2022) 3. Towards Understanding Mechanisms in Transformers: The Case of Copying (2023) 4. Mathematical Foundations of Inductive Biases in Neural Networks (2024) 5. Gradient-Driven Discovery of Emergent World Models in Transformers (2023)