如何训练你的"潜变量"？Google DeepMind 提出 Unified Latents，用扩散模型同时编码、正则化和生成

一个框架统一编码器、扩散先验和扩散解码器，ImageNet-512 FID 1.4，Kinetics-600 FVD 1.3，刷新多项纪录。

从图像到潜在到生成的 UL 管道

Key Takeaways / 核心要点

Unified Latents (UL) 是 Google DeepMind 提出的潜变量学习框架，将编码器、扩散先验（diffusion prior）和扩散解码器（diffusion decoder）联合训练，形成一个统一的生成流水线。
核心创新：将编码器输出的噪声与先验模型的最小噪声水平对齐，使得 KL 正则项简化为加权 MSE，训练目标简洁且可解释。
在 ImageNet-512 上达到 FID 1.4，同时保持高重建质量（PSNR），训练 FLOPs 低于基于 Stable Diffusion 潜变量的模型。
在 Kinetics-600 视频生成上达到 FVD 1.3，刷新当前 SOTA。
提供了对潜变量比特率的可解释控制——通过两个简单超参数（loss factor 和 sigmoid bias）即可调节重建质量与生成质量之间的权衡。
框架具有通用性，理论上可扩展到文本等离散数据的潜变量压缩。

一、为什么我们需要更好的"潜变量"？

如果你关注 AI 图像生成领域，一定对 Latent Diffusion Model（LDM） 不陌生——Stable Diffusion 就是基于它构建的。LDM 的核心思路很简单：先用一个自编码器（AutoEncoder）把高分辨率图像压缩成低分辨率的"潜变量"（latent），再在这个压缩空间里训练扩散模型。这样做的好处是显而易见的：计算量大幅降低，生成效率大幅提升。

但问题来了：这个潜变量到底该怎么学？

这个问题看似简单，实际上困扰了研究者很久。让我们看看现有方案的痛点：

方案一：传统 VAE + KL 正则化（如 Stable Diffusion）

原始 LDM 使用 VAE 风格的 KL 惩罚，让潜变量分布接近标准高斯分布。但由于解码器没有基于似然的损失，KL 项的权重只能手动设置。这就像是在黑暗中调旋钮——你很难精确控制潜变量里到底编码了多少信息。

方案二：语义编码器（如 DINO）

近年来一些工作使用预训练的语义网络（如 DINO）来获取潜变量。这些潜变量信息密度低、结构简单，扩散模型学起来更容易，FID 分数确实很漂亮。但代价是什么？高频细节丢失严重，重建图像看起来"糊"了——PSNR 通常低于 20，细节与原图差异明显。

💡 核心矛盾：潜变量的信息量越少，扩散模型越容易学习（生成质量好）；但信息量太少，重建质量就会下降。如何在这两者之间找到最优平衡点？

这正是 Unified Latents 要解决的问题。

二、背景知识：扩散模型与变分自编码器

在深入 UL 的方法之前，我们先快速回顾两个关键概念。

2.1 变分自编码器（VAE）

VAE 的核心是一个证据下界（ELBO）：

$-\log p(x) \leq \underbrace{-\log p(x|z_0)}_{\text{解码器}} + \underbrace{KL[q(z_0|x) \| p(z_0)]}_{\text{编码器 vs 先验}}$

直觉上：编码器把图像 $x$ 压缩成潜变量 $z_0$ ，解码器从 $z_0$ 重建图像，KL 项确保潜变量分布不会偏离先验太远。

2.2 扩散模型

扩散模型通过逐步加噪→逐步去噪的过程来建模数据分布。关键概念是 logSNR 调度： $\lambda(t) = \log(\alpha_t^2 / \sigma_t^2)$ ，它定义了每个时间步的"破坏程度"。模型学习预测干净数据 $\hat{x}(z_t, t)$ ，整个过程可以分解为不同噪声水平上的损失贡献。

一个重要的细节：扩散模型的损失可以通过不同的加权函数 $w(\lambda)$ 来调整。标准 ELBO 要求 $w(\lambda) = 1$ ，但实际中常用 sigmoid 加权来提升图像质量。

⚠️ 关键洞察：如果用扩散模型作为先验，但对损失做了非均匀加权，编码器可能会"钻空子"——把信息编码到权重最低的噪声水平上，逃避正则化。因此，扩散先验必须使用无加权的 ELBO 损失。

三、Unified Latents：三个关键设计

Unified Latents 的核心思想可以用一句话概括：用扩散先验来正则化潜变量，用扩散解码器来重建图像，两者联合训练。

听起来简单，但魔鬼在细节里。让我们逐一拆解。

UL概述，显示了编码、扩散优先路径和扩散解码器

