One-Step Distillation of Discrete Diffusion Image Generators via Fixed-Point Iteration
作者: Chaoyang Wang, Yunhai Tong
分类: cs.CV
发布日期: 2026-05-20
💡 一句话要点
提出固定点蒸馏(FPD)框架,实现离散扩散图像生成器单步高效蒸馏。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 离散扩散模型 图像生成 模型蒸馏 单步生成 固定点迭代
📋 核心要点
- 离散扩散模型生成图像质量高,但迭代解码速度慢,单步蒸馏方法存在计算量大或优化困难的问题。
- 提出固定点蒸馏(FPD)框架,通过学生网络单步草图和教师网络单步细化构建局部校正目标,实现高效蒸馏。
- 实验表明,FPD在单步推理中实现了有竞争力的视觉保真度和结构对齐,优于现有离散蒸馏基线。
📝 摘要(中文)
离散扩散模型在视觉合成方面表现出色,但依赖于缓慢的迭代解码过程。现有的单步蒸馏方法试图绕过这个瓶颈,但要么通过训练辅助评分网络来有效地加倍计算量,要么引入专门的参数化和多阶段流程来分割优化。本文提出了一种固定点蒸馏(FPD)的端到端框架,该框架通过部分破坏学生网络的单步草图,并使用教师网络进行单步细化,来构建局部校正目标。为了在语义上有意义的空间中计算训练目标,我们将离散token提升到连续特征,并应用多带宽漂移损失来迭代累积这些校正。为了通过离散瓶颈进行反向传播,我们采用直通估计器,在正向传播过程中将精确的硬采样token馈送到教师网络和解码器,确保训练和推理在相同的码本流形上运行,同时将连续梯度路由回学生网络的logits。这种完全可微的路径还可容纳可选的无条件对抗目标,以增强感知真实感。在类条件和文本条件生成上的评估验证了我们框架的有效性。FPD在单步推理中实现了具有竞争力的视觉保真度和结构对齐,缩小了与多步教师网络之间的差距,同时优于现有的离散蒸馏基线。
🔬 方法详解
问题定义:离散扩散模型虽然在图像生成方面表现出色,但其迭代式的解码过程非常耗时。现有的单步蒸馏方法试图加速这一过程,但通常需要训练额外的评分网络,导致计算量翻倍,或者需要引入复杂的参数化和多阶段流程,使得优化过程变得困难。
核心思路:FPD的核心思路是通过模仿教师网络的一步精炼过程来训练学生网络。具体来说,学生网络先生成一个初步的图像草图,然后通过教师网络进行一次精炼,得到一个更接近真实图像的结果。FPD的目标是让学生网络直接生成这个精炼后的结果,从而实现单步生成。
技术框架:FPD框架主要包含以下几个模块:学生网络、教师网络、离散token到连续特征的映射、多带宽漂移损失和直通估计器。学生网络负责生成初始草图,教师网络负责精炼草图。离散token被映射到连续特征空间,以便计算损失。多带宽漂移损失用于迭代累积校正。直通估计器用于在离散瓶颈处进行反向传播。
关键创新:FPD的关键创新在于其局部校正目标的构建方式和直通估计器的使用。通过部分破坏学生网络的单步草图,并使用教师网络进行单步细化,FPD能够构建一个更有效的局部校正目标。直通估计器的使用保证了训练和推理在相同的码本流形上运行,避免了离散瓶颈带来的梯度消失问题。
关键设计:FPD使用多带宽漂移损失来迭代累积校正,该损失函数在连续特征空间中计算,能够更好地捕捉图像的语义信息。此外,FPD还使用了直通估计器,将精确的硬采样token馈送到教师网络和解码器,同时将连续梯度路由回学生网络的logits。这种设计保证了训练和推理的一致性,并避免了梯度消失问题。
🖼️ 关键图片
📊 实验亮点
实验结果表明,FPD在类条件和文本条件图像生成任务上均取得了显著的性能提升。在单步推理中,FPD实现了与多步教师网络相近的视觉保真度和结构对齐,同时优于现有的离散蒸馏基线。具体性能数据未知,但论文强调FPD缩小了单步模型与多步模型之间的差距。
🎯 应用场景
FPD可应用于各种需要快速图像生成的场景,如实时图像编辑、游戏开发、虚拟现实等。通过单步生成高质量图像,可以显著提高用户体验和系统效率。此外,FPD还可以作为一种通用的模型压缩技术,用于减小离散扩散模型的模型大小和计算复杂度。
📄 摘要(原文)
Discrete diffusion models excel at visual synthesis but rely on slow, iterative decoding. Existing single-step distillation methods attempt to bypass this bottleneck, either by training auxiliary score networks that effectively double compute, or by introducing specialized parameterizations and multi-stage pipelines that fragment optimization. In this paper, we introduce Fixed-Point Distillation (FPD), an end-to-end framework that constructs local correction targets by partially corrupting the student's one-step draft and refining it with a single teacher step. To compute the training objective in a semantically meaningful space, we lift discrete tokens into continuous features and apply a multi-bandwidth drift loss that iteratively accumulates these corrections. To backpropagate through the discrete bottleneck, we employ a straight-through estimator that feeds exact hard-sampled tokens to the teacher and decoder during the forward pass, ensuring that training and inference operate on the same codebook manifold, while routing continuous gradients back to the student logits. This fully differentiable pathway additionally accommodates an optional unconditional adversarial objective to enhance perceptual realism. Evaluations on both class- and text-conditional generation validate the effectiveness of our framework. FPD achieves competitive visual fidelity and structural alignment within a single inference step, narrowing the gap to multi-step teachers while outperforming existing discrete distillation baselines.