语言模型通过可解释机制追踪状态变化，且可被引导优化。
(How) Do Language Models Track State?
- 用排列组合任务研究语言模型的状态追踪能力。
- 模型学习到两种机制：关联扫描和基于奇偶性的剪枝+扫描。
- 可通过中间训练任务控制模型选择更优的追踪策略。
Transformer 语言模型在故事生成、代码编写等行为中表现出对未观测世界状态的追踪能力。本文研究了在训练或微调后用于组合排列（即计算一组对象在一系列交换后的顺序）的语言模型如何实现这一能力。尽管该问题具有简单的代数结构，但其可还原为有限自动机模拟和布尔表达式求值等多种任务，因此是通用状态追踪的理想模型。我们发现语言模型在该任务上始终学习到两种状态追踪机制之一：第一种与近期理论工作（Liu et al., 2023；Merrill et al., 2024）提出的“关联扫描”构造高度相似；第二种则利用易于计算的特征（排列奇偶性）部分剪枝输出空间，再通过关联扫描进行精炼。采用前一种机制的模型泛化能力更强、收敛更快，且我们展示了如何通过中间训练任务引导模型选择其中一种。结果表明，无论预训练或微调，语言模型都能学习高效且可解释的状态追踪机制，且其出现可预测、可控制。
原文摘要 · Abstract (English)
Transformer language models (LMs) exhibit behaviors -- from storytelling to code generation -- that seem to require tracking the unobserved state of an evolving world. How do they do this? We study state tracking in LMs trained or fine-tuned to compose permutations (i.e., to compute the order of a set of objects after a sequence of swaps). Despite the simple algebraic structure of this problem, many other tasks (e.g., simulation of finite automata and evaluation of boolean expressions) can be reduced to permutation composition, making it a natural model for state tracking in general. We show that LMs consistently learn one of two state tracking mechanisms for this task. The first closely resembles the "associative scan" construction used in recent theoretical work by Liu et al. (2023) and Merrill et al. (2024). The second uses an easy-to-compute feature (permutation parity) to partially prune the space of outputs, and then refines this with an associative scan. LMs that learn the former algorithm tend to generalize better and converge faster, and we show how to steer LMs toward one or the other with intermediate training tasks that encourage or suppress the heuristics. Our results demonstrate that transformer LMs, whether pre-trained or fine-tuned, can learn to implement efficient and interpretable state-tracking mechanisms, and the emergence of these mechanisms can be predicted and controlled.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。