Change of Thought: Adaptive Test-Time Computation

📄 arXiv: 2507.13569v1 📥 PDF

作者: Mrinal Mathur, Mike Doan, Barak Pearlmutter, Sergey Plis

分类: cs.LG, cs.AI

发布日期: 2025-07-17


💡 一句话要点

提出SELF-Transformer,通过自适应迭代注意力权重提升编码器Transformer的性能。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: Transformer 自注意力机制 迭代推理 自适应计算 编码器 深度学习 自然语言处理

📋 核心要点

  1. 传统Transformer编码器受限于表达能力,单次计算难以处理复杂任务。
  2. SELF-Transformer通过迭代细化注意力权重,实现输入自适应的计算量调整。
  3. 实验表明,该方法在编码器任务上显著提升了准确率,且未增加参数量。

📝 摘要(中文)

在单次固定深度传递中评估的Transformer在表达能力上被证明受限于常数深度电路类TC0。自回归地运行Transformer消除了这一限制,首先是在下一个token预测中,最近是在思维链推理中。这两种机制都依赖于反馈循环,将内部状态解码为token,然后在后续步骤中重新编码它们。虽然这种“大声思考”反映了人类的推理方式,但生物大脑在迭代时不会将中间状态外化为语言。为了在不诉诸token级别自回归的情况下提高编码器Transformer的表达能力,我们引入了SELF-Transformer:一个迭代地细化其自身注意力权重到固定点的编码器层。SELF-Transformer不是一次性生成重组输入序列的对齐矩阵,而是在内部迭代更新该矩阵,从而根据输入难度调整测试时的计算量。这种自适应性在不增加参数数量的情况下,在编码器风格的基准测试中产生了高达20%的准确率提升,表明测试时输入自适应对齐为适度的额外计算预算提供了显著的好处。因此,Self-Transformer在保持纯编码器架构的简单性的同时,恢复了迭代推理的大部分表达能力。

🔬 方法详解

问题定义:Transformer编码器在单次前向传播中表达能力有限,难以处理需要复杂推理的任务。现有的自回归方法虽然可以提升性能,但引入了额外的token生成和编码过程,增加了计算开销和延迟。因此,如何在不增加参数量和引入自回归的情况下,提升编码器Transformer的表达能力是一个关键问题。

核心思路:SELF-Transformer的核心思路是通过迭代地细化注意力权重,使模型能够根据输入数据的复杂程度自适应地调整计算量。通过内部迭代更新注意力矩阵,模型可以在不增加网络深度的情况下,逐步逼近最优的注意力分配,从而提升性能。这种方法模仿了人类大脑的迭代推理过程,但避免了将中间状态显式地表达为语言。

技术框架:SELF-Transformer的核心是一个特殊的编码器层,该层包含一个迭代更新注意力权重的模块。该模块接收输入序列,并初始化一个注意力矩阵。然后,该模块通过多次迭代,逐步更新注意力矩阵,直到其收敛到一个固定点。在每次迭代中,注意力矩阵被用于重新加权输入序列,并生成新的表示。最终,收敛的注意力矩阵被用于生成最终的输出表示。整个过程保持了编码器架构的简洁性。

关键创新:SELF-Transformer的关键创新在于其迭代更新注意力权重的机制。与传统的Transformer编码器一次性计算注意力权重不同,SELF-Transformer通过多次迭代逐步细化注意力权重,从而实现输入自适应的计算量调整。这种迭代更新机制使得模型能够处理更复杂的依赖关系,并提升表达能力。

关键设计:SELF-Transformer的关键设计包括迭代次数的选择、注意力权重的更新方式以及收敛判据的设定。迭代次数需要根据任务的复杂程度进行调整,过少的迭代可能无法充分利用模型的表达能力,而过多的迭代则可能导致计算开销过大。注意力权重的更新方式可以采用不同的优化算法,例如梯度下降或动量法。收敛判据可以基于注意力权重的变化量或输出表示的变化量来设定。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,SELF-Transformer在多个编码器风格的基准测试中取得了显著的性能提升,最高可达20%的准确率提升。与传统的Transformer编码器相比,SELF-Transformer在不增加参数数量的情况下,实现了更高的性能。这表明SELF-Transformer能够有效地利用计算资源,并提升模型的表达能力。

🎯 应用场景

SELF-Transformer具有广泛的应用前景,可用于机器翻译、文本分类、情感分析等各种编码器任务。其自适应计算的特性使其在资源受限的场景下具有优势,例如移动设备或边缘计算。此外,该方法还可以应用于其他类型的神经网络,以提升其表达能力和效率。

📄 摘要(原文)

Transformers evaluated in a single, fixed-depth pass are provably limited in expressive power to the constant-depth circuit class TC0. Running a Transformer autoregressively removes that ceiling -- first in next-token prediction and, more recently, in chain-of-thought reasoning. Both regimes rely on feedback loops that decode internal states into tokens only to re-encode them in subsequent steps. While this "thinking aloud" mirrors human reasoning, biological brains iterate without externalising intermediate states as language. To boost the expressive power of encoder Transformers without resorting to token-level autoregression, we introduce the SELF-Transformer: an encoder layer that iteratively refines its own attention weights to a fixed point. Instead of producing -- in one pass -- the alignment matrix that remixes the input sequence, the SELF-Transformer iteratively updates that matrix internally, scaling test-time computation with input difficulty. This adaptivity yields up to 20\% accuracy gains on encoder-style benchmarks without increasing parameter count, demonstrating that input-adaptive alignment at test time offers substantial benefits for only a modest extra compute budget. Self-Transformers thus recover much of the expressive power of iterative reasoning while preserving the simplicity of pure encoder architectures.