Table of Contents
Fetching ...

Annealed Relaxation of Speculative Decoding for Faster Autoregressive Image Generation

Xingyao Li, Fengzhuo Zhang, Cunxiao Du, Hui Ji

TL;DR

This paper tackles the latency challenge of autoregressive (AR) image generation by grounding speculative decoding (SD) in theory and introducing COOL-SD, an annealed relaxation of SD. It derives a near-tight upper bound on the TV distance between the relaxed output and the target distribution, identifies an optimal resampling distribution G_i^*, and demonstrates that a decaying (annealed) acceptance schedule further reduces distributional bias. The proposed COOL-SD framework consistently improves the speed–quality Pareto frontier over prior SD methods on two large AR image models (LlamaGen-XL and Lumina-mGPT), and even enhances performance when applied to LANTERN++. The combination of principled bias control and adaptive acceptance yields faster image generation with minimal quality loss, making AR models more viable for real-time applications.

Abstract

Despite significant progress in autoregressive image generation, inference remains slow due to the sequential nature of AR models and the ambiguity of image tokens, even when using speculative decoding. Recent works attempt to address this with relaxed speculative decoding but lack theoretical grounding. In this paper, we establish the theoretical basis of relaxed SD and propose COOL-SD, an annealed relaxation of speculative decoding built on two key insights. The first analyzes the total variation (TV) distance between the target model and relaxed speculative decoding and yields an optimal resampling distribution that minimizes an upper bound of the distance. The second uses perturbation analysis to reveal an annealing behaviour in relaxed speculative decoding, motivating our annealed design. Together, these insights enable COOL-SD to generate images faster with comparable quality, or achieve better quality at similar latency. Experiments validate the effectiveness of COOL-SD, showing consistent improvements over prior methods in speed-quality trade-offs.

Annealed Relaxation of Speculative Decoding for Faster Autoregressive Image Generation

TL;DR

This paper tackles the latency challenge of autoregressive (AR) image generation by grounding speculative decoding (SD) in theory and introducing COOL-SD, an annealed relaxation of SD. It derives a near-tight upper bound on the TV distance between the relaxed output and the target distribution, identifies an optimal resampling distribution G_i^*, and demonstrates that a decaying (annealed) acceptance schedule further reduces distributional bias. The proposed COOL-SD framework consistently improves the speed–quality Pareto frontier over prior SD methods on two large AR image models (LlamaGen-XL and Lumina-mGPT), and even enhances performance when applied to LANTERN++. The combination of principled bias control and adaptive acceptance yields faster image generation with minimal quality loss, making AR models more viable for real-time applications.

Abstract

Despite significant progress in autoregressive image generation, inference remains slow due to the sequential nature of AR models and the ambiguity of image tokens, even when using speculative decoding. Recent works attempt to address this with relaxed speculative decoding but lack theoretical grounding. In this paper, we establish the theoretical basis of relaxed SD and propose COOL-SD, an annealed relaxation of speculative decoding built on two key insights. The first analyzes the total variation (TV) distance between the target model and relaxed speculative decoding and yields an optimal resampling distribution that minimizes an upper bound of the distance. The second uses perturbation analysis to reveal an annealing behaviour in relaxed speculative decoding, motivating our annealed design. Together, these insights enable COOL-SD to generate images faster with comparable quality, or achieve better quality at similar latency. Experiments validate the effectiveness of COOL-SD, showing consistent improvements over prior methods in speed-quality trade-offs.
Paper Structure (41 sections, 7 theorems, 60 equations, 6 figures, 3 tables, 1 algorithm)

This paper contains 41 sections, 7 theorems, 60 equations, 6 figures, 3 tables, 1 algorithm.

Key Result

Theorem 1

The output tokens $x_{1:\tau+1}$ obtained with vanilla sd follows the target model distribution $x_{1:\tau+1}\sim P(\cdot\,|\,\texttt{pt})$.

Figures (6)

  • Figure 1: Illustration of vanilla SD and Cool-SD. By increasing the acceptance criterion $f_{i}$ according to an annealing schedule, along with a principled resampling distribution $G_{i}^{*}$, we can further increase the inference speed.
  • Figure 2: Qualitative results of Cool-SD. We demonstrate the trade-off between generation efficiency and image quality by comparing outputs from Cool-SD on Lumina-mGPT under different parameter settings. The speed-up factor is annotated to the left of each row. The first row shows images generated by Eagle-1 without any relaxation, serving as the baseline.
  • Figure 3: The trade-off curves between imaging quality (evaluated by FID) and accepted length/lantency. We compared Cool-SD with UniformRSD and LANTERN++ with $k=10$ on two target models: LlamaGen-XL and Lumina-mGPT. We tested on $5000$ random sampled captions from MSCOCO validation set.
  • Figure 4: Ablation study on LANTERN++ with our resampling distribution $G_{i}^{*}$, under different settings of $k$.
  • Figure 5: Ablation study on the effects of the annealing schedule on trade-off curves. We compare Cool-SD using a linear or exponential annealing schedule against the uniform schedule UniformRSD, across two target models: LlamaGen-XL and Lumina-mGPT. Evaluations are conducted on $5,000$ randomly sampled captions from the MS-COCO validation set.
  • ...and 1 more figures

Theorems & Definitions (7)

  • Theorem 1: Speculative Decoding Recovers Unbiased Target Distribution chen2023accelerating
  • Theorem 2
  • Proposition 1
  • Proposition 2: Annealing Property of the Relaxation Criterion
  • Proposition 3
  • Proposition 4
  • Proposition 5