重建与生成之间的权衡：在潜空间扩散模型中驯服优化困境

paper是华发表在CVPR 2025的工作

paper title：Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

Code：链接

Abstract

采用 Transformer 架构的潜空间扩散模型在生成高保真图像方面表现突出。然而，近期研究揭示了这种两阶段设计中的一个优化困境：在视觉分词器中增加每个 token 的特征维度虽然能够提升重建质量，但同时需要更大规模的扩散模型以及更多训练迭代，才能获得与之相当的生成性能。因此，现有系统往往不得不在次优解之间折中：要么由于分词器中的信息损失而产生视觉伪影，要么由于计算成本过高而无法充分收敛。我们认为，这一困境源于在无约束的高维潜空间中进行学习的固有难度。为了解决这一问题，我们提出在训练视觉分词器时，将潜空间与预训练视觉基础模型对齐。我们提出的 VA-VAE（Vision foundation model Aligned Variational AutoEncoder）显著拓展了潜空间扩散模型在重建–生成性能上的前沿，使得 Diffusion Transformers（DiT）能够在高维潜空间中更快速地收敛。为了充分发挥 VA-VAE 的潜力，我们构建了一个具有改进训练策略和架构设计的增强版 DiT 基线，称为 LightningDiT。该集成系统在 ImageNet 256×256 图像生成任务上取得了最先进（SOTA）性能，FID 为 1.35；同时展现出惊人的训练效率，仅用 64 个 epoch 即达到 FID 2.11 —— 相较于原始 DiT 提供超过 21× 的收敛加速。

1. Introduction

潜空间扩散模型（latent diffusion model）[34] 使用连续值变分自编码器（VAE）[19]，或称视觉分词器，来压缩视觉信号，从而降低高分辨率图像生成的计算需求。视觉分词器的性能，尤其是其压缩与重建能力，在很大程度上决定了整个系统的有效性 [6, 8]。一种增强重建能力的直接方法是增加视觉 token 的特征维度，从而有效扩展潜表示的信息容量。近年来，多篇具有影响力的文本生成图像工作 [6, 8, 21] 探索了比 Stable Diffusion [31, 34] 广泛采用的 VAE 更高维度的分词器，因为高维分词器提供更好的细节重建能力，从而实现更高质量的生成效果。

然而，随着研究的深入，在潜空间扩散模型中，重建与生成性能之间出现了一个优化困境 [8, 15, 18]。具体而言，虽然增加 token 特征维度能够提升分词器的重建精度，但它会显著降低生成性能（见图 1）。当前有两类常见策略应对这一问题：第一类是扩大模型参数规模，例如 Stable Diffusion 3 [8] 展示了更高维分词器在模型容量大幅增大时可以达到更强的生成性能——然而，这种方法需要极高的训练计算量，使其在大多数实际应用中成本过高。第二类策略是有意限制分词器的重建能力，例如 Sana [4, 42]、W.A.L.T [15]，以便加速扩散模型训练的收敛。但受限的重建质量天然限制了生成性能的上限，导致生成结果中的视觉细节不够完美。这两种方法都存在内在折中，无法有效控制潜在的优化困境。

图 1. 潜空间扩散模型中的优化困境。 在潜空间扩散模型中，提高视觉分词器的特征维度能够增强细节重建能力，但会显著降低生成质量。（在分词器的规格说明中，“f” 与 “d” 分别表示下采样率和维度。所有结果均在 ImageNet 256×256 数据集上评估，且扩散模型训练的计算量保持固定。）

在本文中，我们提出了一种简单而有效的方式来解决这一优化困境。我们的灵感来源于自回归（AR）生成领域，其中增大离散值 VAE 的码本规模会导致低码本使用率 [44, 50]。通过可视化不同特征维度下的潜空间分布（见图 1），我们观察到高维分词器学习到的潜表示分布更不分散，其分布可视化中存在更集中、更高强度的区域。该分析表明，优化困境源于从零开始学习无约束高维潜空间的固有困难。为此，我们为潜扩散模型中的连续 VAE [19] 设计了一种由视觉基础模型引导的优化策略。我们的关键发现指出，借助视觉基础模型引导学习潜表示，能够显著提升高维分词器的生成性能，同时保持其原有的重建能力（见图 2）。

