来源：媒矿工厂本文共5000字，建议阅读10分钟
本文着重于潜在扩散系统的优化困境。
题目： Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models
作者： Jingfeng Yao, Bin Yang, Xinggang Wang
论文链接： http://arxiv.org/abs/2501.01423
来源：CVPR 2025
内容整理： 周楚骎具有 Transformer 架构的潜在扩散模型擅长生成高保真图像。然而，最近的研究揭示了这种两阶段设计中的优化困境：虽然在视觉分词器中增加每个标记的特征维度可以提高重建质量，但它需要更大的扩散模型和更多的训练迭代才能实现可比的生成性能。因此，现有系统经常满足于次优解决方案，要么由于分词器中的信息丢失而产生视觉伪影，要么由于昂贵的计算成本而无法完全收敛。本文认为，这种困境源于学习无约束高维潜在空间的固有困难。为了解决这个问题，本文建议在训练视觉分词器时，将潜在空间与预先训练的视觉基础模型对齐。本文提出的 VA-VAE（视觉基础模型对齐变分自动编码器）显著扩展了潜在扩散模型的重建生成前沿，使扩散变压器 （DiT） 在高维潜在空间中的收敛速度更快。为了充分利用 VA-VAE 的潜力，本文通过改进的训练策略和架构设计构建了增强的 DiT 基线，称为 LightningDiT。该集成系统在 的ImageNet 上实现了SOTA 性能，FID 得分为 1.35，同时在短短 64 个时期内就达到了 2.11 的 FID 得分，展示了卓越的训练效率——与原始 DiT 相比，收敛速度提高了21倍以上。
目录
- 简介
- 使 VAE 与 Vision Foundation 模型保持一致
- Marginal Cosine Similarity Loss
- Marginal Distance Matrix Similarity Loss
- 自适应权重
- Improved Diffusion Transformer
- 实验
- 实现细节
- 基础模型改善收敛性
- 基础模型提高可扩展性
- 消融和讨论
- 生成友好的VA-VAE
- 视觉基础模型
- 总结
简介
图1 潜在扩散模型中的优化困境
潜在扩散模型利用连续值变分自动编码器 （VAE）或视觉分词器来压缩视觉信号，从而减少高分辨率图像生成的计算需求。这些视觉分词器的性能，特别是它们的压缩和重建能力，在决定整体系统的有效性方面起着至关重要的作用。增强重建能力的直接方法是增加视觉标记的特征维度，从而有效地扩展 Latent 表示的信息容量。最近，一些有影响力的文本到图像作品与 Stable Diffusion 中广泛采用的 VAE 相比，探索了更高维的分词器，因为这些分词器提供了改进的细节重建，实现了更精细的生成质量。
然而，随着研究的深入，在潜在扩散模型中，重建和生成性能之间出现了一个优化困境。具体来说，虽然增加标记特征维度可以提高标记器的重建精度，但它会显著降低生成性能（见上图）。目前，有两种常见的策略可以解决这个问题：第一种涉及扩大模型参数，如 Stable Diffusion 3，它表明更高维的分词器可以通过明显更大的模型容量实现更强的生成性能——然而，这种方法需要更多的训练计算，这使得它对于大多数实际应用来说非常昂贵。第二种策略是故意限制分词器的重建能力，例如 Sana、W.A.L.T，以更快地收敛扩散模型训练。然而，这种受损的重建质量本身就限制了生成性能的上限，导致生成结果中的视觉细节不完美。这两种方法都涉及固有的权衡，并且无法有效控制潜在的优化困境。
本文提出了一种简单而有效的方法来解决这种优化困境。从自回归 （AR） 生成中汲取灵感，其中增加离散值 VAE 的码簿大小会导致码簿利用率低。通过可视化不同特征维度的潜在空间分布（见上图），本文观察到高维分词器以不那么分散的方式学习潜在表示，分布可视化中更集中的高强度区域证明了这一点。该分析表明，优化困境源于从头开始学习无约束高维潜在空间的固有困难。为了解决这个问题，本文为潜在扩散模型中的连续VAEs开发了一种视觉基础模型指导的优化策略。本文的主要发现表明，由视觉基础模型指导的学习潜在表征显著提高了高维分词器的生成性能，同时保留了其原始的重建能力（如下图所示）。
图2 潜在扩散模型的重建生成能力
本文的主要技术贡献是 Vision Foundation 模型对齐损失 （VF Loss），这是一个即插即用的模块，可在分词器训练期间将潜在表示与预先训练的视觉基础模型保持一致。而事实证明，使用预先训练的 Vision Foundation 模型天真地初始化 VAE 编码器是无效的——可能是因为潜在表示会迅速偏离其初始状态以优化重建——本文发现精心设计的关节重建和对齐损失至关重要。本文的对齐损失是专门为规范高维潜在空间而设计的，而不会过度限制它们的容量。首先，本文强制执行元素级和成对相似性，以确保特征空间中全局和局部结构的全面正则化。其次，本文在相似性成本中引入了一个边际，以提供受控的比对灵活性，从而防止过度正则化。此外，本文还研究了不同视力基础模型的影响。
为了评估生成性能，本文将拟议的 Vision 基础模型对齐 VAE （VA-VAE） 与扩散 Transformer （DiT）耦合起来，以创建潜在扩散模型。为了充分利用 VA-VAE 的潜力，本文通过先进的扩散训练策略和 Transformer 架构改进构建了一个增强的 DiT 框架，本文将其命名为 LightningDiT。本恩的贡献实现了以下重要的里程碑：
- 所提出的 VF Loss 有效解决了潜在扩散模型中的优化困境，使用高维分词器使 DiT 训练速度提高了 2.5× 以上;
- 集成系统仅用 64 个训练 epoch 就达到了 2.11 的 FID，与原始 DiT 相比，收敛速度提高了 21× 以上;
- 该集成系统在 ImageNet-256 图像生成方面实现了 1.35 的SOTA FID 分数。
使 VAE 与 Vision Foundation 模型保持一致
本节介绍了 VA-VAE，这是一种通过视觉基础模型对齐训练的视觉分词器。关键方法包括通过利用基础模型的特征空间来限制分词器的潜在空间，从而增强其对生成任务的适用性。
图3 提出的VA-VAE
如上图所示，本文的架构和训练过程主要遵循 LDM，采用具有连续潜在空间的 VQGAN 模型架构，受 KL 损失的约束。本文的主要贡献在于 Vision Foundation 模型对齐损失、VF 损失的设计，它有效地在不改变模型架构或训练管道的情况下，优化了潜在空间解决上文提到的优化困境。
VF 损失由两个部分组成：marginal cosine similarity loss 和 marginal distance matrix similarity loss。这些组件经过精心设计，是一个简单明了的即插即用模块，与 VAE 架构分离。
Marginal Cosine Similarity Loss
在训练过程中，给定的图像 均由视觉令牌的编码器和冷冻视觉基础模型处理，从而导致图像潜在的 和基础视觉表示。,投影 以使用线性变换匹配 的维度，其中，来产生。
损失函数 最小化相应特征 和 之间的相似性差距。对于每对，计算余弦相似性，并减去边缘 。 ReLU可确保只有低于 的成对有助于损失，从而集中在不太相似的对上。最终损失是在 的特征网格中的所有位置上平均的。
Marginal Distance Matrix Similarity Loss
补充 （强制点对点绝对比对），本文还旨在使特征中的相对分布距离矩阵尽可能相似。为此提出了边缘距离矩阵相似性损失。
距离矩阵相似性损耗使特征矩阵 和 的内部分布对齐。在这里， 表示每个扁平特征图中元素的总数。对于每对，计算特征矩阵 和 中相应向量之间余弦相似性差的绝对值，从而促进了其相对结构的更紧密比对。同样，减去边界 以放松约束。 ReLU功能可确保只有超过 的差异对损失的差异。
自适应权重
原始的重建损失和KL损失都是总和损失，这使VF损失完全不同量级，这使调整稳定训练的重量变得具有挑战性。受GAN损失的启发，采用了一种自适应加权机制。在反向传播之前，在编码器的最后一个卷积层上计算和的梯度。自适应加权设置为这两个梯度的比率，以确保和对模型优化具有相似的影响。这种比对大大减少了VF损耗的调整范围。
然后，本文将通过自适应加权获得VF损失。自适应加权的目的是快速对不同的VAE训练管道进行损失量表。在此基础上，本文仍然可以使用手动调整的超参数来进一步提高性能。
本文将评估VF损失在即将进行的实验中重建和产生的潜在扩散帕累托前沿中的重要作用。
Improved Diffusion Transformer
本节介绍了本文的LightningDiT。DiT 作为文本-图像和文本-视频任务的基础模型取得了巨大成功。但是，其收敛速度显著的慢，导致了高实验迭代成本。以前有影响力的工作Dinov2，Convnext 和Eva 展示了如何融合高级设计策略可以振兴经典方法。在本文的工作中，旨在扩大DiT体系结构的潜力，并探索DiT可以走多远的边界。虽然本文没有声称任何个人优化细节是最初贡献，但本文认为，开源、快速收敛的DiT训练管道将极大地支持社区对DiT的持续研究。
表1 DiT的表现
本文利用以 f8d4 为visual tokenizer的SD-VAE ，并使用DiT-XL/2作为实验模型。在表1中显示了优化例程。每个模型均已训练为80个epoch，并用dopri5 integrator进行采样，该模型的 NFE 比原始 DiT 少用于快速推断。为了确保进行公平的比较，不使用诸如CFG间隔和时间段偏移之类的样本优化方法。采用三类优化策略。在计算级别，实施torch.compile和bfloat16加速训练。此外，将批量的大小增加，并将 AdamW 的 降低到0.95，从先前的 Auraflow 中汲取了灵感。为了进行扩散优化，结合了Rectified Flow，logit正态分布（LogNorm）采样和速度方向损失。在模型体系结构级别，我们应用常用的 Transformer 优化，包括 RMSNorm，SWiGLU和RoPE。在训练期间，观察到某些加速策略不是正交的。例如，单独使用时梯度剪辑是有效的，但在对数点采样和速度方向损失后组合时倾向于降低性能。
本文优化的模型LightningDiT在Imagenet类条件生成上达到了7.13（CFG = 1）的FID，仅用80个epoch，仅是原始DiT和SiT所需的1400个epoch训练量的6％。以前的伟大工作MDT 或REPA 在蒙版图像建模（MIM）和表示对齐的帮助下达到了类似的收敛性能。我们的结果表明，即使没有任何复杂的培训管道，简单的DiT仍然可以取得非常具竞争力的表现。这种优化的体系结构在接下来的快速实验验证方面有很大帮助。
实验
本节主要目标是通过利用提出的VF损失来实现潜在扩散系统内重建和生成的重建和生成前沿。引入的LightningDit中，证明了VF损失如何有效地解决优化的障碍，从而从收敛性，可伸缩性和整体系统性能的角度来看有效地解决了优化的障碍。
实现细节
详细介绍了潜在扩散系统。对于visual tokenizer，采用主要遵循LDM的体系结构和培训策略。具体而言，利用 VQGAN 网络结构，省略量化并应用KL损失来调节连续的潜在空间。为了实现多节点训练，在MAR的设置上分别将学习率和全局批量规模扩展到1e-4和256。训练三个不同的 f16 tokenizer：一种没有VF损失，一种使用VF损失（MAE），另一种使用VF损失（Dinov2）。这里f表示下采样率，D表示潜在尺寸。根据经验，。对于生成模型，采用Lightningdit，进一步完善了概述的设计技术。从tokenizer中提取所有潜在特征，并以80或160个epoch的分辨率为256上的ImageNet上的LightningDiT的各种版本的LightningDiT。将DiT的patch size设置为1，以确保整个系统的下采样率为16。所有压缩步骤均由VAE处理。除非另有说明，否则我们模型的其他架构参数与DiT的架构参数一致。
基础模型改善收敛性
表2 VF loss提升性能
上表列出了对八种不同的象征器的重建和产生的评估，所有生成模型均经过160个 epoch（LightningDit-B）或80个 epoch（LightningDit-L＆LightningDit-XL）的训练。提供以下发现：
结果突出了潜在扩散系统中的优化困境。表中以蓝色突出显示的结果说明了重建性能（rFID）和相应的生成性能（FID）。可以观察到，随着tokenizer牌尺寸的增加，其rFID减小，而相应的生成FID增加。
VF损失可以有效地增强高维tokenizer的生成性能。在表的F16D32和F16D64部分中，VF损失（Dinov2）和VF损失（MAE）都显着改善了不同尺度上DiT模型的生成性能。这使得以更高的重建性能和更高的生成性能（即引言中提到的重建生成前沿）实现系统。但是，值得注意的是，对于低维数字的VF损失是不必要的，例如通常使用的 f16d16。这与图1中的潜在分布观察一致。我们认为这是因为较低维空间可以学习更多合理的分布，而无需其他额外的监督信号。
图4 VF Loss加速收敛
此外，上图中提出了FID的收敛图。在f16d32和d16f64上，使用VF损失的使用分别按2.54和2.76的倍数加速。这些还表明，VF损失显著提高了高维tokenizer的生成性能和收敛速度。
基础模型提高可扩展性
增加模型参数计数是提高高维tokenizer生成性能的一种方法。使用尺寸为0.1b至1.6B的LightningDiT模型来评估3种不同的tokenizer的生成性能。为了促进缩放中幂律的观察，使用轴的对数尺度。注意到，随着参数数量的增加，蓝色和绿色线之间存在轻微的收敛趋势，但仍然存在很大的差距。这意味着，高维f16d32 tokenizer 对产生的负面影响即使在1.6B处也没有完全缓解，这是一个已经被认为是在Imagenet上被认为很重要的参数大小。发现VF损失有效地弥合了这一差距。低于0.6B，橙色和蓝线的性能相似。但是，随着模型尺度超过1B，f16d32 VF DINOV2逐渐与f16d16差距变大，表明可扩展性更强。
消融和讨论
在本节中，对VF损失的设计进行消融实验，以评估各种基础模型和损失公式的影响。然后们对VF损失的潜在机制进行更深入的分析，从而提供可能有所帮助的其他见解。
生成友好的VA-VAE
与patch size为2的SD-VAE相比，patch size为1的VA-VAE具有出色的生成性能。用VA-VAE代替了SD-VAE [34]，导致FID-50K从7.13降低到7.13至4.29。这种改进可以归因于两个主要原因。首先，观察到使用f16训练的DiT与使用f8的DIT相比，使用f16且patch size为1的DiT训练比 f8 patch size为2的DIT更容易收敛。
视觉基础模型
表3 基础模型消融
我们使用三种类型的基础模型训练VA-VAE：具有掩盖图像模型的自我监督模型，图像文本对比度学习模型CLIP和SAM。为了加速收敛，分别将学习率和全球批量大小调整为1e-4和256。与以前的设置相反，每个tokenizer在ImageNet 256×256上进行了50个epoch的训练。对于每个tokenizer，在相应的潜在空间中训练LightningDit-B，以160个epoch的速度训练。上表总结了我们的发现，表明所有这些视觉基础模型都增强了扩散模型的生成性能。其中，自我监管的预训练的模型DinoV2取得了卓越的生成结果。
总结
本文着重于潜在扩散系统的优化困境。为了解决这个问题，提出了与视觉基础模型保持一致的VA-VAE，以及一项优化的DIT，其中包括了高级设计策略。在VA-VAEe中，VF损失结合边缘余弦的相似性和距离矩阵损失 - 与视觉模型保持一致，从而导致更均匀的特征分布和高达2.8倍更快的收敛性。通过LightningDiT，整合了先进的训练技术和架构改进，以实现更快的DIiT收敛。将VA-VAE（rFID = 0.28）的高重建能力与LightningDit的快速收敛结合在一起，本文的方法在Imagenet 256上实现了1.35的最新FID。此外，本文的方法可实现2.11 FID，仅使用64个epoch，证明了21.8倍提速。
