LoopAhead:面向 Looped Transformer 的推测解码框架,加速最高 7.82 倍
LoopAhead: Speculating Ahead for Faster Decoding in Looped Language Models
用循环层的早期状态提前起草 token,解码最高快 7.82 倍,输出分布不变,研究解码优化的可以看看
Looped Transformer 通过重复使用同一组 Transformer 层来提高参数效率,但每个 token 生成仍需串行计算和反复访问权重,解码延迟较高。LoopAhead 是一个针对循环模型的推测解码框架,包含轻量级 lookahead drafter(LAD),用早期循环状态一次预测多个未来 token。更深层的循环反馈决定哪些草稿分支继续、哪些暂停,最终 token 由全深度模型确认。在保持目标模型输出分布不变的前提下,LoopAhead 相比自回归解码最高实现 7.82 倍加速。
LoopAhead: Speculating Ahead for Faster Decoding in Looped Language Models
Looped Transformers improve parameter efficiency by repeatedly applying the same Transformer layers. However, generating each token still involves sequential computation and repeated access to model weights, leading to high decoding latency. Recent pipelined self-speculative decoding methods reduce this delay by drafting tokens from early-loop predictions and asynchronously advancing tokens at different recurrent depths in the same batch, but they do not fully exploit early-loop features for longer lookahead and still spend computation verifying incorrect speculative branches. We introduce \method, a speculative decoding framework for looped transformers that enables earlier drafting of future tokens and assigns complementary roles to different recurrent depths. A lightweight lookahead drafter (LAD) predicts multiple future tokens with early-loop states, and deeper-loop feedback helps decide which draft branches to continue and which to pause, while the full-depth model makes the final token decisions. LoopAhead achieves up to 7.82x decoding speedup over autoregressive decoding while preserving the target model's output distribution.