图 2. 潜空间扩散模型的重建–生成前沿。 VA-VAE 改善了高维潜表示的特征分布。通过与视觉基础模型的对齐，我们拓展了潜空间扩散模型中重建与生成之间的前沿。

我们的主要技术贡献是 Vision Foundation model alignment Loss（VF Loss），一种可插拔模块，用于在训练分词器时将潜表示对齐到预训练视觉基础模型 [17, 29]。尽管直接用预训练视觉基础模型初始化 VAE 编码器被证明无效 [23]——可能是因为潜表示会迅速偏离初始状态以优化重建——我们发现精心设计的联合重建与对齐损失至关重要。我们的对齐损失专为在不过度约束容量的前提下，对高维潜空间进行正则化而设计。首先，我们同时施加逐元素与成对相似性约束，以保证在特征空间中对全局与局部结构的全面正则化。其次，我们在相似性代价中引入 margin，以提供受控的对齐灵活性，从而避免过度正则化。此外，我们还探讨了不同视觉基础模型的影响。

为评估生成性能，我们将提出的 Vision foundation model Aligned VAE（VA-VAE）与 Diffusion Transformers（DiT）[30] 配对，构建潜空间扩散模型。为了充分发挥 VA-VAE 的潜力，我们基于先进的扩散训练策略和 Transformer 架构改进，构建了增强版 DiT 框架，我们称之为 LightningDiT。我们的贡献达成了以下重要成果：

• 所提出的 VF Loss 有效解决了潜扩散模型中的优化困境，使高维分词器的 DiT 训练速度提升超过 2.5×；
• 集成系统在仅 64 个训练 epoch 内达到 FID 2.11，相比原始 DiT 实现了超过 21× 的收敛加速；
• 集成系统在 ImageNet-256 图像生成上取得 SOTA 的 FID 1.35。

2. Related Work

2.1. Tokenizers for Visual Generation

---

以变分自编码器（VAE）[19] 为代表的视觉自编码器利用编码器–解码器结构来生成连续的视觉表示，从而实现视觉信号的压缩与重建。VQVAE [40] 首次引入码本（codebook），将 VAE 的潜表示量化为离散格式。在此基础上，VQGAN [7] 通过加入 GAN 损失提升了重建质量，并成为自回归（AR）生成模型 [2] 的关键组成部分。然而，近期工作 [14, 44] 指出，码本利用率的不足会带来意外的折中：尽管更大的码本提升了重建保真度，但却会削弱生成性能。对此问题，Look-up Free Quantization（LFQ）被提出用于改善该现象。VQGAN-LC [50] 也识别了这一挑战，但采取了另一种方法：利用 CLIP [33] 视觉基础模型初始化码本，使码本利用率提升至 99%。基于连续 VAE 的潜空间扩散系统 [34] 也报告了类似的现象：随着分词器维度增大，重建质量提升，但生成性能下降，同时训练成本显著增加 [4, 8, 23]。目前，对这一问题的细粒度分析与针对性解决方案仍然不足。

2.2. Fast Convergence of Diffusion Transformers

---

Diffusion Transformers（DiT）[30] 是当前视觉生成任务中最受欢迎的基础扩散模型 [3, 8, 11, 26, 28, 49]。由于其卓越的可扩展性，它们在多种文本生成图像（text-to-image）和文本生成视频（text-to-video）任务中都表现出强大的效果。然而，DiT 的收敛速度较慢。相关研究如 SiT 通过引入 Rectified Flow 来提升 DiT 的效率，而 MDT [12] 和 MaskDiT [47] 则将掩码图像建模与扩散结合，以实现更快的收敛。类似地，REPA [45] 使用表示对齐方法，在训练过程中将 DiT 的特征对齐到 DINOv2 [29] 等基础 MIM 模型，以加速其收敛。然而，无论是引入 MIM 还是特征对齐技术，都会不可避免地增加训练成本，并提高整个系统的复杂度，尽管这些方法在更少的 epoch 中实现了更快的收敛。因此，一些先进的自回归（AR）模型 [22, 39] 开始展现出更具竞争力的性能。基于 DINOv2 [29]、ConvNeXt [25] 与 EVA [10] 等重要工作的成功经验——这些工作展示了将先进设计技巧融入经典方法能够赋予其新的生命力 [16, 17, 48]——我们希望进一步挖掘 DiT 架构的潜力。我们的思路是在不引入复杂设计的条件下优化 DiT，从而探索其全部能力边界。

