FastGuide speeds up reward guidance for diffusion language models
FastGuide: Accelerating Reward Guidance for Diffusion Large Language Models
Computation and LanguageMachine Learning
Summary
Guiding large language models to produce desired outputs can be slow because it requires repeated and costly computations. The authors present FastGuide, a method that combines parallel and step-by-step decoding to make this guiding process faster. It uses clever caching and only updates uncertain words later to save time. Tests show FastGuide generates similar quality text much quicker than older methods.
What this means in practice
- •For language model deployers: Reduce inference time for language models guided by reward functions, speeding up controlled text generation in real-world applications.
- •For natural language generation engineers: Implement faster decoding strategies when combining diffusion models with reward models to improve product workflows involving guided text output.
Authors
Darshan Thaker, Lachlan Ewen MacDonald, René Vidal
Abstract
Gradient-based reward guidance provides a flexible way to use downstream reward models to control masked diffusion language models at inference time. However, its computational cost remains high as each decoding iteration incurs expensive diffusion model forward passes and reward model backpropagation steps. To address this, we introduce FastGuide, an adaptive hybrid of parallel and autoregressive decoding to accelerate reward guidance for diffusion language models. In analogy to parallel decoding, FastGuide amortizes the cost of reward model backpropagation by computing guidance once per decoding step and reusing it to generate multiple tokens. Within each decoding step, FastGuide makes diffusion forward passes autoregressive by unmasking tokens one at a time while efficiently recomputing token distributions after each unmasking by utilizing KV caching techniques and sparse recomputation of attention. Lastly, to adapt hybrid decoding to the model's confidence, FastGuide defers any token that the model is unconfident about under its recomputed distribution. Experiments on three reward benchmarks demonstrate that FastGuide is up to $4.4\times$ faster than sequential reward-guided decoding while retaining similar generation quality.