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

📅 发布时间:2026/7/16 22:51:11 👁️ 浏览次数:
如何训练你的“潜变量“?Google DeepMind 提出 Unified Latents,用扩散模型同时编码、正则化和生成
如何训练你的潜变量Google DeepMind 提出 Unified Latents用扩散模型同时编码、正则化和生成一个框架统一编码器、扩散先验和扩散解码器ImageNet-512 FID 1.4Kinetics-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 ModelLDM不陌生——Stable Diffusion 就是基于它构建的。LDM 的核心思路很简单先用一个自编码器AutoEncoder把高分辨率图像压缩成低分辨率的潜变量latent再在这个压缩空间里训练扩散模型。这样做的好处是显而易见的计算量大幅降低生成效率大幅提升。但问题来了这个潜变量到底该怎么学这个问题看似简单实际上困扰了研究者很久。让我们看看现有方案的痛点方案一传统 VAE KL 正则化如 Stable Diffusion原始 LDM 使用 VAE 风格的 KL 惩罚让潜变量分布接近标准高斯分布。但由于解码器没有基于似然的损失KL 项的权重只能手动设置。这就像是在黑暗中调旋钮——你很难精确控制潜变量里到底编码了多少信息。方案二语义编码器如 DINO近年来一些工作使用预训练的语义网络如 DINO来获取潜变量。这些潜变量信息密度低、结构简单扩散模型学起来更容易FID 分数确实很漂亮。但代价是什么高频细节丢失严重重建图像看起来糊了——PSNR 通常低于 20细节与原图差异明显。 核心矛盾潜变量的信息量越少扩散模型越容易学习生成质量好但信息量太少重建质量就会下降。如何在这两者之间找到最优平衡点这正是Unified Latents要解决的问题。二、背景知识扩散模型与变分自编码器在深入 UL 的方法之前我们先快速回顾两个关键概念。2.1 变分自编码器VAEVAE 的核心是一个证据下界ELBO−log⁡p(x)≤−log⁡p(x∣z0)⏟解码器KL[q(z0∣x)∥p(z0)]⏟编码器 vs 先验-\log p(x) \leq \underbrace{-\log p(x|z_0)}_{\text{解码器}} \underbrace{KL[q(z_0|x) \| p(z_0)]}_{\text{编码器 vs 先验}}−logp(x)≤解码器−logp(x∣z0​)​​编码器vs先验KL[q(z0​∣x)∥p(z0​)]​​直觉上编码器把图像xxx压缩成潜变量z0z_0z0​解码器从z0z_0z0​重建图像KL 项确保潜变量分布不会偏离先验太远。2.2 扩散模型扩散模型通过逐步加噪→逐步去噪的过程来建模数据分布。关键概念是logSNR 调度λ(t)log⁡(αt2/σt2)\lambda(t) \log(\alpha_t^2 / \sigma_t^2)λ(t)log(αt2​/σt2​)它定义了每个时间步的破坏程度。模型学习预测干净数据x^(zt,t)\hat{x}(z_t, t)x^(zt​,t)整个过程可以分解为不同噪声水平上的损失贡献。一个重要的细节扩散模型的损失可以通过不同的加权函数w(λ)w(\lambda)w(λ)来调整。标准 ELBO 要求w(λ)1w(\lambda) 1w(λ)1但实际中常用 sigmoid 加权来提升图像质量。⚠️ 关键洞察如果用扩散模型作为先验但对损失做了非均匀加权编码器可能会钻空子——把信息编码到权重最低的噪声水平上逃避正则化。因此扩散先验必须使用无加权的 ELBO 损失。三、Unified Latents三个关键设计Unified Latents 的核心思想可以用一句话概括用扩散先验来正则化潜变量用扩散解码器来重建图像两者联合训练。听起来简单但魔鬼在细节里。让我们逐一拆解。UL概述显示了编码、扩散优先路径和扩散解码器展示图像编码为zcleanz_{\text{clean}}zclean​扩散先验从纯噪声到z0z_0z0​的路径以及扩散解码器重建图像的完整流程。3.1 编码与先验将编码噪声与扩散精度对齐这是 UL 最精妙的设计。传统 VAE 中编码器输出均值和方差方差控制潜变量的信息量。Stable Diffusion 让模型自己学习这个方差。但 UL 采取了完全不同的策略编码器只输出确定性潜变量zcleanE(x,θ)z_{\text{clean}} E(x, \theta)zclean​E(x,θ)不预测方差。对潜变量加固定量的高斯噪声得到z0z_0z0​。具体使用 logSNR 5 的噪声水平对应α0≈1.0\alpha_0 \approx 1.0α0​≈1.0σ0≈0.08\sigma_0 \approx 0.08σ0​≈0.08。潜变量z0z_0z0​带有约 8% 的噪声——不多不少刚好让扩散先验能有效建模。扩散先验从纯噪声z1z_1z1​出发逐步去噪到z0z_0z0​。由于噪声水平完全对齐KL 正则项简化为一个加权 MSE 损失——先验模型在各噪声水平上预测干净潜变量的误差之和。 三个好处训练稳定不需要学习编码器方差避免了 LSGM 中熵项估计不稳定的问题。可解释的比特率KL 项直接给出潜变量的信息上界bits per dimension。简洁的目标函数整个训练只有两项损失——先验损失 解码器损失。3.2 解码扩散解码器解码器同样是扩散模型但工作在图像空间。它接收带噪声的潜变量z0z_0z0​作为条件从纯噪声出发逐步生成图像。与先验不同解码器可以使用重加权的 ELBO如 sigmoid 加权。这意味着解码器可以打折处理高频细节——每比特信息的成本更低高频内容自然由解码器负责。解码器加权显示哪些噪声水平受到惩罚哪些是折扣展示解码器加权函数哪些噪声水平被惩罚潜变量优先哪些被折扣解码器优先。这里有一个关键超参数loss factor (lf)。它本质上是对解码器损失的放大等价于缩小 KL 项。为什么需要它因为当解码器很强大时VAE 容易出现后验坍塌posterior collapse——编码器干脆不编码任何信息全靠解码器硬生成。loss factor 通过适度放大解码器损失来缓解这个问题。实验中loss factor 通常设在1.3 到 1.7之间配合 sigmoid bias 参数就能精确控制潜变量的比特率。这两个超参数是 UL 框架中最重要的调节旋钮。3.3 第二阶段基础模型训练理论上第一阶段训练的先验模型就可以直接用来生成z0z_0z0​再通过解码器生成图像。但作者发现用 ELBO 损失训练的先验模型生成质量不佳。原因在于ELBO 损失对低频和高频内容一视同仁而生成高质量图像需要更关注感知上重要的中低频信息。解决方案冻结编码器和解码器重新训练一个基础模型base model使用 sigmoid 加权。由于只需要冻结的编码器基础模型的规模和 batch size 可以比第一阶段大得多。与标准 LDM 的唯一区别是UL 的基础模型有一个固定的最大 logSNR与先验相同采样时使用模型的最终预测z^\hat{z}z^而非z0z_0z0​。 打个比方第一阶段像是教潜变量该长什么样第二阶段像是教一个更大的模型如何画出这些潜变量。两阶段分工明确各司其职。四、实验结果全面碾压的训练效率作者在 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。潜变量类型gFID30KCLIPUL (LF1.5)4.127.1Pixel (无潜变量)5.027.0StableDiffusion6.827.0UL 在感知质量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)PSNRgFIDlf1.25, sb0~0.5较低较好lf1.5, sb0~1.0中等最优lf1.7, sb2~2.0较高略差 直觉理解比特率就像一个信息水龙头。开小了低比特率潜变量只保留粗略轮廓扩散模型容易学但细节靠解码器脑补开大了高比特率潜变量保留了太多细节扩散模型学习负担重。最优点在中间——大约 1 bpd 左右。一个有趣的发现最优比特率取决于基础模型的规模。小模型偏好低比特率信息少、容易学大模型可以消化更高比特率的潜变量。这为实际部署提供了重要指导。六、Mini Case Study从零搭建一个 UL 图像生成系统假设你是一位研究者想用 UL 框架在自己的数据集上训练一个图像生成模型。以下是一个简化的步骤指南第一步训练自编码器Stage 1准备编码器如 ViT-B和扩散解码器如 U-ViT设置 loss factor 1.5sigmoid bias 0中等比特率起步先验模型使用小型 U-NetELBO 损失无加权编码器输出 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训练至收敛第四步推理基础模型从纯噪声采样z^\hat{z}z^将z^\hat{z}z^加上 logSNR5 对应的噪声得到z0z_0z0​解码器以z0z_0z0​为条件从纯噪声生成最终图像⚠️ 注意事项先验模型的训练必须使用无加权 ELBO否则编码器会作弊编码器的 logSNR 不要设太低噪声太大否则信息损失严重解码器的采样步数通常需要 50-100 步才能获得高质量输出七、优势、局限与未来方向优势统一框架编码器、先验、解码器联合优化避免了分阶段训练的信息不匹配问题。训练效率高在 ImageNet-512 上UL 在相同 FLOPs 下始终优于 SD 潜变量和像素扩散。可解释的比特率控制两个超参数即可精确调节无需反复试错。通用性强同一框架适用于图像和视频理论上可扩展到其他模态。局限推理成本解码器是扩散模型需要多步采样比单步前馈解码器慢。论文中提到可以用蒸馏技术加速但未深入探讨。两阶段训练虽然第一阶段的先验模型可以丢弃但仍需两阶段训练流程。超参数敏感性最优比特率依赖于基础模型规模需要针对不同场景调优。未来方向论文提出了几个值得关注的方向离散数据的潜变量压缩将 UL 扩展到文本等离散数据用连续潜变量压缩 token 序列。多模态统一图像、视频、文本共享同一潜变量空间。更高效的解码器通过蒸馏或一致性模型减少解码步数。自适应比特率根据图像复杂度动态调整每张图的比特率。八、如何应用这些想法研究者UL 提供了一个干净的理论框架来理解潜变量学习。核心公式简洁适合在此基础上做改进。特别值得关注的是比特率与模型规模的关系——这可能是 scaling law 在潜变量空间的一个新维度。工程师/从业者如果你在训练自己的生成模型UL 的自编码器可以替代 SD 的 VAE获得更好的训练效率。关键配置建议loss factor 1.3-1.5sigmoid bias 0logSNR 516x 下采样。注意UL 目前没有开源代码论文来自 Google DeepMind但核心思想可以在现有框架上复现。产品经理/决策者UL 的核心价值是用更少的计算资源训练更好的生成模型。这对降低训练成本有直接意义。视频生成的 SOTA 结果FVD 1.3表明这个方向在视频领域尤其有潜力。九、常见误解与 FAQQ1UL 和 Stable Diffusion 的 VAE 有什么本质区别SD 的 VAE 使用前馈解码器一次性输出图像 手动设置的 KL 权重。UL 使用扩散解码器多步生成图像 由扩散先验自动确定的 KL 正则。最大的区别在于SD 的 KL 权重是拍脑袋定的而 UL 的正则化有严格的理论基础变分下界。此外UL 的解码器更强大可以在更低比特率下保持高重建质量。Q2为什么不直接在像素空间做扩散非要搞潜变量计算效率。一张 512×512 的图像有 786,432 个像素值而 16x 下采样后的潜变量只有 32×32×C数据量减少了几十倍。扩散模型的计算量与数据维度密切相关在潜变量空间训练可以节省大量算力。论文实验也证实在相同计算预算下UL 潜变量上的扩散模型显著优于像素空间扩散。Q3扩散解码器不会很慢吗推理时需要采样两次先验 解码器推理时不需要先验模型。流程是基础模型采样 → 得到z^\hat{z}z^→ 加噪得到z0z_0z0​→ 解码器采样生成图像。确实扩散解码器比前馈解码器慢但作者指出可以通过蒸馏技术如一致性蒸馏将解码器加速到少数几步。而且扩散解码器的优势在于它能在低比特率潜变量上脑补出高质量细节这是前馈解码器做不到的。Q4比特率越高越好吗不是。论文明确指出存在一个最优比特率取决于基础模型的容量。小模型偏好低比特率~0.5 bpd大模型可以利用更高比特率~1-2 bpd。超过最优点后FID 反而会变差因为扩散模型需要学习的信息量超出了其能力。Q5UL 能用于文本生成吗论文在讨论部分提到了这个可能性。理论上可以用 UL 框架将离散 token 序列压缩为连续潜变量然后用扩散模型建模。这类似于连续潜变量语言模型的思路但论文本身没有做这方面的实验。Q6UL 和 DALL-E 的 dVAE 有什么关系DALL-E 使用离散 VAEdVAE将图像压缩为离散 token然后用自回归模型生成。UL 使用连续潜变量 扩散模型两者的压缩方式和生成范式完全不同。UL 的优势在于连续潜变量可以更灵活地控制信息量且扩散模型在图像生成质量上已经超越了自回归方法。十、总结Unified Latents 是一篇把事情做对的论文。它没有提出花哨的新架构而是回归变分推断的基本原理用一个干净的数学框架统一了潜变量学习的三个核心组件。最让人印象深刻的不是某个单一数字而是整个框架的一致性理论推导清晰、实验设计严谨、消融实验充分、结果在图像和视频上都很强。这种从第一性原理出发的研究风格在当前追求 SOTA 的浮躁氛围中尤为难得。对于生成模型的从业者来说UL 传递的核心信息是潜变量的质量比模型的大小更重要。与其盲目堆参数不如先把潜变量学好。ReferencesHeek, J., Hoogeboom, E., Mensink, T., Salimans, T. (2026).Unified Latents (UL): How to train your latents. arXiv preprint arXiv:2602.17270. Google DeepMind.更多资源获取欢迎关注我的公众号「木子吉星」公众号二维码