Finite State Automata Inside Transformers with Chain-of-Thought: A Mechanistic Study on State Tracking
作者: Yifan Zhang, Wenyu Du, Dongming Jin, Jie Fu, Zhi Jin
分类: cs.CL, cs.LG
发布日期: 2025-02-27 (更新: 2025-06-03)
🔗 代码/项目: GITHUB
💡 一句话要点
研究Transformer+CoT在有限状态自动机中的状态跟踪能力，揭示其内部机制。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer Chain-of-Thought 有限状态自动机 状态跟踪 可解释性 神经元激活 模型鲁棒性
📋 核心要点
- 大型语言模型结合CoT展现出强大的能力，但对其内部算法机制的理解仍然有限，特别是状态跟踪方面。
- 本文通过分析Transformer+CoT在状态跟踪任务中的表现，揭示了模型内部隐式学习的有限状态自动机（FSA）。
- 实验表明，Transformer+CoT在各种挑战性场景下，如噪声数据和长度泛化，均表现出良好的鲁棒性。
📝 摘要（中文）
本文研究了Transformer+CoT及其变体在状态跟踪方面的能力，验证了CoT的有效性。通过识别负责跟踪世界状态的电路（模型组件的子集），发现后期MLP神经元起着关键作用。提出了压缩和区分两个指标，并表明每个状态的神经元集合几乎达到100%的准确率，证明了模型内部嵌入了一个隐式的有限状态自动机（FSA）。此外，还探讨了三种具有挑战性的设置：跳过中间步骤、引入数据噪声和测试长度泛化。结果表明，Transformer+CoT学习了鲁棒的算法（FSA），突出了其在具有挑战性的场景中的弹性。代码已开源。
🔬 方法详解
问题定义：现有研究对Transformer+CoT的算法机制理解不足，尤其是在状态跟踪方面。虽然CoT在理论上可以提高表达能力，但缺乏对Transformer+CoT如何学习和执行状态跟踪任务的深入理解。现有方法难以解释模型内部如何表示和更新状态信息，以及其鲁棒性和泛化能力。
核心思路：本文的核心思路是通过将Transformer+CoT在状态跟踪任务中的行为与有限状态自动机（FSA）进行类比，来揭示其内部机制。通过分析模型内部神经元的激活模式，识别出负责跟踪状态的电路，并量化这些电路的压缩和区分能力，从而验证模型是否学习到了一个隐式的FSA。这样设计的目的是为了将复杂的模型行为简化为可解释的离散状态转换。
技术框架：整体研究框架包括以下几个阶段：1) 构建基于有限状态自动机的状态跟踪任务。2) 使用Transformer+CoT及其变体训练模型来解决这些任务。3) 设计指标（压缩和区分）来量化模型内部神经元对不同状态的表示能力。4) 通过消融实验和干预实验，验证关键神经元在状态跟踪中的作用。5) 在具有挑战性的场景下（如噪声数据和长度泛化）测试模型的鲁棒性。
关键创新：本文最重要的技术创新点在于提出了使用压缩和区分两个指标来量化Transformer+CoT内部神经元对不同状态的表示能力。这两个指标能够有效地衡量模型是否学习到了一个隐式的FSA，并为理解模型的内部机制提供了新的视角。与现有方法相比，本文不仅关注模型的整体性能，更深入地研究了模型内部的神经元行为，从而揭示了其状态跟踪的原理。
关键设计：本文的关键设计包括：1) 精心设计的状态跟踪任务，使其能够反映FSA的行为。2) 压缩指标的计算方式：衡量同一状态下神经元激活模式的相似性。3) 区分指标的计算方式：衡量不同状态下神经元激活模式的差异性。4) 使用后期MLP神经元作为状态表示的关键组件，并通过实验验证其重要性。5) 在训练过程中，使用了标准的Transformer架构和CoT技术，没有引入特殊的损失函数或网络结构。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Transformer+CoT在状态跟踪任务中表现出色，其内部神经元能够以接近100%的准确率区分不同的状态。在引入数据噪声和测试长度泛化时，模型仍然表现出很强的鲁棒性，证明了其学习到的FSA具有良好的泛化能力。这些结果突出了CoT在提升模型算法学习能力方面的作用。
🎯 应用场景
该研究成果有助于提升大型语言模型的可解释性和可控性，并为开发更鲁棒、泛化能力更强的模型提供指导。潜在应用包括：智能对话系统、机器人控制、流程自动化等领域，通过理解模型内部的状态表示，可以更好地进行模型调试、优化和安全验证。
📄 摘要（原文）
Chain-of-thought (CoT) significantly enhances the performance of large language models (LLMs) across a wide range of tasks, and prior research shows that CoT can theoretically increase expressiveness. However, there is limited mechanistic understanding of the algorithms that Transformer+CoT can learn. Our key contributions are: (1) We evaluate the state tracking capabilities of Transformer+CoT and its variants, confirming the effectiveness of CoT. (2) Next, we identify the circuit (a subset of model components, responsible for tracking the world state), indicating that late-layer MLP neurons play a key role. We propose two metrics, compression and distinction, and show that the neuron sets for each state achieve nearly 100% accuracy, providing evidence of an implicit finite state automaton (FSA) embedded within the model. (3) Additionally, we explore three challenging settings: skipping intermediate steps, introducing data noises, and testing length generalization. Our results demonstrate that Transformer+CoT learns robust algorithms (FSAs), highlighting its resilience in challenging scenarios. Our code is available at https://github.com/IvanChangPKU/FSA.