引言

在生成式 AI 飞速发展的今天，像 Stable Diffusion 和 Sora 这样的潜在扩散模型 (Latent Diffusion Models, LDMs) 已成为创建高保真图像和视频的黄金标准。这些模型的魔力在于它们不直接处理像素，而是在一个压缩的“潜在空间 (latent space) ”中运行。这种压缩由一个称为视觉分词器 (Visual Tokenizer) 的组件处理，通常是一个变分自编码器 (VAE) 。

长期以来，人们有一个简单的假设: 如果我们想要更好的图像，就需要更好的分词器。具体来说，我们假设增加分词器的容量 (维度) 将使其能够捕捉更多细节，从而让扩散模型生成更逼真的图像。

然而，最近的研究发现了一个令人沮丧的悖论。研究人员发现，虽然增加分词器的细节确实提高了其重建图像的能力，但它实际上损害了扩散模型生成新图像的能力。这就是优化困境 (Optimization Dilemma) 。 为了解决这个问题，传统做法是大规模增加扩散模型的参数量，但这会消耗大量的计算资源。

在这篇文章中，我们将深入探讨一篇新论文: “Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models” (重建与生成: 驯服潜在扩散模型中的优化困境) 。 我们将探索作者是如何发现这一瓶颈，并提出一种新颖的解决方案: VA-VAE (视觉基础模型对齐的 VAE) 。通过教导分词器像预训练的视觉模型 (如 DINOv2) 一样“思考”，他们仅用了一小部分的训练时间就实现了最先进 (SOTA) 的生成结果。

背景: 两阶段之舞

要理解这项创新，我们首先需要了解潜在扩散模型的架构。它分两个明显的阶段运行:

  1. 视觉分词器 (VAE) : 它充当翻译官的角色。它接收图像 (像素空间) 并将其压缩成更小、更密集的表示 (潜在空间) 。它还包含一个解码器，用于将该潜在表示转换回图像。
  2. 生成模型 (DiT) : 这代表了“大脑”。它通常是一个扩散 Transformer (DiT) 。它学习如何从噪声中创建新的潜在表示。

该系统的效率依赖于 VAE。如果 VAE 压缩图像过度，我们会丢失精细的细节 (如头发的纹理或文本) 。如果压缩不够，“大脑”需要处理的数据就会过多，导致训练缓慢且昂贵。

优化困境

研究人员首先测试了一个假设: 如果我们增加视觉 Token 的特征维度 (本质上是使压缩后的潜在数据的“文件大小”变大) ，我们应该能得到更好的图像。

他们测试了三种分词器规格，以其下采样因子 (\(f\)) 和维度 (\(d\)) 来标记:

  • f16d16: 标准压缩。
  • f16d32: 更高维度。
  • f16d64: 极高维度。

结果如下图所示，揭示了核心问题。

图 1. 潜在扩散模型中的优化困境。

请仔细观察图 1 。

  • 重建 (绿色箭头) : 当我们从 d16 移动到 d64 时，*重建图像 (Reconstructed Images) * 变得更清晰。rFID (重建 FID 分数，越低越好) 从 0.49 降至 0.18。模型几乎保留了完美的细节。
  • 生成 (红色箭头) : 然而，看看*生成图像 (Generated Images) *。它们崩溃了。gFID (生成 FID 分数) 从 20.3 飙升至 45.8。扩散模型未能有效地学习分布，导致生成了杂乱的噪声。

图 1底部的散点图暗示了原因。在高维空间 (f16d64) 中，潜在分布变得“结块”，高强度区域集中在很小的范围内。这使得空间变得不受约束，且对于扩散 Transformer 来说极其难以导航和学习。

解决方案: VA-VAE

作者认为，困境之所以存在，是因为我们要求 VAE 在没有足够指导的情况下从零开始学习高维潜在空间。这使得空间变得过于复杂和不规则。

他们的解决方案非常优雅: 不要从零开始学习。

他们提出了 VA-VAE (视觉基础模型对齐的 VAE) 。其核心思想是使用一位“老师”——一个预训练的视觉基础模型 (VFM) ，如 DINOv2 或 MAE——来指导 VAE 潜在空间的几何结构。这些基础模型已经观察了数百万张图像，并学习了极佳的、结构化的视觉数据表示。

图 3. 提出的视觉基础模型对齐 VAE (VA-VAE)。

