本文提出Looped Transformers方法，通过循环结构和自适应步数显著提升了Transformer在算法任务上的长度泛化能力，在多种任务中优于传统方法。
Transformer, Reasoning, Prediction, Efficiency
Ying Fan, Yilun Du, Kannan Ramchandran, Kangwook Lee
University of Wisconsin-Madison, Massachusetts Institute of Technology, UC Berkeley
Generated by grok-3
Background Problem
Transformer模型在处理算法任务（如加法、奇偶校验等）时，尽管在训练长度范围内表现良好，但在面对未见过的更长输入时往往无法实现长度泛化（length generalization）。这一问题表明模型可能并未真正学习到任务的通用算法解法，而只是针对特定长度的数据进行了拟合。论文的出发点是探索如何通过架构创新和训练策略改进Transformer的长度泛化能力，解决的核心问题是：如何让模型在仅接触有限长度训练数据的情况下，学习到适用于任意长度的算法解法。
Method
论文提出了Looped Transformers方法，核心思想是通过循环结构和自适应步数来处理输入长度变化带来的计算复杂度差异。具体步骤如下：
- 定义n-RASP-L问题：将任务分解为可通过有限深度Transformer实现的迭代步骤（RASP-L操作），并根据输入长度n调整迭代步数T(n)。
- 模型架构：采用仅解码器（decoder-only）的Transformer结构，通过重复应用同一解码器块（decoder block）实现循环，每次迭代将原始输入嵌入与前一步输出嵌入相加（input injection），以保持对原始输入的连接。
- 训练策略：在无中间步骤监督的情况下，仅对最终输出进行端到端监督，同时利用预定义的步数T(n)指导训练，确保模型学习到与长度无关的中间步骤。
- 推理策略：推理时根据预定义步数（Oracle）或最大置信度准则（Maximum Confidence）自适应调整循环步数，以适应不同长度输入。
Experiment
实验在多个算法任务上验证了Looped Transformers的有效性：
- 任务与数据集：包括Parity、Copy、Addition、Binary Sum、Multiplication和Unique Set，训练数据长度范围为1到20（部分任务到12），测试长度远超训练长度（例如到40以上）。
- 实验设置：采用GPT-2解码器架构，使用课程学习策略逐步增加训练长度，与多种基线方法（如Vanilla NTP、NTP-Pause、NTP-Loop）进行对比，评估指标为精确匹配准确率（exact match accuracy）。
- 结果分析：Looped Transformers在所有任务上均表现出显著的长度泛化能力，例如在Parity任务中，训练到20位时可近乎完美泛化到40位以上，而传统NTP方法在训练长度+10时即表现不佳。相比NTP变体（如NTP-Pause和NTP-Loop），Looped Transformers的改进更为明显。
- 消融研究：验证了输入注入（input injection）和自适应深度的重要性，显示去除这些组件会降低泛化性能。
- 停止准则：最大置信度停止准则在部分任务（如Addition、Copy）上表现接近Oracle准则，但在非收敛任务（如Parity）上效果稍差。
- 评估合理性：实验设置较为全面，涵盖多种任务和基线方法，数据集设计考虑了长度分布的多样性；但训练数据中预定义步数的要求可能限制了实际应用场景，计算复杂度在推理时也较高。
Further Thoughts
Looped Transformers的循环结构和自适应步数为解决长度泛化问题提供了一个有前景的方向，但其对预定义步数T(n)的依赖可能限制了其在更复杂或未知任务上的应用。未来可以探索如何通过无监督或自适应方法自动学习所需的步数，而无需额外标注。此外，该方法与位置编码（positional encoding）的研究方向是正交的，结合最新的位置编码技术（如RoFormer或随机位置编码）可能进一步提升性能，尤其是在处理更复杂的序列任务时。我还想到其循环结构与RNN的相似性，是否可以借鉴RNN在时间序列任务中的优化技巧（如梯度截断）来缓解长步数训练时的计算负担？另外，论文中提到的不支持多重循环任务的局限性启发了我，或许可以结合多智能体系统（Multi-Agent Systems）或分层推理框架，将复杂任务分解为多个独立循环模块，逐步解决更广泛的算法问题。这些方向值得进一步探索。