通过教师对齐的端到端蒸馏实现高保真两步图像生成 High-Fidelity Two-Step Image Generation via Teacher-Aligned End-to-End Distillation

方法的第一步是准备8步教师模型。作者假设已经通过已建立的蒸馏技术（如解耦DMD）获得一个8步Z-Image Turbo教师模型，这一阶段的生产已经很成熟，不进一步讨论。第二步是执行分布对齐对抗学习。在这一步中，作者不是使用外部高质量数据集的图像作为GAN判别器的真实样本，而是使用8步教师模型生成的输出作为真实样本。具体而言，判别器$D$区分"真实"和"虚假"图像，并为生成器提供对抗梯度：$\mathcal{L}_{GAN} = \mathbb{E}_{x \sim p_{real}}[\log D(x)] + \mathbb{E}_{x \sim p_{fake}}[\log(1 - D(x))]$。通过使用教师生成的图像作为真实样本，训练变得更加稳定，最终生成质量也得到显著提升。第三步是实施步骤解耦参数化。两个去噪步骤的模型参数被完全解耦，两个步骤的模型权重都从8步教师的权重初始化，然后独立训练。这实际上翻倍了专门用于2步生成任务的模型容量。第四步是进行带迭代正则化的端到端训练。整个梯度路径——从初始噪声通过步骤1模型、中间表示、步骤2模型，到最终输出和损失——是简洁和可追踪的，使完全跟踪通过两个步骤的梯度成为可能。为了避免完全连接的两步图的峰值内存成本，作者首先通过步骤1预测的分离克隆传播第二分支，存储结果梯度，然后通过继承损失将其注入步骤1模型，该继承损失加到步骤1的原始局部损失上。在实践中，作者将继承项缩放一个小权重（0.1）以防止梯度爆炸。最后，总体训练目标结合了分布对齐对抗损失和DMD损失：$\mathcal{L} = \mathcal{L}_{GAN} + \lambda \mathcal{L}_{DMD}$。

本文的技术新颖性体现在三个方面。首先，在对抗学习的真实样本选择上，本文首次系统性地研究了使用教师生成图像与外部真实图像的差异，并从分布对齐的角度解释了为什么教师生成样本能够提供更稳定和可达到的优化路径。作者观察到，当使用8步教师输出作为真实样本时，生成器的GAN损失表现出健康的模式：随着判别器增强，它最初增加，然后随着生成器成功关闭分布差距而平台化。相比之下，使用外部真实图像时，损失不仅高得多，而且在整个训练过程中持续增长，表明学生无法克服的不可逾越的分布分歧。其次，在模型参数化方面，本文系统地比较了全权重解耦、共享权重和每步LoRA三种方案，并通过实验证据表明全权重解耦在最具容量需求和文本导向的指标上明显优于其他方案。作者测试了较弱形式的解耦：两个步骤共享同一主干但使用任务特定的LoRA模块，该变体在GAN损失方面表现不佳，最终损失甚至高于共享模型的完全微调。最后，在端到端训练方面，本文首次在2步设置下实现了完全的端到端训练，并系统性地研究了步骤1损失的必要性。作者的实验表明，即使在端到端训练下，单独的步骤1损失仍然是不可或缺的，去除它会导致步骤2生成器的GAN损失激增，并伴有生成质量的明显退化。作者从迁移学习的角度解释了这一行为：蒸馏模型的最终性能取决于下游任务的学习能力和利用预训练期间积累的知识的能力，而扩散模型拥有 deeply ingrained 的迭代性质，去除步骤1损失会破坏预训练权重中编码的归纳偏差。

独立分析的弱点主要集中在四个方面。首先，参数存储开销显著增加，这对于移动端和边缘设备部署是一个实际障碍。在内存受限的环境中，存储两个独立的模型（每个都从6B参数的教师初始化）可能是不可行的。一个改进方向是探索更高效的参数共享方案，如混合专家（MoE）架构或轻量级适配器，在保持解耦的容量优势的同时减少存储开销。其次，在密集文本渲染和复杂场景生成方面，2步模型仍然不如8步教师可靠。这可能部分归因于2步生成的有限迭代深度，这使得模型难以捕捉细粒度的细节和复杂的空间关系。一个改进方向是引入辅助任务或损失函数，专门针对文本渲染和复杂场景推理进行训练，或者使用渐进式细化策略，在关键区域增加额外的推理步骤。第三，端到端训练虽然有效，但可能仍然受到梯度传播的限制，特别是在长程依赖方面。一个改进方向是探索更复杂的梯度传播机制，如分层梯度流或多尺度特征融合，以改善两个步骤之间的协调。最后，训练稳定性和收敛速度可能仍然是一个挑战，特别是在大规模部署中。一个改进方向是设计更稳定的训练策略，如课程学习或自适应损失加权，以改善训练动态并减少对超参数的敏感性。

未来研究方向包括作者提出的和基于成果可延伸的多个方向。作者明确指出的方向包括在质量、存储和部署效率之间找到更好的平衡，这是参数解耦带来的额外存储成本的一个重要挑战。一个具体方向是探索更高效的参数化方案，如混合专家架构、轻量级适配器或知识蒸馏技术，在保持解耦的容量优势的同时减少存储开销。作者还提到，本文的端到端训练框架可能为基于强化学习的少步生成器优化提供有用基础，这是一个有前景但未探索的方向。基于本文成果，可延伸的另一个研究方向是将这些技术应用于其他步数压缩场景，如3步或4步生成，或者应用于其他模态（如视频、3D资产）的生成。此外，可以探索更复杂的教师-学生蒸馏策略，如多教师蒸馏或渐进式蒸馏，其中学生首先学习与4步教师对齐，然后与2步教师对齐。另一个方向是研究更精细的重噪调度策略，超越本文使用的约束调度$\tau_{CA} = \tau_{DM} > t$，以更好地平衡CFG增强和分布匹配的作用。最后，可以探索在推理时自适应地选择步数，根据输入复杂度和质量要求在2步到8步之间动态调整，以在效率和质量之间实现更好的权衡。

关于复现性的评估，本文提供了相对详细的实验设置和结果，但开源情况需要进一步确认。作者在附录中提供了端到端训练实现的伪代码，包括内存优化的关键细节，如梯度检查点和FSDP的使用，这有助于复现者理解如何在有限的内存内实现端到端训练。实验使用了多个公开基准测试（DPGBench、GenEval、OneIGBench、LongTextBench），所有提示词直接来自官方基准测试集，没有提示增强，这确保了结果的公平性和可复现性。消融实验系统地研究了每个组件的贡献，包括GAN目标分布（教师生成 vs 外部真实图像）、参数解耦（全权重解耦 vs 共享权重 vs 每步LoRA）和步骤1损失（有 vs 无），这些详细的消融分析有助于理解每个组件的作用并指导复现。然而，论文没有明确提供代码、模型权重或训练脚本的开源链接，也没有详细说明训练所需的计算资源（如GPU数量、训练时间）。考虑到模型规模（6B参数教师）和端到端训练的计算成本，复现可能需要 substantial 的计算资源。此外，论文没有提供数据集的详细信息，虽然提到使用教师生成的图像作为真实样本，但没有说明生成了多少样本、使用的提示词分布等。总体而言，虽然论文提供了相对详细的方法描述和实验分析，但缺乏代码和训练细节的开源可能会增加复现的难度。