3. Align VAE with Vision Foundation Models

在本节中，我们介绍 VA-VAE，这是一种通过视觉基础模型对齐训练的视觉分词器。其关键方法是利用基础模型的特征空间来约束分词器的潜在空间，从而提升其在生成任务中的适用性。如图 3 所示，我们的架构和训练过程主要遵循 LDM [34]，采用具有连续潜在空间的 VQGAN [7] 模型架构，并由 KL 损失进行约束。我们的关键贡献在于设计了视觉基础模型对齐损失 VF 损失，它在不改变模型架构或训练流程的前提下，大幅优化了潜在空间，有效解决了第 1 节中讨论的优化困境。

VF 损失由两个部分组成：边际余弦相似度损失和边际距离矩阵相似度损失。这两个组件被精心设计为一个简单的、可即插即用的模块，并与 VAE 架构解耦。我们将在下文详细解释每个部分。

图 3. 所提出的视觉基础模型对齐 VAE（VA-VAE）。视觉基础模型用于引导高维视觉分词器的训练，有效缓解优化困境并提升生成性能。

3.1. Marginal Cosine Similarity Loss

---

在训练过程中，给定图像 $I$ 会被视觉分词器的编码器和冻结的视觉基础模型同时处理，从而得到图像潜变量 $Z$ 和基础视觉表征 $F$。如公式 1 所示，我们使用一个固定投影 $F$，通过线性变换来匹配 $Z$ 的维度，其中 $W\in {\mathbb{R}}^{d_z\times d_f}$ 和 $b\in {\mathbb{R}}^{d_z}$ 分别是权重矩阵和偏置项，从而生成 $F^{\prime }\in {\mathbb{R}}^{d_z}$。

$$F^{\prime }=WF+B(1)$$

如公式 2 所定义，损失函数 ${\mathcal{L}}_{mcos}$ 最小化特征矩阵 $Z$ 与 $F^{\prime }$ 在每个空间位置 $(i,j)$ 的对应特征 $z_{ij}$ 和 $f_{ij}^{\prime }$ 之间的相似度差距。对于每一对特征，我们计算余弦相似度 $\frac{z_{ij}\cdot f_{ij}^{\prime }}{\Vert z_{ij}\Vert \Vert f_{ij}^{\prime }\Vert }$，并减去一个边际值 $m_1$。ReLU 函数确保只有相似度低于 $m_1$ 的特征对才会对损失产生贡献，从而将对齐重点放在相似度较低的特征对上。最终损失在 $h\times w$ 的特征网格上的所有位置取平均。

$$
\mathcal{L}_{mcos}

\frac{1}{h \times w}
\sum_{i=1}^{h}
\sum_{j=1}^{w}
\mathrm{ReLU}
\left(
1 - m_1 -
\frac{z_{ij} \cdot f’{ij}}{|z{ij}||f’_{ij}|}
\right)
\quad (2)
$$

3.2. Marginal Distance Matrix Similarity Loss

---

作为对 ${\mathcal{L}}_{mcos}$ 的补充（该损失强调点对点的绝对对齐），我们还希望特征内部的相对分布距离矩阵尽可能相似。为此，我们提出了边际距离矩阵相似度损失。

在公式 3 中，距离矩阵相似度损失用于对齐特征矩阵 $z$ 和 $f$ 的内部分布。这里 $N=h\times w$ 表示每个展开后的特征图中的元素总数。对于每一对位置 $(i,j)$，我们计算特征矩阵 $z$ 和 $f$ 中对应向量之间余弦相似度差的绝对值，从而促进它们在相对结构上的更紧密对齐。同样，我们减去一个边际值 $m_2$ 来放宽约束。ReLU 函数确保只有差异超过 $m_2$ 的特征对才会对损失产生贡献。

$$
\mathcal{L}_{mdms}

\frac{1}{N^2}
\sum_{i,j}
\mathrm{ReLU}
\left(
\left|
\frac{z_i \cdot z_j}{|z_i||z_j|}

\frac{f_i \cdot f_j}{|f_i||f_j|}
\right|

m_2
\right)
\quad (3)
$$