展示图像编码为 $z_{\text{clean}}$ ，扩散先验从纯噪声到 $z_0$ 的路径，以及扩散解码器重建图像的完整流程。

3.1 编码与先验：将编码噪声与扩散精度对齐

这是 UL 最精妙的设计。

传统 VAE 中，编码器输出均值和方差，方差控制潜变量的信息量。Stable Diffusion 让模型自己学习这个方差。但 UL 采取了完全不同的策略：

编码器只输出确定性潜变量 $z_{\text{clean}} = E(x, \theta)$ ，不预测方差。
对潜变量加固定量的高斯噪声，得到 $z_0$ 。具体使用 logSNR = 5 的噪声水平，对应 $\alpha_0 \approx 1.0$ ， $\sigma_0 \approx 0.08$ 。

潜变量 $z_0$ 带有约 8% 的噪声——不多不少，刚好让扩散先验能有效建模。

扩散先验从纯噪声 $z_1$ 出发，逐步去噪到 $z_0$ 。由于噪声水平完全对齐，KL 正则项简化为一个加权 MSE 损失——先验模型在各噪声水平上预测干净潜变量的误差之和。

💡 三个好处：

训练稳定：不需要学习编码器方差，避免了 LSGM 中熵项估计不稳定的问题。

可解释的比特率：KL 项直接给出潜变量的信息上界（bits per dimension）。

简洁的目标函数：整个训练只有两项损失——先验损失 + 解码器损失。

3.2 解码：扩散解码器

解码器同样是扩散模型，但工作在图像空间。它接收带噪声的潜变量 $z_0$ 作为条件，从纯噪声出发逐步生成图像。

与先验不同，解码器可以使用重加权的 ELBO（如 sigmoid 加权）。这意味着解码器可以"打折"处理高频细节——每比特信息的成本更低，高频内容自然由解码器负责。

解码器加权，显示哪些噪声水平受到惩罚，哪些是折扣

展示解码器加权函数，哪些噪声水平被惩罚（潜变量优先），哪些被折扣（解码器优先）。

这里有一个关键超参数：loss factor (lf)。它本质上是对解码器损失的放大（等价于缩小 KL 项）。为什么需要它？因为当解码器很强大时，VAE 容易出现后验坍塌（posterior collapse）——编码器干脆不编码任何信息，全靠解码器硬生成。loss factor 通过适度放大解码器损失来缓解这个问题。

实验中，loss factor 通常设在 1.3 到 1.7 之间，配合 sigmoid bias 参数，就能精确控制潜变量的比特率。这两个超参数是 UL 框架中最重要的"调节旋钮"。

3.3 第二阶段：基础模型训练

理论上，第一阶段训练的先验模型就可以直接用来生成 $z_0$ ，再通过解码器生成图像。但作者发现，用 ELBO 损失训练的先验模型生成质量不佳。

原因在于：ELBO 损失对低频和高频内容一视同仁，而生成高质量图像需要更关注感知上重要的中低频信息。

解决方案：冻结编码器和解码器，重新训练一个基础模型（base model），使用 sigmoid 加权。由于只需要冻结的编码器，基础模型的规模和 batch size 可以比第一阶段大得多。

与标准 LDM 的唯一区别是：UL 的基础模型有一个固定的最大 logSNR（与先验相同），采样时使用模型的最终预测 $\hat{z}$ 而非 $z_0$ 。

💡 打个比方：第一阶段像是"教潜变量该长什么样"，第二阶段像是"教一个更大的模型如何画出这些潜变量"。两阶段分工明确，各司其职。

四、实验结果：全面碾压的训练效率

作者在 ImageNet-512、Kinetics-600 以及大规模文生图/文生视频数据集上进行了实验。

4.1 ImageNet-512 图像生成

这是最核心的实验。自编码器在 512×512 分辨率上操作，下采样 16 倍得到 32×32 的潜变量。

在ImageNet-512上的FID与训练成本

训练成本 vs 生成 FID 的对比图，展示 UL 在各种训练预算下均优于其他方法。

关键发现：

UL 在训练成本 vs 生成质量的权衡上全面优于现有方法，是 ImageNet-512 上最高效的预训练方案。
为了公平比较，作者用完全相同的架构（2-level ViT）分别在 UL 潜变量和 Stable Diffusion 潜变量上训练。结果：UL 的优势更加明显。
最终达到 gFID 1.4，同时保持高 PSNR 重建质量。
有趣的是，patching 操作对 SD 潜变量上的基础模型性能有害，而 UL 不受此影响。

4.2 文生图（Text-to-Image）

