Annealed Relaxation of Speculative Decoding for Faster Autoregressive Image Generation¶

会议: AAAI 2026
arXiv: 2601.09212v1
代码: https://github.com/xingyaoL/COOL-SD (有)
领域: 图像生成 / 高效推理
关键词: speculative decoding, 自回归图像生成, 退火松弛, 总变差距离, 推理加速

一句话总结¶

提出Cool-SD，一种有理论支撑的退火松弛speculative decoding框架：通过推导TV距离上界得到最优重采样分布，并证明接受概率递减调度比均匀调度产生更小的分布偏移，在LlamaGen和Lumina-mGPT上实现了比LANTERN++更优的速度-质量权衡。

背景与动机¶

自回归(AR)图像生成模型（LlamaGen、Lumina-mGPT等）质量媲美扩散模型，但逐token生成导致推理很慢。Speculative Decoding(SD)是LLM领域成熟的加速技术，用小的draft模型生成候选token序列，再由target模型并行验证。然而SD在图像生成中效果有限，原因是图像token的语义模糊性——很多位置有多个概率相近的候选token，导致draft模型难以命中target模型的选择，接受率低。

LANTERN和LANTERN++通过松弛验证条件（允许近邻token被接受）来提升接受率，但它们的设计都是启发式的，缺乏理论基础来保证松弛后的分布与目标分布的接近程度。而且它们的重采样分布也不是最优的，导致分布偏移难以控制。

核心问题¶

如何在松弛speculative decoding的验证条件以提高接受率的同时，有理论保证地最小化输出分布与目标分布之间的偏差？ 需要同时解决两个子问题：给定松弛的接受标准，什么是最优的重采样分布？接受概率应该如何随位置变化？

方法详解¶

整体框架¶

Cool-SD在标准SD框架上做两处改进： - 输入: target模型 \(P\)、draft模型 \(Q\)、prefix、draft长度 \(L\) - 修改1: 将接受概率 \(f_i\) 从vanilla SD的严格匹配替换为乘以松弛因子 \(\omega_i\) 的版本，且 \(\omega_i\) 按指数衰减 - 修改2: 当token被拒绝时，使用推导出的最优重采样分布 \(G_i^*\) 来纠正偏差 - 输出: 生成的token序列，接近目标分布且接受长度更长

关键设计¶

TV距离上界与最优重采样分布 (Theorem 2): 对于任意松弛接受标准 \(\{f_i\}_{i=1}^L\)，推导了输出分布 \(\hat{P}\) 与目标分布 \(P\) 之间的TV距离几乎紧的上界。最小化这个上界得到每个位置的最优重采样分布：\(G_{i+1}^* = \text{Norm}([P(\cdot|x_{1:i}) - Q(\cdot|x_{1:i}) \cdot f_{i+1}((x_{1:i}, \cdot), P, Q)]_+)\)。当接受概率不低于vanilla SD时，这个最优分布恰好与vanilla SD的重采样分布一致（Proposition 1），实现简洁。
退火性质与指数衰减调度 (Proposition 2): 通过扰动分析证明了一个退火性质：在固定期望接受token数的条件下，前面位置松弛多、后面位置松弛少比反过来产生更小的TV距离上界。基于此设计指数衰减调度 \(\omega_i = \delta \cdot \exp(-\nu \cdot i - \mu)\)，其中 \(\delta\) 控制松弛总量，\(\nu = 0.7\) 控制衰减速度。
参数化接受标准: \(f_i(x_{1:i}, P, Q; \omega_i) = \min\{1, \omega_i \cdot P(\tilde{x}_i | x_{1:i-1}) / Q(\tilde{x}_i | x_{1:i-1})\}\)，当 \(\omega_i > 1\) 时放宽接受条件，\(\omega_i = 1\) 退化为vanilla SD。\(\delta\) 是唯一需要调节的超参数，控制速度-质量权衡。

损失函数 / 训练策略¶

Cool-SD本身不需要额外训练，只改变SD的验证和重采样策略。Draft模型训练遵循Eagle流程。实现比LANTERN++更简洁，无需近邻token的概率检索和聚合。

实验关键数据¶

