Decoding Looped Transformers Better for (Almost) Free

Weihao Liu ,

Huangjie Zheng ,

Tianrong Chen ,

Rohit Dilip ,

Richard He Bai ,

Yizhu Jiao ,

Yuyang Wang ,

Ruixiang Zhang

AK's Picks LLM MQIA ML TAVR

循环式Transformer（Looped Transformers）通过在多个循环中重复执行同一个共享模块，从而实现参数高效性。每个循环都会生成一个中间表征，该表征均可用于解码下一个词元（token）；然而，标准解码过程会丢弃早期循环所生成的状态。由于早期循环所承载的计算量更小，因此这种循环机制天然地提供了对齐良好的“弱预测—强预测”配对，无需引入辅助模型或额外的外部训练。我们提出了LoopCD——一种无需额外训练的对比式解码框架：它通过将最终预测结果与某一次较早的循环输出进行对比，来引导词元选择。LoopCD可在两种模式下运行：一是在logit空间中操作，需额外进行一次输出前向传播（称为LoopCD-Logits）；二是在隐状态（hidden-state）空间中操作，不增加任何输出计算开销（称为LoopCD-Hidden）。在四种不同架构的循环式Transformer模型上，LoopCD在完整循环深度下均展现出显著且稳定的效果提升：例如，LoopCD-Logits将Ouro-2.6B-Thinking模型在AIME 2024考试上的pass@1准确率从61.88%提升至73.33%；LoopCD-Hidden则将Huginn模型在HumanEval基准上的pass@1准确率从22.56%提升至31.71%。尤为关键的是，这些性能提升使得模型仅需一半数量的循环次数，即可达到甚至超越未加引导的完整深度基线模型的表现，从而将前向推理所需的浮点运算量（FLOPs）降低22.5%至48.2%。LoopCD通过将循环过程中产生的中间隐状态转化为有效的解码引导信号，在大幅提升解码质量的同时，显著降低了推理阶段的计算开销。

Looped Transformers achieve parameter efficiency by repeatedly executing a shared block across recurrent loops. Each loop yields an intermediate representation decodable for the same next token, yet standard decoding discards earlier states. Because earlier loops embody less computation, recurrence inherently supplies aligned weak-and-strong prediction pairs without auxiliary models or external training. We introduce LoopCD, a training-free contrastive decoding framework that guides token selection by contrasting the final prediction with an earlier recurrent pass, operating either in logit space with one extra output pass (LoopCD-Logits) or in hidden-state space with zero output overhead (LoopCD-Hidden). Across four looped Transformer families, LoopCD delivers substantial, consistent gains at full recurrent depth: LoopCD-Logits raises Ouro-2.6B-Thinking's AIME 2024 pass@1 from 61.88% to 73.33%, while LoopCD-Hidden lifts Huginn's HumanEval pass@1 from 22.56% to 31.71%. Crucially, these performance gains enable halving the number of recurrent loops while still matching or exceeding full-depth unguided baselines, reducing forward FLOPs by 22.5% to 48.2%. By transforming intermediate recurrent states into effective guidance signals, LoopCD achieves superior decoding quality while substantially reducing inference compute.

许愿