作者在内部大规模文生图数据集上训练了多个自编码器（loss factor 从 1.25 到 1.7），并为每个自编码器训练了不同规模的基础模型（100、300、970 GFlops）。

潜变量类型	gFID@30K	CLIP
UL (LF=1.5)	4.1	27.1
Pixel (无潜变量)	5.0	27.0
StableDiffusion	6.8	27.0

UL 在感知质量（gFID）上显著优于像素扩散和 SD 潜变量，文本对齐（CLIP）也略有优势。

4.3 Kinetics-600 视频生成

在视频生成上，UL 使用 4×8×8 的下采样处理 16 帧 128×128 的视频。

Kinetics-600 上训练成本 vs FVD 的对比图

小模型就达到了 FVD 1.7
中等模型达到 FVD 1.3，刷新 SOTA
在相同训练预算下，UL 大幅领先 Video Diffusion、RIN、MAGVIT-v2 和 W.A.L.T.

五、比特率的艺术：如何精确调控潜变量信息量

UL 框架最优雅的特性之一，是对潜变量比特率的可解释控制。

作者发现，两个超参数几乎完全决定了比特率：

Loss factor (lf)：放大解码器损失的系数。lf 越大，编码器越倾向于编码更多信息（比特率越高）。
Sigmoid bias (sb)：控制解码器加权函数的偏移。sb 越大，解码器在低噪声水平上的权重越高，编码器需要编码更多细节。

不同 loss factor 和 sigmoid bias 下的比特率 vs 重建质量曲线

实验结果非常清晰：

配置	比特率 (bpd)	PSNR	gFID
lf=1.25, sb=0	~0.5	较低	较好
lf=1.5, sb=0	~1.0	中等	最优
lf=1.7, sb=2	~2.0	较高	略差

💡 直觉理解：比特率就像一个"信息水龙头"。开小了（低比特率），潜变量只保留粗略轮廓，扩散模型容易学但细节靠解码器脑补；开大了（高比特率），潜变量保留了太多细节，扩散模型学习负担重。最优点在中间——大约 1 bpd 左右。

一个有趣的发现：最优比特率取决于基础模型的规模。小模型偏好低比特率（信息少、容易学），大模型可以消化更高比特率的潜变量。这为实际部署提供了重要指导。

六、Mini Case Study：从零搭建一个 UL 图像生成系统

假设你是一位研究者，想用 UL 框架在自己的数据集上训练一个图像生成模型。以下是一个简化的步骤指南：

第一步：训练自编码器（Stage 1）

准备编码器（如 ViT-B）和扩散解码器（如 U-ViT）
设置 loss factor = 1.5，sigmoid bias = 0（中等比特率起步）
先验模型使用小型 U-Net，ELBO 损失（无加权）
编码器输出 logSNR = 5 的噪声潜变量
联合训练编码器 + 先验 + 解码器，约 400K 步

第二步：评估自编码器质量

检查重建 PSNR（目标 > 25 dB）
检查先验 KL（目标 ~1 bpd）
如果 PSNR 太低，增大 loss factor 或 sigmoid bias
如果 FID 太高，减小 loss factor

第三步：训练基础模型（Stage 2）

冻结编码器，丢弃先验和解码器
用更大的模型（如 2-level ViT-L）在编码后的潜变量上训练
使用 sigmoid 加权，可以用更大的 batch size
训练至收敛

第四步：推理

基础模型从纯噪声采样 $\hat{z}$
将 $\hat{z}$ 加上 logSNR=5 对应的噪声得到 $z_0$
解码器以 $z_0$ 为条件，从纯噪声生成最终图像

⚠️ 注意事项：

先验模型的训练必须使用无加权 ELBO，否则编码器会"作弊"

编码器的 logSNR 不要设太低（噪声太大），否则信息损失严重

解码器的采样步数通常需要 50-100 步才能获得高质量输出

七、优势、局限与未来方向

优势

统一框架：编码器、先验、解码器联合优化，避免了分阶段训练的信息不匹配问题。
训练效率高：在 ImageNet-512 上，UL 在相同 FLOPs 下始终优于 SD 潜变量和像素扩散。
可解释的比特率控制：两个超参数即可精确调节，无需反复试错。
通用性强：同一框架适用于图像和视频，理论上可扩展到其他模态。

局限

推理成本：解码器是扩散模型，需要多步采样，比单步前馈解码器慢。论文中提到可以用蒸馏技术加速，但未深入探讨。
两阶段训练：虽然第一阶段的先验模型可以丢弃，但仍需两阶段训练流程。
超参数敏感性：最优比特率依赖于基础模型规模，需要针对不同场景调优。

未来方向

论文提出了几个值得关注的方向：

