论文
循环 Transformer 中的自适应深度:诊断学习停止门和轨迹读数
Adaptive Depth in Looped Transformers: Diagnosing Learned Halting Gates and Trajectory Readouts
摘要
循环 Transformer 通过重复应用共享循环块来增加测试时间计算。循环 Transformer 中的学习停止目标通常使用单个退出分布作为推理时间停止规则和每个深度损失的训练时间加权。这将退出选择与轨迹形成联系在一起:门不仅选择要使用的循环状态,还决定每个中间状态的监督强度。因此,读出、诱导轨迹或它们的相互作用可能会导致自适应计算性能不佳。我们通过这个轨迹读出透镜,跨受控合成任务(模块化算术和二进制奇偶校验)以及大规模 Ouro-1.4B 和 2.6B 检查点,研究循环 Transformer 中的自适应深度。我们发现,固定先验深度监督在没有依赖于输入的停止策略的情况下塑造轨迹,产生困难感知轨迹,其中间状态暴露出有用的停止信号,并且简单的事后置信度读数通常匹配或优于学习的线性和 MLP 门。在冻结轨迹上安装门可以定位故障:它似乎主要源于联合门训练引起的轨迹,而不是有限的门表达能力。 Ouro 评估中也存在相同的模式,其中预训练的思考门具有竞争力,但并非一致帕累托最优,并且测量的延迟证实,平均退出深度的减少可转化为实际的推理时间节省。我们的系统诊断评估将循环 Transformer 中的自适应深度重新构建为轨迹形成和退出读出的联合问题,而不是单独的门学习,突出了先前的学习停止工作经常隐含的区别。