Wavefront decoding speeds up looped language model text generation
WaveFront Decoding: Parallelized Self-Speculative Decoding for Looped Language Models
Machine Learning
Summary
Generating text with looped language models takes a long time because the model repeats the same process many times for each word. The authors created Wavefront Decoding, a way to speed up this text generation without retraining the model. It works by making quick draft guesses and checking them simultaneously, fixing mistakes as they go. Their method runs over twice as fast as the usual way on some large language models.
What this means in practice
- •For language model engineers: Reduce text generation delay in looped language models by integrating efficient parallel decoding.
- •For cloud ai service operators: Cut inference costs and serve model output faster by implementing wavefront decoding with KV sharing in deployed language models.$Commercial implications: This method reduces compute time for language model serving, enabling cloud providers to offer faster, cheaper AI text generation services.
Authors
Hyeongju Ha, Jae-Joon Kim
Abstract
Looped language models repeatedly apply a weight-shared block to increase effective depth without increasing parameter count, but the resulting T sequential recurrent-block calls per generated token substantially increase decoding latency. To address the issue, we introduce Wavefront Decoding (WFD), a training-free self-speculative decoding framework designed for looped language models. WFD exploits two properties of these architectures: intermediate recurrence outputs provide effective draft predictions, and weight sharing allows token states at different positions and recurrence depths to be processed in one batched recurrent-block call. WFD organizes these mixed-depth states into a diagonal wavefront, continuously drafting new positions at shallow depth while advancing earlier positions toward full-depth verification. Unlike the phase-separated draft-then-verify schedule, WFD therefore co-batches drafting and verification within the same recurrent calls, while rejected drafts are corrected using full-depth predictions. Across six Spec-Bench task categories, WFD achieves 2.42x speedup on Ouro-2.6B and 3.54x on Huginn-3.5B over autoregressive decoding, consistently outperforming draft-then-verify. Cross-recurrence KV sharing further reduces wavefront KV traffic and increases WFD's speedup to 4.81x on Huginn-3.5B.