离散数据的潜变量压缩：将 UL 扩展到文本等离散数据，用连续潜变量压缩 token 序列。
多模态统一：图像、视频、文本共享同一潜变量空间。
更高效的解码器：通过蒸馏或一致性模型减少解码步数。
自适应比特率：根据图像复杂度动态调整每张图的比特率。

八、如何应用这些想法？

研究者：

UL 提供了一个干净的理论框架来理解潜变量学习。核心公式简洁，适合在此基础上做改进。
特别值得关注的是比特率与模型规模的关系——这可能是 scaling law 在潜变量空间的一个新维度。

工程师/从业者：

如果你在训练自己的生成模型，UL 的自编码器可以替代 SD 的 VAE，获得更好的训练效率。
关键配置建议：loss factor 1.3-1.5，sigmoid bias 0，logSNR 5，16x 下采样。
注意：UL 目前没有开源代码（论文来自 Google DeepMind），但核心思想可以在现有框架上复现。

产品经理/决策者：

UL 的核心价值是"用更少的计算资源训练更好的生成模型"。这对降低训练成本有直接意义。
视频生成的 SOTA 结果（FVD 1.3）表明这个方向在视频领域尤其有潜力。

九、常见误解与 FAQ

Q1：UL 和 Stable Diffusion 的 VAE 有什么本质区别？

SD 的 VAE 使用前馈解码器（一次性输出图像）+ 手动设置的 KL 权重。UL 使用扩散解码器（多步生成图像）+ 由扩散先验自动确定的 KL 正则。最大的区别在于：SD 的 KL 权重是"拍脑袋"定的，而 UL 的正则化有严格的理论基础（变分下界）。此外，UL 的解码器更强大，可以在更低比特率下保持高重建质量。

Q2：为什么不直接在像素空间做扩散，非要搞潜变量？

计算效率。一张 512×512 的图像有 786,432 个像素值，而 16x 下采样后的潜变量只有 32×32×C，数据量减少了几十倍。扩散模型的计算量与数据维度密切相关，在潜变量空间训练可以节省大量算力。论文实验也证实：在相同计算预算下，UL 潜变量上的扩散模型显著优于像素空间扩散。

Q3：扩散解码器不会很慢吗？推理时需要采样两次（先验 + 解码器）？

推理时不需要先验模型。流程是：基础模型采样 → 得到 $\hat{z}$ → 加噪得到 $z_0$ → 解码器采样生成图像。确实，扩散解码器比前馈解码器慢，但作者指出可以通过蒸馏技术（如一致性蒸馏）将解码器加速到少数几步。而且，扩散解码器的优势在于它能在低比特率潜变量上"脑补"出高质量细节，这是前馈解码器做不到的。

Q4：比特率越高越好吗？

不是。论文明确指出存在一个最优比特率，取决于基础模型的容量。小模型偏好低比特率（~0.5 bpd），大模型可以利用更高比特率（~1-2 bpd）。超过最优点后，FID 反而会变差，因为扩散模型需要学习的信息量超出了其能力。

Q5：UL 能用于文本生成吗？

论文在讨论部分提到了这个可能性。理论上，可以用 UL 框架将离散 token 序列压缩为连续潜变量，然后用扩散模型建模。这类似于"连续潜变量语言模型"的思路，但论文本身没有做这方面的实验。

Q6：UL 和 DALL-E 的 dVAE 有什么关系？

DALL-E 使用离散 VAE（dVAE）将图像压缩为离散 token，然后用自回归模型生成。UL 使用连续潜变量 + 扩散模型，两者的压缩方式和生成范式完全不同。UL 的优势在于连续潜变量可以更灵活地控制信息量，且扩散模型在图像生成质量上已经超越了自回归方法。

十、总结

Unified Latents 是一篇"把事情做对"的论文。它没有提出花哨的新架构，而是回归变分推断的基本原理，用一个干净的数学框架统一了潜变量学习的三个核心组件。

最让人印象深刻的不是某个单一数字，而是整个框架的一致性：理论推导清晰、实验设计严谨、消融实验充分、结果在图像和视频上都很强。这种"从第一性原理出发"的研究风格，在当前追求 SOTA 的浮躁氛围中尤为难得。

对于生成模型的从业者来说，UL 传递的核心信息是：潜变量的质量比模型的大小更重要。与其盲目堆参数，不如先把潜变量学好。

References

Heek, J., Hoogeboom, E., Mensink, T., & Salimans, T. (2026). Unified Latents (UL): How to train your latents. arXiv preprint arXiv:2602.17270. Google DeepMind.

更多资源获取欢迎关注我的公众号：「木子吉星」

公众号二维码