Target模型	方法	CLIP↑	FID↓	IR↑	接受长度↑	延迟(s)↓	加速比↑
Lumina-mGPT (7B)	Target Model	0.3330	28.99	0.6855	1.00	170.14	1.00×
Lumina-mGPT	Eagle-1 (无松弛)	0.3330	29.05	0.6883	2.76	71.66	2.37×
Lumina-mGPT	LANTERN++	0.3328	30.31	0.6697	2.99	68.64	2.48×
Lumina-mGPT	Cool-SD (δ=1.1)	0.3325	30.30	0.6699	3.11	63.24	2.69×
LlamaGen-XL (775M)	Eagle-1	0.3157	20.97	-0.0859	2.42	4.99	2.03×
LlamaGen-XL	LANTERN++	0.3157	21.17	-0.1155	2.67	4.70	2.15×
LlamaGen-XL	Cool-SD (δ=1.1)	0.3167	21.02	-0.0997	2.73	4.46	2.27×
LlamaGen-XL	Cool-SD (δ=2)	0.3154	21.20	-0.1353	3.34	3.72	2.72×

消融实验要点¶

退火 vs 均匀松弛: Cool-SD（指数衰减）在FID-接受长度曲线上全面优于UniformRSD（均匀松弛），验证了退火性质的实际效果
最优重采样分布对LANTERN++的提升: 将 \(G_i^*\) 应用到LANTERN++上，在不同 \(k\) 设置下都改善了LANTERN++的表现，且 \(k\) 越大改善越明显（因为LANTERN++自身的重采样偏差随 \(k\) 增大而增大）
线性 vs 指数衰减: 在LlamaGen上指数调度优于线性调度，在Lumina-mGPT上两者接近，均优于均匀调度
与SJD对比: Cool-SD (δ=1.1) 在Lumina-mGPT上达到2.69×加速 vs SJD的2.06×，质量相当
统计显著性: Cool-SD vs LANTERN++的接受长度差异在1000样本上 \(t=18.11, p < 10^{-60}\)

亮点¶

理论与实践完美结合: 不是先有方法再补理论，而是从TV距离上界出发推导最优重采样分布，从扰动分析得到退火性质，再据此设计算法。理论直接指导设计
最优重采样分布的即插即用性: \(G_i^*\) 可以直接提升LANTERN++等已有方法的性能，是一个通用的改进组件
退火直觉的理论基础: "前松后紧"的接受策略有直觉上的合理性（前面位置影响范围更大，放宽后面位置收益更小），论文给出了严格证明
实现更简洁: 与LANTERN++需要构建k-近邻和概率聚合不同，Cool-SD只需修改接受阈值和重采样分布，额外开销几乎为零

局限性 / 可改进方向¶

超参数 \(\delta\) 仍需人工调节: 虽然只有一个超参数，但最优值因模型和应用场景不同，无法自适应选择
退火性质的严格证明仅限 \(L=2\): 虽然论文声称可通过逐对比较推广到 \(L > 2\)，但正式证明仅覆盖两步情况
假设draft和target分布接近 (Assumption 2: TV ≤ 2/5): 当draft模型质量差时理论保证可能减弱
仅测试图像生成: 理论框架适用于任何AR任务，但实验仅在图像生成上验证，LLM文本生成上的效果未知
\(\delta\) 增大时质量退化: 在高加速比（>3.5×）下图像质量明显下降，无法无限制地提速

与相关工作的对比¶

方法	核心思路	与Cool-SD的关键差异
Vanilla SD (Eagle-1)	严格匹配的无偏解码	Cool-SD通过可控松弛换取更高加速，有理论界保证偏差程度
LANTERN	基于k-近邻的概率聚合松弛	启发式设计无理论保证，重采样分布非最优
LANTERN++	LANTERN + 静态树结构	Cool-SD理论更强，实现更简单（无需近邻检索），Pareto前沿更优
SJD	Jacobi解码+speculative	无损但加速比有限(2.06× vs Cool-SD 2.69×)

启发与关联¶

最优重采样分布的推导思路可以迁移到其他需要近似采样的场景（如guided diffusion中的重采样）
退火调度的思想与模拟退火、温度调度在优化中的应用有类似精神——"前松后紧"是一种通用的启发
理论框架本身适用于任何AR模型的speculative decoding加速，不限于图像

评分¶

新颖性: ⭐⭐⭐⭐ TV距离上界推导和退火性质证明是实质性的理论贡献，但整体改进方向并非全新
实验充分度: ⭐⭐⭐⭐ 两个target模型、多种比较方法、消融分析全面，但缺少非图像域的验证
写作质量: ⭐⭐⭐⭐⭐ 理论推导严谨，实验展示清晰，理论到实践的过渡自然
价值: ⭐⭐⭐⭐ 为松弛SD提供了首个理论框架，最优重采样分布的即插即用性对社区有直接价值