用记忆令牌提升通用Transformer的推理能力，解决深度与计算量的权衡问题。
Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning
- 引入记忆令牌作为计算草稿，让单块Transformer通过自适应计算时间完成复杂推理。
- 记忆令牌数量在8-32之间表现稳定（准确率57.4%±0.7%），过少或过多均导致性能下降。
- 发现初始化陷阱，负偏置可避免训练失败，且注意力头会随递归深度分工协作。
我们在Sudoku-Extreme这一组合推理基准上，研究了学习型记忆令牌作为单块通用Transformer结合自适应计算时间（ACT）的计算草稿。实验证明记忆令牌不可或缺：无记忆配置无法达到非平凡性能。最优记忆数存在明显下限（T=0始终失败，T=8可靠成功），随后进入稳定平台期（T=8-32，精确匹配率57.4%±0.7%），并在T=64时出现稀释效应。在停止侧压力（lambda warmup）下，平均停止深度随记忆规模单调下降（从T=8时的11.6降至T=64时的8.3），表明记忆与思考深度可在固定精度下互换资源。我们还发现路由初始化陷阱：默认零偏置和Graves推荐的正偏置均导致模型陷入浅层停止均衡，无法逃脱。将偏置反转为-3（“深起点”）可消除该失败模式，消融实验显示此陷阱源于ACT初始化本身而非架构设计。可靠训练下，ACT的种子方差比固定深度处理低一个数量级（±0.7对比±9.3个百分点）；lambda warmup在相同精度下节省34%算力；注意力头在递归深度中分化为记忆读取、约束传播和整合模块。代码：https://github.com/che-shr-cat/utm-jax。
原文摘要 · Abstract (English)
We study learned memory tokens as a computational scratchpad for a single-block Universal Transformer with Adaptive Computation Time (ACT) on Sudoku-Extreme, a combinatorial reasoning benchmark. Memory tokens are empirically necessary: no configuration without them reaches non-trivial performance. The optimal count has a sharp lower threshold (T=0 always fails, T=8 reliably succeeds) followed by a stable plateau (T=8-32, 57.4% +/- 0.7% exact-match) and a dilution boundary at T=64. Under halt-side pressure (lambda warmup), mean halt drops monotonically with memory size across the plateau (from 11.6 at T=8 to 8.3 at T=64), showing that memory tokens and ponder depth substitute as resources at fixed accuracy. We also identify a router initialization trap that causes the majority of training runs to fail: both default zero-bias and Graves' recommended positive bias settle into a shallow halt equilibrium the model cannot escape. Inverting the bias to -3 ("deep start") eliminates the failure mode, and ablation shows the trap is inherent to ACT initialization rather than an artifact of our architecture. With reliable training, ACT yields an order of magnitude lower seed variance than fixed-depth processing (+/-0.7 vs +/-9.3 pp); lambda warmup recovers 34% of compute at matched accuracy; and attention heads specialize into memory readers, constraint propagators, and integrators across recursive depth. Code: https://github.com/che-shr-cat/utm-jax.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。