3.3. AdaptiveWeighting

---

在图 3 中，原始的重建损失和 KL 损失都属于求和型损失，这使得 VF 损失处于完全不同的尺度上，从而难以调整其权重以实现稳定训练。受到 GAN Loss [7] 的启发，我们采用了一种自适应加权机制。在反向传播之前，我们计算 $L_{vf}$ 和 $L_{rec}$ 在编码器最后一个卷积层上的梯度，如公式 4 所示。自适应权重被设置为这两个梯度的比值，确保 $L_{vf}$ 和 $L_{rec}$ 对模型优化具有相似影响。这种对齐显著减少了 VF 损失的调整范围。

$$
w_{adaptive}

\frac{|\nabla L_{rec}|}{|\nabla L_{vf}|}
\quad (4)
$$

随后我们得到使用自适应加权后的 VF 损失，如公式 5 所示。自适应加权的目的在于快速对齐不同 VAE 训练流程中的损失尺度。在此基础上，我们仍然可以使用人工调节的超参数进一步提升性能。

我们将在后续实验中评估 VF 损失在重建与生成任务上实现潜扩散帕累托前沿中的重要作用。

$$
L_{vf}

w_{hyper} * w_{adaptive}
\left(
L_{mcos} + L_{mdms}
\right)
\quad (5)
$$

4. Improved Diffusion Transformer

在本节中，我们介绍我们的 LightningDiT。Diffusion Transformers（DiT）[30] 已在文本生成图像 [3, 8] 和文本生成视频任务 [28, 32] 中取得显著成功。然而，其在 ImageNet 上的收敛速度明显慢于自回归（AR）[22] 模型，导致实验迭代成本较高。先前具有影响力的工作，如 DINOv2 [29]、ConvNeXt [25] 和 EVA [10] 展示了如何通过引入先进的设计策略来重振经典方法 [16, 48]。在我们的工作中，我们旨在拓展 DiT 架构的潜力，并探索 DiT 能够达到的极限。虽然我们并不声称任何单独的优化细节是我们的原创贡献，但我们相信，一个开源、快速收敛的 DiT 训练流程将极大支持社区在扩散 Transformer 方向上的持续研究。

我们使用 SD-VAE [34] 的 f8d4 规格作为视觉分词器，并采用 DiT-XL/2 作为实验模型。我们在表 1 中展示优化流程。

我们采用三类优化策略。在计算层面，我们使用 torch.compile [38] 和 bfloat16 训练来加速训练。此外，我们增大 batch size，并将 AdamW 的 beta2 降至 0.95，借鉴了此前成功的工作 AuraFlow [9]。在扩散优化方面，我们引入 Rectified Flow [24, 27]、logit normal distribution（lognorm）采样 [8] 和 velocity direction loss [43]。在模型架构层面，我们应用常见的 Transformer 优化，包括 RMSNorm [46]、SwiGLU [35] 和 ROPE [36]。在训练过程中，我们观察到某些加速策略并非正交。例如，梯度裁剪在单独使用时有效，但在 lognorm 采样和 velocity direction loss 之后组合使用时，反而会降低性能。

我们优化后的模型 LightningDiT 在 ImageNet 类别条件生成上约 80 个 epoch 即达到了 FID 7.13（cfg=1），仅为原始 DiT 和 SiT 模型在 1400 个 epoch 所需训练量的 6%。此前优秀的工作如 MDT [12] 或 REPA [45] 在 Mask Image Modeling（MIM）和表征对齐的帮助下取得了类似的收敛性能。我们的结果显示，即便没有复杂的训练流程，朴素的 DiT 也仍然能够实现非常有竞争力的性能。这一优化后的架构对我们后续的快速实验验证提供了极大帮助。

表 1. LightningDiT 的性能。LightningDiT 在 ImageNet 类别条件生成任务上获得了 FID-50k=7.13，相比原始 DiT [30] 使用了 94% 更少的训练样本。我们的方法超越了精心设计的 “MIM+Diffusion”† 方法（如 MDT [12]）和 “Representation+Diffusion”* 方法（如 REPA [45]），同时仍保持极高的简洁性。我们展示了原始 DiT 只需结合先进的设计技术，也能够达到卓越的性能。