Technical Summary: Annealed Relaxation of Speculative Decoding for Faster Autoregressive Image Generation

Problem Statement

Autoregressive (AR) models have emerged as a powerful paradigm for image generation, often matching or surpassing diffusion-based methods. However, their sequential nature—predicting one token at a time—results in substantial inference latency, particularly for high-resolution images requiring thousands of tokens. While Speculative Decoding (SD) has successfully accelerated Large Language Model (LLM) inference by using a fast draft model to propose tokens for verification by a slower target model, its application to image generation has been limited.

Prior research attributes this limitation to token selection ambiguity: in image generation, multiple candidate tokens often receive nearly identical probabilities, making it difficult for the draft model to accurately predict the target model's top-ranked tokens. This mismatch leads to low token acceptance rates during the verification stage. Existing attempts to address this, such as LANTERN and LANTERN++, introduce "relaxed" verification schemes that permit limited divergence from the target distribution. However, these methods rely on heuristic, empirically driven designs lacking theoretical grounding regarding their divergence from the target distribution, resulting in modest performance gains.

Methodology: COOL-SD

The paper proposes COOL-SD (Annealed Relaxation of Speculative Decoding), a framework built on two key theoretical insights to improve the speed-quality trade-off in AR image generation.

1. Theoretical Foundation of Relaxed SD

The authors generalize the standard SD framework by replacing the rigid, exact-match acceptance test with a controllable acceptance probability $f_i$. This relaxation allows for longer expected accepted prefixes but introduces distributional drift (bias).

• Optimal Resampling: The authors derive an almost-tight upper bound on the Total Variation (TV) distance between the output distribution of the relaxed SD and the target distribution. By minimizing this bound, they derive token-wise optimal resampling distributions ($G^*_i$). These distributions are designed to correct the bias introduced by relaxed acceptance, ensuring the output closely approximates the target model.
• Annealing Property: Through perturbation analysis, the authors reveal that for a fixed expected number of accepted tokens, the TV bound is minimized when acceptance probabilities decrease monotonically across token positions. This suggests that relaxing the acceptance criterion more aggressively at the beginning of a sequence and tightening it later is theoretically superior to a uniform relaxation.

2. The COOL-SD Algorithm

Based on these insights, COOL-SD integrates:

• Exponentially Decaying Acceptance Schedule: Instead of a fixed relaxation parameter, COOL-SD employs an acceptance probability $f_i$ that decays exponentially as the decoding position $i$ increases ($\omega_i = \delta \cdot \exp(-\nu \cdot i - \mu)$). This "annealing" strategy reduces distributional shift compared to uniform relaxation.
• Principled Resampling: The algorithm utilizes the derived optimal resampling distribution $G^*_i$ (Eqn. 5) to correct rejected tokens, rather than relying on heuristic adjustments used in prior works.

Key Contributions

1. Theoretical Analysis of Relaxed SD: The paper establishes a rigorous theoretical basis for relaxed speculative decoding. It derives the expected number of accepted tokens and a nearly-tight upper bound on the TV distance between the relaxed output and the target distribution.
2. Optimal Resampling Derivation: It provides a closed-form solution for token-wise optimal resampling distributions that minimize the derived TV bound, offering a principled alternative to heuristic resampling strategies.
3. Annealing Insight: The authors prove that a monotonically decaying acceptance schedule minimizes the distributional bias bound, motivating the design of the annealed acceptance criterion.
4. COOL-SD Framework: They propose COOL-SD, which combines the derived optimal resampling distributions with an exponentially decaying acceptance schedule.
5. Empirical Validation: Extensive experiments demonstrate that COOL-SD consistently improves the speed-quality Pareto frontier over existing SD variants (Eagle-1, LANTERN++) on AR image generation tasks.

Experimental Results

The authors evaluated COOL-SD on two target models: LlamaGen-XL (775M) and Lumina-mGPT (7B), using the MS-COCO 2017 validation set.

• Speed-Quality Trade-off: COOL-SD achieves significantly faster inference than lossless baselines (Eagle-1) and prior relaxed methods (LANTERN++) with only minor quality degradation.
  • On Lumina-mGPT, COOL-SD increased the average accepted length from 2.76 (Eagle-1) to 3.11, reducing latency from ~71.6s to ~63.2s (a 2.69x speed-up) while maintaining comparable FID and CLIP scores.
  • On LlamaGen-XL, COOL-SD achieved a 2.72x speed-up (3.72s latency) with an accepted length of 3.34, showing negligible degradation in FID and CLIP scores compared to the baseline.
• Comparison with LANTERN++: COOL-SD consistently produced higher-quality images at the same latency or achieved greater acceleration at similar quality levels. Qualitative results showed that LANTERN++ suffered from visible artifacts and patchy backgrounds at higher speed-ups, whereas COOL-SD preserved fidelity up to a 3.10x speed-up.
• Ablation Studies:
  • Annealing vs. Uniform: Replacing the exponential decay with a uniform relaxation parameter (UNIFORMRSD) resulted in inferior FID-Accepted Length trade-offs, confirming the benefit of the annealing schedule.
  • Resampling Distribution: Applying the derived optimal resampling distribution $G^*_i$ to LANTERN++ significantly improved its performance, validating the importance of the principled resampling strategy over the heuristic one used in LANTERN++.

Significance and Claims

The paper claims that COOL-SD offers a principled and theoretically grounded approach to accelerating autoregressive image generation. Unlike previous heuristic methods, COOL-SD is supported by mathematical analysis of distributional fidelity. The authors assert that their method:

• Provides a tighter bound on distributional bias through optimal resampling.
• Achieves a superior speed-quality trade-off by leveraging an annealing acceptance schedule.
• Offers a modular foundation for future advances, as the derived resampling distributions can be seamlessly integrated into other relaxed SD variants (demonstrated by improvements when applied to LANTERN++).

The work positions COOL-SD as a significant step toward making AR image models more practical for real-time and interactive applications by effectively addressing the token ambiguity problem inherent in image generation.