arXiv:2605.18797cs.LGcs.AI2026-05被引 3

通过改进循环结构让模型更稳定,实现更高性能。

Simply Stabilizing the Loop via Fully Looped Transformer

论文配图:Simply Stabilizing the Loop via Fully Looped Transformer
图 1 · 摘自论文原文
  • 将循环信号分布到所有层,缓解残差爆炸问题。
  • 用注意力注入抑制梯度振荡,支持12次循环训练。
  • 适合需要灵活计算资源的部署场景。

提升模型性能通常依赖增加参数量,而循环Transformer通过重复使用相同Transformer模块,在不增加参数量或上下文长度的前提下,以额外计算换取性能提升,并可通过调整推理时循环次数自然平衡性能与计算开销。然而,随着循环次数增加,现有循环Transformer仍存在训练不稳定问题。分析表明,其根源在于梯度振荡和残差爆炸。为此,我们提出全循环Transformer(Fully Looped Transformer),引入两项无参改进:(1) 全循环架构,将层间信号分布至所有层以缓解残差爆炸;(2) 注意力注入,复用现有注意力模块抑制梯度振荡。该方法使模型可稳定训练至12次循环,而基线模型在此情况下崩溃。在较温和场景下,全循环变压器仍能将下游任务平均性能提升高达13.2%。实验表明,该方法提升了训练稳定性、增强了下游性能,并可在不同测试时计算预算下通过调整循环次数实现初步自适应。

原文摘要 · Abstract (English)

Scaling model performance typically requires increasing model size. Looped Transformer offers a compelling alternative by iteratively reusing the same Transformer blocks, trading additional computation for improved performance without increasing parameter count or context length. Because the number of loop iterations can be adjusted at inference, it also provides a natural mechanism for balancing performance and test-time compute. However, Looped Transformer still suffers from training instability when the number of loop iterations increases. Our analysis reveals that this instability stems from two sources: gradient oscillation and residual explosion. To address these two problems, we propose the Fully Looped Transformer, which introduces two parameter-free modifications: (1) Fully Looped Architecture, which distributes inter-loop signals across all layers to mitigate residual explosion; (2) Attention Injection, which reuses the existing attention block to suppress gradient oscillation. These modifications stabilize training dynamics, enabling the Fully Looped Transformer to be trained stably up to 12 loop iterations, whereas other baseline looped models collapse in this regime. In milder settings where Looped Transformer does not collapse, Fully Looped Transformer still improves average downstream-task performance by up to 13.2\%. Overall, our experiments demonstrate that Fully Looped Transformer improves training stability, enhances downstream performance, and provides preliminary adaptability under different test-time compute budgets by varying loop iterations at inference.

Transformer循环结构训练稳定推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。