如图 3 所示，架构仍然是标准的编码器-解码器设置。但是，这里增加了一个新的分支。在训练期间，输入图像通过视觉基础模型以提取特征。作者随后引入了一种新的损失函数，即 VF Loss (视觉基础损失) ，强制 VAE 的潜在 Token 与这些基础特征对齐。

这种对齐为 VAE 创建了一个“地图”，确保高维空间保持结构化，便于扩散模型后续的学习。

深入解析: VF Loss

本文的核心技术贡献是 VF Loss 的数学公式。仅仅说“让特征相似”是不够的。你需要约束空间，但又不能破坏 VAE 重建像素级细节的能力。

VF Loss 由两个特定部分组成。

1. 边际余弦相似度损失 (\(\mathcal{L}_{\mathrm{mcos}}\))

首先，来自 VAE 的潜在特征 \(Z\) 需要被投影到与基础模型特征 \(F\) 相同的维度。这是通过一个可学习的线性矩阵 \(W\) 完成的:

公式 1: 特征投影

现在，我们在每个空间位置 \((i, j)\) 比较投影后的 VAE 特征 (\(Z'\)) 与基础特征 (\(F\))。

公式 2: 边际余弦相似度损失

这里发生了什么?

  • 该公式计算 VAE Token 与基础 Token 之间的余弦相似度 。
  • 边际 (\(m_1\)): 这是天才之处。他们使用了 ReLU(1 - m - similarity) 结构。这实际上是在说: “我们要让特征相似，但不需要完全相同。”
  • 一旦相似度超过某个阈值 (\(1 - m_1\))，损失就变为零。这给了 VAE “回旋空间”来保留重建所需的局部细节，而不是强迫它完全复制基础模型 (基础模型可能会忽略像素级纹理) 。

2. 边际距离矩阵相似度损失 (\(\mathcal{L}_{\mathrm{mdms}}\))

虽然第一个损失对齐了单个点，但它没有考虑到点之间的关系。为了解决这个问题，作者引入了一种损失来保持特征空间的几何结构 。

公式 3: 边际距离矩阵相似度损失

它是如何工作的:

  • 这计算了 VAE 特征内 Token 对 (位置 \(i\) 与位置 \(j\)) 之间的相似度，并将其与基础特征中同一对的相似度进行比较。
  • 它确保如果图像的两个块在 DINOv2 眼中语义相关，那么它们在 VAE 的潜在空间中也应该是相关的。
  • 与前一个损失一样，它包含一个边际 (\(m_2\)) 以防止过度正则化。

3. 自适应加权

最后，出现了一个实际挑战: 这些新损失的规模可能与标准的像素重建损失完全不同。为了避免为每个实验手动调整权重，作者提出了自适应加权 (Adaptive Weighting) 。

公式 4: 自适应加权

公式 5: 总 VF Loss

该机制根据重建损失的梯度大小动态调整 VF Loss 的权重。它确保模型在整个训练过程中平衡“看起来像图像” (重建) 和“结构化空间” (VF Loss) 。

LightningDiT: 更快的基线

为了证明 VA-VAE 有效，研究人员需要训练扩散模型来使用这些新 Token 生成图像。然而，在 ImageNet 上训练扩散 Transformer (DiT) 出了名的慢。

为了加快反馈循环，他们构建了 LightningDiT , 这是标准 DiT 架构的优化版本。虽然这不是主要的理论贡献，但它是一项重大的工程成就。他们结合了几种现代训练“技巧”来加速收敛:

  1. Rectified Flow (矫正流) : 扩散过程的一种更高效的公式。
  2. 架构升级: SwiGLU 激活函数、RMSNorm 和旋转位置编码 (RoPE)。
  3. Lognorm 采样: 训练期间更好的时间步采样。
  4. 速度方向损失 (Velocity Direction Loss) : 一种专门的损失函数，用于拉直生成轨迹。

有了这些改进，LightningDiT 的收敛速度比原始 DiT 快得多，使团队能够对新分词器的有效性进行广泛的实验。

实验结果

对齐真的能解决优化困境吗？结果表明这无疑是肯定的。

拓展前沿

最重要的结果是重建与生成的“帕累托前沿 (Pareto Frontier) ”。理想情况下，我们要处于左下角 (低重建误差，低生成误差) 。

图 2. 潜在扩散模型的重建-生成前沿。

在图 2 中，注意点的聚类。

  • 左下角: 标准的 f16d16 模型表现尚可，但达到了极限。
  • 右下角: 未对齐的高维模型 (f16d32) 向右漂移——重建得到改善 (x 轴) ，但生成变差 (y 轴) 。
  • 右上角 (绿色箭头) : VA-VAE (对齐) 模型推动了边界。它们保持了高维模型的出色重建能力，同时将生成质量拉回到极佳的水平。

定量提升

我们可以在表 2 中看到原始数据。

表 2. VF loss 提升生成性能。

关注 f16d64 (最高维度) 的行。

  • 标准 LDM: rFID 为 0.17 (极佳的重建) ，但生成 FID 为 36.83 (糟糕的生成) 。
  • LDM + VF Loss (DINOv2): rFID 保持在 0.14 的低位，但生成 FID 降至 24.00 。
  • 这证明了 VF Loss 使得使用高维 Token 成为可能，而不会破坏生成能力。

收敛速度

对于研究人员来说，最实际的好处可能是训练速度。因为潜在空间更干净，扩散 Transformer 学习得更快。

图 4. VF Loss 加速收敛。

图 4 显示了训练曲线。橙色线 (带有 VF Loss) 立即下降，与蓝色线 (标准) 相比，仅用一小部分的步数就达到了更低的 FID 分数。作者报告称，对于高维分词器，收敛速度实现了 2.7 倍的加速 。

最先进的性能

当结合 VA-VAE 和优化的 LightningDiT 时，该系统在 ImageNet 256x256 上取得了显著的成果。

表 3. ImageNet 256x256 上的系统级性能。

如表 3 所示，该系统实现了 1.35 的 FID，击败了 SiT 和 MDT 等之前的 SOTA 方法。更令人印象深刻的是，它仅在 64 个 epoch 内就达到了 2.11 的竞争性 FID，相比原始 DiT 论文实现了 21 倍的加速 。

视觉质量

数字看起来不错，但图像如何呢？

图 5. 可视化结果。

图 5 展示了系统生成的样本。北极熊皮毛的细节、竹子的纹理以及汉堡上复杂的光照表明，该模型已成功利用高维潜在空间捕捉到了细粒度的细节。

为什么它有效？

最后，我们要问: 为什么 DINOv2 对齐会有这么大的不同？作者使用 t-SNE (一种在 2D 中可视化高维数据的方法) 提供了一个令人信服的可视化。

图 6. 使用 t-SNE 可视化潜在空间。

在图 6 中，看右侧的图( VF DINOv2 )。

  • 朴素 (Naive) 模型 (左) 显示的聚类有些混乱和重叠。
  • 对齐 (Aligned) 模型 (右) 显示了清晰、分离良好的聚类。

这证实了 VF Loss 有效地“组织”了潜在空间。它迫使 VAE 将语义相似的项目分组在一起 (例如，所有“狗”的 Token 在一个区域，“车”的 Token 在另一个区域) 。当扩散模型稍后尝试学习这个空间时，它的工作要轻松得多，因为底层的地图是合乎逻辑且结构化的。

这得到了 表 6 的进一步支持，该表测量了空间的“均匀性 (Uniformity) ”。

表 6. 均匀性与生成性能之间的关系。

数据显示出明显的相关性: 更低的密度变化 (更均匀的空间) 导致更好的生成 FID 分数。

哪个基础模型最好？

最后，你使用哪位老师重要吗？

表 4. 基础模型消融实验。

表 4 比较了使用 DINOv2、MAE、SAM 和 CLIP 的情况。虽然它们都有帮助，但 DINOv2 (一种自监督模型) 提供了最大的提升。这表明 DINOv2 学习到的鲁棒、语义特征是图像生成任务的最佳“向导”。

结论

“优化困境”长期以来一直是高分辨率图像生成的棘手问题。研究人员被迫在清晰的重建 (使用高维 Token) 或稳定的生成 (使用低维 Token) 之间做出选择，或者付出巨大的计算成本来强行解决。

本文提出的 VA-VAE 和 VF Loss 提供了一种更明智的前进道路。通过将 VAE 的潜在空间与预训练的视觉基础模型对齐，我们可以享受两全其美: 高维 Token 的高保真度和紧凑 Token 的训练效率。

结合 LightningDiT 中的工程改进，这项工作为更快、更便宜和更高质量的生成模型打开了大门。它提醒我们，有时，学习一项复杂任务的最佳方法是拥有一个好老师。