论文
通过学习随机停止稳定循环 Transformer 中的外推
Stabilizing Extrapolation in Looped Transformers via Learned Stochastic Stopping
摘要
循环 Transformer 重复应用共享的 Transformer 块,在架构上自然适合可变长度算法任务。尽管它们可以表现出超出训练序列长度的强大长度泛化能力,但这种行为很脆弱,即使在性能良好的分布内解决方案中也会产生较高的分布外(OOD)方差。我们将这种差异追溯到简单算法任务中序列长度和循环数之间的虚假相关性。在训练期间将随机性引入循环数量可大幅减少 OOD 方差并稳定整个推理时间循环计数的预测。为了改进启发式随机化方案,我们进一步将 RL-Halting 作为学习随机计划进行分析,发现它通常可以改善准确性与稳定性的权衡。在二进制加法、Dyck-1、Unique Set 和 Copy 中,学习随机停止通常可以改善这种权衡,但也可以稳定次优计算。我们的工作表明,“何时停止”应被视为训练时设计选择,而不仅仅是推理时计算分配规则。