LoopFold speeds up looped transformers for better language tasks
LoopFold: Parallelizing Recurrent Depth in Looped Transformers
Computation and Language
Summary
Language models that repeat computations multiple times can become slow because each step depends on the previous one. The authors created LoopFold, which predicts what later steps will need in advance and corrects them after, allowing some steps to run at the same time. This reduces wait time while still improving the model’s accuracy on several tasks. LoopFold showed better accuracy gains and almost halved the delay compared to older methods.
What this means in practice
- •For ai model engineers: Reduce inference latency in repeated-computation transformers to speed up language model deployment while improving prediction quality.
- •For natural language processing teams: Enhance zero-shot task performance by integrating LoopFold to increase model accuracy without increasing parameter size significantly.
Authors
Feichen Song, Boyi Zeng, Yiqin Hao, Shixiang Song, He Li, Yizhou Yang, Zitong Wang, Yifan Liu, Ziwei He, Zhouhan Lin
Abstract
Looped Transformers increase computation per token by repeatedly applying a shared backbone, offering a way to improve language modeling and downstream task performance at a fixed backbone parameter count. However, each loop depends on the preceding loop, so increasing recurrent depth also increases decoding latency. Existing approaches include Parallel Loop Transformers, which reduce latency by overlapping loop computation across token positions, but yield rapidly diminishing marginal accuracy gains at larger depths. To address this challenge, we introduce LoopFold, which folds computation at multiple recurrent depths into parallel batches for the same token prefix. A lightweight predictor estimates inputs to later loops so they can start before earlier loops finish. A corrector then refines the outputs of these later loops using newly available results from earlier loops. All components are trained jointly using only the final next-token prediction loss. In experiments with 1.4B models, deeper LoopFold configurations achieve progressively higher average zero-shot accuracy across eight tasks, rising from 53.41% to 56.38%. Vanilla looped models and PLT show limited gains over the same depth range. Compared with a nine-loop vanilla model, LoopFold uses five sequential stages and additional parallel computation to reduce decoding latency by 46%.