Loading post…
ModalExplainer
Multi-token Residual Prediction
Summary
Multi-Token Residual Prediction (MRP) adapts multi-token prediction for diffusion language models by training a small module to predict the residual between adjacent denoising steps rather than the full next-step distribution. This residual target is easier to learn, allowing a tiny 3-layer transformer to accurately predict corrections across multiple steps. In the static regime, MRP enables lossless speculative decoding with up to 1.56× throughput in SGLang, or direct decoding with a small quality cost. In the dynamic regime, MRP recovers up to +16 accuracy points lost by aggressive low-threshold decoding through a remasking mechanism. The module is trained on a frozen backbone and composes cleanly with existing inference methods.
Context
Diffusion language models (DLMs) decode by denoising a fully masked sequence, unmasking high-confidence positions per step. This creates a tradeoff: unmasking many tokens per step increases throughput but degrades quality. Existing acceleration methods live on this Pareto curve. Multi-token prediction (MTP) speeds up autoregressive models by predicting several tokens per forward pass, but applying MTP directly to DLMs fails because predicting the full next-step distribution across multiple steps causes error accumulation and collapse.
Approach / What changed
MRP attaches a small 3-layer transformer to a frozen DLM backbone. The module reads the backbone’s hidden states, predicts the inter-step logit residual, and adds it to the backbone’s own logits. Training uses a residual KL divergence objective: the frozen backbone is run twice (before and after revealing tokens), and MRP minimizes the difference between the two outputs on still-masked positions. Only the MRP module is trained; the backbone, LM head, and token embeddings remain frozen.
Takeaways
MRP predicts the residual between adjacent denoising steps rather than the full next-step distribution, which is a much easier target for a small module and prevents error accumulation across multiple steps.
In static denoising, MRP enables speculative decoding for lossless speedup (up to 1.56× throughput) or direct decoding for further speed with a small quality cost, giving a tunable tradeoff.
In dynamic denoising, MRP recovers up to +16 accuracy points lost by aggressive low-threshold decoding by using the inter-step residual to remask over-committed tokens without a second backbone pass.
FAQ
Why does direct MTP fail on diffusion LMs?
Direct MTP distills the full next-step distribution from scratch at each step, which requires reproducing a large, high-dynamic-range target. Per-step errors compound, causing accuracy to collapse from 84.8% to single digits by step 4 on GSM8K with SDAR-4B.
How does MRP achieve lossless speedup in static denoising?
MRP acts as a speculative draft module: it cheaply predicts several future denoising steps, and the backbone verifies them in parallel. Because MRP's predictions are accurate, verification rarely rejects them, yielding throughput up to 1.56× without quality loss.
What is the remasking mechanism in dynamic denoising?
After the backbone reveals tokens above a confidence threshold, MRP reads the inter-step residual to compute corrected confidences. Positions whose corrected confidence falls back below the threshold are remasked, recovering quality lost by aggressive unmasking without a second backbone forward pass.
Source and freshness
- Source
- modal.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.