AbbIE: Autoregressive Block-Based Iterative Encoder for Efficient Sequence Modeling¶

会议: NeurIPS 2025
arXiv: 2507.08567
代码: 无
领域: LLM推理
关键词: recurrent Transformer, iterative encoder, test-time scaling, fixed point, upward generalization

一句话总结¶

提出 AbbIE，一种将 decoder-only Transformer 的中间层（Body）进行递归迭代的架构，只需训练时用 2 次迭代，推理时即可通过增加迭代次数实现 upward generalization，在语言建模困惑度和 zero-shot ICL 任务上均超过标准 Transformer，且可作为标准 Transformer 的 drop-in 替代。

研究背景与动机¶

领域现状：Transformer 性能传统上通过增大模型参数量和训练数据来提升（scaling law）。近年 test-time scaling 成为新方向，但现有递归 Transformer（如 Geiping et al. 2025）需要大量迭代次数训练，且通常限于特定任务。

现有痛点：(1) GPU 显存增长慢于算力增长，限制了模型规模扩展；(2) 现有递归 Transformer 训练代价高（需要很多迭代），无法用作标准 Transformer 的通用替代；(3) 多数递归方法在训练迭代次数之外的推理迭代次数上无法泛化（upward generalization 失败）。

核心矛盾：如何在不大幅增加训练成本的前提下让 Transformer 具备 test-time compute scaling 能力？

本文目标 设计一种递归 Transformer 使得：(a) 单次迭代时等价于标准 Transformer；(b) 仅需 2 次迭代训练；(c) 推理时可扩展到任意迭代次数且性能持续提升。

切入角度：观察到 Transformer 的残差流自然地将原始输入信息注入每一层，这可能足以实现 Path Independence（收敛到不动点），从而无需额外投影矩阵就能递归迭代。

核心 idea：将 Transformer 分为 Head-Body-Tail 三段，只对 Body 做递归迭代，利用 inter-iteration residual connection 确保收敛，2 次训练迭代即可实现推理时的 upward generalization。

方法详解¶

整体框架¶

输入 token 经 embedding + Head（\(N_h\) 个 Transformer block）映射到 concept space，然后由 Body（\(N_b\) 个 block）递归迭代 \(r\) 次，最后由 Tail（\(N_t\) 个 block）映射回 token space 并 unembed。当 \(r=1\) 时完全等价于标准 Transformer。

关键设计¶

Head-Body-Tail 分割:
- 功能：将 Transformer 的层分为三组——Head 负责 tokenization 到 concept space，Body 负责迭代推理，Tail 负责 concept space 到 token space。
- 核心思路：Head 和 Tail 只执行一次，Body 重复执行 \(r\) 次。这种分割基于 Kaplan et al. 2024 的 concept space 理论——Transformer 的前几层在做去 tokenization，后几层在做 re-tokenization，中间层在概念空间中操作。
- 设计动机：避免对整个模型做递归（会破坏 tokenization），只在合适的抽象层级做迭代。
AbbIE-D（Diffusion-inspired 变体）:
- 功能：在 Body 的每次迭代之间添加 inter-iteration residual connection。
- 核心思路：\(h_{k+1} = B(h_k) + h_k\)，其中 \(B(\cdot)\) 是 Body 组件。相比 AbbIE-C 的 \(h_{k+1} = B(h_k)\)（仅依赖 Body 内部残差），AbbIE-D 额外增强了原始输入信息的相对比重，使得 \(h_0\) 的信号在多次迭代中不被稀释。
- 设计动机：实现 Path Independence（不动点收敛）需要每次迭代都有足够的原始输入信号。实验证明 AbbIE-C 发散，AbbIE-D 收敛。
仅 2 次训练迭代:
- 功能：训练时只使用 \(r=2\)（Body 执行 2 次），但推理时可以用 \(r=4, 8\) 等更多次。
- 核心思路：由于 AbbIE-D 满足不动点性质，2 次迭代足以让模型学会如何利用额外迭代改进表征。更大的模型（350M）在 \(r=4\) 时达到最低困惑度，说明成功实现了 upward generalization。
- 设计动机：降低训练成本至接近标准 Transformer 水平，同时保留 test-time scaling 能力。

