论文
通用 Transformer 需要内存:自适应递归推理中的深度状态权衡
Universal Transformers Need Memory: Depth-State Trade-offs in Adaptive Recursive Reasoning
摘要
我们研究学习记忆标记作为单块通用 Transformer 的计算暂存器,在 Sudoku-Extreme(一种组合推理基准)上具有自适应计算时间(ACT)。内存词元在经验上是必要的:没有它们,任何配置都无法达到不平凡的性能。最佳计数具有一个急剧较低的阈值(T=0 总是失败,T=8 可靠地成功),然后是稳定的平台(T=8-32,57.4% +/- 0.7% 精确匹配)和 T=64 的稀释边界。在停止侧压力(lambda 预热)下,平均停止随着整个平台的内存大小单调下降(从 T=8 时的 11.6 到 T=64 时的 8.3),这表明内存词元和思考深度以固定精度替代为资源。我们还发现了一个导致大多数训练失败的路由器初始化陷阱:默认的零偏差和格雷夫斯建议的正偏差都陷入了模型无法逃脱的浅停止平衡。将偏差反转为 -3(“深度启动”)可以消除故障模式,并且消融表明陷阱是 ACT 初始化所固有的,而不是我们架构的产物。通过可靠的训练,ACT 产生的种子方差比固定深度处理低一个数量级(+/-0.7 与 +/-9.3 pp); lambda 预热以匹配的精度恢复了 34% 的计算量;注意力头专门研究内存读取器、约束传播器和递归深度的积分器。代码:https://github.com/che-shr-cat/utm-jax。