Transformers represent belief state geometry in their residual stream
作者: Adam S. Shai, Sarah E. Marzen, Lucas Teixeira, Alexander Gietelink Oldenziel, Paul M. Riechers
分类: cs.LG, cs.CL
发布日期: 2024-05-24 (更新: 2025-02-04)
💡 一句话要点
Transformer在残差流中以线性方式表征信念状态几何结构，蕴含未来信息。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 信念状态 残差流 线性表示 未来预测
📋 核心要点
- 现有方法缺乏对Transformer内部表征数据生成过程信念状态的理解。
- 该论文提出信念状态在Transformer残差流中线性表示，即使几何结构复杂。
- 实验证明残差流中的信念状态包含超越下一词预测的未来信息。
📝 摘要（中文）
本文研究了大型语言模型在进行下一词预测训练时所构建的计算结构。研究表明，这种结构是由数据生成过程隐藏状态上的信念更新的元动力学决定的。基于最优预测理论，研究者预测并发现信念状态在Transformer的残差流中以线性方式表示，即使预测的信念状态几何结构具有高度非平凡的分形结构。研究进一步探讨了信念状态几何结构在最终残差流中表示或分布在多个层的残差流中的情况，并提供了一个解释这些观察结果的框架。此外，研究表明推断的信念状态包含关于整个未来的信息，超越了Transformer显式训练的局部下一词预测。这项工作提供了一个通用框架，将训练数据的结构与Transformer内部激活的几何结构联系起来。
🔬 方法详解
问题定义：现有大型语言模型，特别是Transformer，在进行下一词预测训练时，其内部构建了怎样的计算结构？现有方法缺乏对这种结构与训练数据内在联系的深入理解，特别是如何表征数据生成过程的信念状态，以及这种表征是否蕴含了超越局部预测的信息。
核心思路：论文的核心思路是，Transformer在训练过程中学习到的结构反映了数据生成过程隐藏状态上的信念更新的元动力学。具体来说，Transformer的残差流以线性方式表征了信念状态的几何结构，即使这种几何结构非常复杂（如分形结构）。这种表征不仅用于预测下一个词，还包含了关于整个未来的信息。
技术框架：该研究主要通过理论分析和实验验证相结合的方式进行。首先，基于最优预测理论，研究者预测信念状态应该在残差流中线性表示。然后，通过分析Transformer的内部激活，验证了这一预测。研究还考察了信念状态几何结构在不同层残差流中的分布情况，并提出了一个解释框架。
关键创新：该论文最重要的创新在于揭示了Transformer残差流与信念状态几何结构之间的线性关系。这种关系不仅提供了一种理解Transformer内部工作机制的新视角，也为研究Transformer如何学习和表示复杂数据结构提供了一个通用框架。此外，论文还发现Transformer学习到的信念状态包含了超越局部预测的未来信息，这表明Transformer具有更强的推理能力。
关键设计：论文的关键设计在于如何从Transformer的残差流中提取和分析信念状态。具体的技术细节（如参数设置、损失函数、网络结构等）在论文中没有详细描述，但可以推断研究者使用了标准的Transformer架构和训练方法，重点在于如何解释和分析训练后的模型内部激活。
🖼️ 关键图片
📊 实验亮点
研究发现，即使信念状态几何结构具有高度非平凡的分形结构，Transformer的残差流仍然能够以线性方式表示这些信念状态。更重要的是，研究表明，Transformer学习到的信念状态包含了关于整个未来的信息，超越了其显式训练的局部下一词预测目标。这些发现为理解Transformer的内部工作机制提供了重要的线索。
🎯 应用场景
该研究成果可应用于提升语言模型的理解和推理能力，例如，通过显式地操纵残差流中的信念状态表征，可以引导模型生成更符合逻辑和上下文的文本。此外，该研究也为理解和改进其他类型的神经网络提供了新的思路，例如，可以借鉴该方法来分析神经网络如何学习和表示复杂的数据结构。
📄 摘要（原文）
What computational structure are we building into large language models when we train them on next-token prediction? Here, we present evidence that this structure is given by the meta-dynamics of belief updating over hidden states of the data-generating process. Leveraging the theory of optimal prediction, we anticipate and then find that belief states are linearly represented in the residual stream of transformers, even in cases where the predicted belief state geometry has highly nontrivial fractal structure. We investigate cases where the belief state geometry is represented in the final residual stream or distributed across the residual streams of multiple layers, providing a framework to explain these observations. Furthermore we demonstrate that the inferred belief states contain information about the entire future, beyond the local next-token prediction that the transformers are explicitly trained on. Our work provides a general framework connecting the structure of training data to the geometric structure of activations inside transformers.