损失函数 / 训练策略¶

标准 next-token prediction (NLL)。使用 AdamW (\(\beta_1=0.9, \beta_2=0.95\))，Warmup-Stable-Decay 学习率调度。训练 token 预算为 20 tokens/parameter（compute-optimal）。所有模型共享 tied embedding。

实验关键数据¶

主实验¶

Benchmark	指标	AbbIE-D (r=8)	AbbIE-D (r=2)	Std	0pt (r=2)
HellaSwag (350M)	Acc	36.6	33.8	30.1	29.7
LAMBADA (350M)	Acc	30.8	29.8	24.2	22.2
ARC-Easy (350M)	Acc	53.2	48.9	45.6	46.3
CommonsenseQA (350M)	Acc	23.7	20.0	20.0	20.0

注：CommonsenseQA 上标准 Transformer 和 0pt 均停留在随机基线(20%)，只有 AbbIE-D 在 r=8 时超过随机水平，表明迭代提升了涌现推理能力。

消融实验¶

配置	不动点收敛?	Upward Generalization?	说明
AbbIE-D	收敛	是 (350M at r=4)	inter-iteration residual 保证收敛
AbbIE-C	发散	否	仅靠 Body 内部残差不够
0pt (Geiping et al.)	收敛	否 (r!=2时困惑度崩溃)	收敛但无法泛化到训练外迭代次数

关键发现¶

AbbIE-D 是唯一在 2 次训练迭代下实现 upward generalization 的通用递归 Transformer，在 4x 训练迭代次数下 ICL 性能仍在提升。
困惑度比标准 Transformer 低约 5%，且遵循相同的 scaling law。
FLOP 效率随训练时间改善：虽然 AbbIE-D 的训练 FLOP 略高，但差距在长训练 run 中收敛。
关键发现：即使困惑度在 r=8 时略微回升，ICL 任务性能仍然持续提升——说明困惑度和下游任务性能之间的关系不是严格单调的。

亮点与洞察¶

r=1 时等价于标准 Transformer 是极好的工程性质：意味着可以先用标准方式训练，需要时再开启迭代，零风险采纳。
仅 2 次训练迭代的设计巧妙地平衡了训练成本和推理能力。对比 0pt 需要大量迭代训练但仍无法泛化，说明关键在于架构设计（inter-iteration residual）而非训练配方。
concept space 的理论框架为 Head-Body-Tail 分割提供了合理解释，也为未来自适应选择分割点提供了方向。

局限与展望¶

仅验证到 350M 模型：scaling 到 1B+ 模型是否仍然有效尚不清楚。作者提到 200M 模型的 upward generalization 不如 350M，暗示存在一个临界模型规模。
推理延迟线性增长：\(r\) 次迭代意味着推理延迟乘以 \(r\)（虽然参数不增加），对延迟敏感的场景不友好。
ICL 改善幅度有限：最大改善约 12%（HellaSwag），但绝对性能仍与同等规模标准模型持平。
未验证 generation 任务：所有评估都是 zero-shot ICL，缺少生成任务（如翻译、摘要）的评估。

评分¶

新颖性: ⭐⭐⭐⭐ Head-Body-Tail 分割 + inter-iteration residual 简洁有效，但递归 Transformer 方向已有不少工作
实验充分度: ⭐⭐⭐ 仅到 350M，缺少生成任务评估
写作质量: ⭐⭐⭐⭐ 逻辑清晰，理论（Path Independence）和实验结合好
价值: ⭐⭐⭐⭐ 提出了一种实用的递归 Transformer 替代方案，drop-in 特性很有吸引力