统一潜变量（UL）：如何训练你的潜变量 Unified Latents (UL): How to train your latents
联合训练扩散先验和解码器学习潜变量表示，实现高效生成
前置知识
变分自编码器（VAE）
变分自编码器是一种生成模型，通过学习数据的潜变量表示来建模数据分布。它由编码器和解码器组成，编码器将输入数据映射到潜空间的分布参数（通常是均值和方差），解码器从潜变量重构数据。训练目标是最大化证据下界（ELBO），包含重构损失和KL散度正则项。在潜变量扩散模型（LDM）中，VAE用于学习紧凑的潜表示，使扩散模型能在低维空间高效运行。
本文的Unified Latents框架本质上是对VAE-扩散模型范式的革新，理解传统VAE的编码-解码流程和ELBO损失是理解本文方法改进的基础。
扩散模型（Diffusion Model）
扩散模型是一类通过逐步向数据添加高斯噪声（前向过程），然后学习逆转该过程（反向去噪）来生成样本的生成模型。给定数据分布 $q(x)$，前向过程定义为 $x_t = \alpha(t)x + \sigma(t)\epsilon$，其中 $\epsilon \sim \mathcal{N}(0, 1)$。通过学习预测干净数据 $\hat{x}(x_t, \theta)$ 来实现生成。常用的训练目标包括加权ELBO损失，如 $\text{sigmoid}(\lambda(t) - b)$ 加权。
本文的核心创新在于将扩散模型同时用作先验和解码器，需要深入理解扩散模型的前向/反向过程、噪声调度以及ELBO损失分解。
潜变量扩散模型（Latent Diffusion Model, LDM）
潜变量扩散模型（如Stable Diffusion）是将扩散过程应用于潜空间而非像素空间的生成模型。首先使用自编码器将高维图像压缩到低维潜空间，然后在潜空间中训练扩散模型。这种方法大幅降低了计算成本，使高分辨率图像生成成为可能。标准LDM使用VAE配合KL惩罚来正则化潜空间，但KL权重需要手动调节，难以控制潜变量的信息含量。
本文是对LDM范式的根本性改进，提出用扩散先验替代传统KL正则化，需要理解标准LDM的架构和局限性。
信噪比（Log-SNR）与噪声调度
在扩散模型中，信噪比 $\lambda(t) = \log(\alpha_t^2 / \sigma_t^2)$ 衡量时间步 $t$ 处信号相对于噪声的强度。噪声调度决定了从数据到纯噪声的退化过程。方差保持调度满足 $\alpha_t^2 + \sigma_t^2 = 1$。较高的 $\lambda$ 表示低噪声（更多信号），较低的 $\lambda$ 表示高噪声。信噪比是理解本文如何将编码器噪声与先验最小噪声级别对齐的关键概念。
本文通过固定编码器的最终信噪比 $\lambda(0) = 5$ 来控制潜变量精度，这是方法的核心设计决策，需要理解SNR的物理含义。
ELBO加权策略
扩散模型的训练可以看作是ELBO在不同噪声水平上的分解。标准ELBO要求 $w(\lambda_t) = 1$，但实践中常使用sigmoid加权 $w(\lambda(t)) = \text{sigmoid}(\lambda(t) - b)$ 来降低高频细节的损失权重，从而提升生成质量。这种加权对先验模型和解码器有不同影响：先验必须使用无加权ELBO以保证比特率界的有效性，而解码器可以使用加权ELBO来优化视觉质量。
理解不同组件（先验vs解码器）对加权策略的不同要求是掌握本文方法的关键，加权策略直接影响潜变量的信息分布和生成质量。
FID与PSNR评估指标
FID（Fréchet Inception Distance）通过比较生成样本和真实样本在Inception网络特征空间中的分布差异来衡量生成质量，越低越好。PSNR（峰值信噪比）衡量重构图像与原始图像的像素级相似度，越高表示重构越精确。在潜变量模型中，gFID衡量生成质量，rFID衡量重构质量。两者存在权衡关系：高信息潜变量导致低rFID但可能增加建模难度从而提高gFID。
本文大量使用这些指标来评估方法效果，理解gFID/rFID/PSNR的权衡关系是理解论文核心论点的前提。
研究动机
当前潜变量扩散模型面临一个根本性的设计困境：如何学习最优的潜变量表示？原始的Latent Diffusion Model（Rombach et al., 2022）使用VAE风格的KL惩罚将潜分布与标准高斯对齐，但由于解码器缺乏基于似然的损失，KL项的权重必须手动设置，使得潜变量的信息含量难以精确控制。近年来的研究转向使用预训练网络（如DINO）或强正则化的自编码器来获取语义表示，虽然获得了令人印象深刻的FID分数（如某些方法在ImageNet上达到较低FID），但高频信息严重丢失，PSNR分数通常低于20，导致重构结果在细节上与原始图像明显不同。具体来说，现有方法存在一个明确的权衡：潜变量的信息内容与输出重构质量之间的矛盾。如果潜结构更容易学习，通常会导致更好的生成性能，但这是以牺牲重构质量为代价的。理论上，即使是单个未正则化的潜变量通道也能编码无限量的信息，但实际上受机器精度和编码器平滑度的限制。潜通道的数量因此决定了信息容量：更少的通道产生更易建模但重构质量较低的潜变量，更多的通道则允许近乎完美的重构但需要更强的建模能力。目前没有一个原则性的方法来系统地导航这个权衡。
本文的目标是本文的核心目标是回答一个关键问题：当潜变量随后将由扩散模型建模时，应该如何对其进行正则化？具体来说，作者希望设计一个统一的框架，能够（1）提供对潜变量比特率的可解释上界，（2）通过简单的超参数控制重构-建模权衡，（3）在训练效率和生成质量上超越现有方法。在ImageNet-512上，目标是达到具有竞争力的FID分数同时保持高重构质量（高PSNR），并且在训练FLOPs上优于在Stable Diffusion潜变量上训练的模型。在视频生成领域（Kinetics-600），目标是设置新的FVD state-of-the-art。此外，作者还希望方法能够在大规模文生图和文生视频任务上进行扩展验证。
与已有工作不同的是，本文的独特切入角度在于将编码器、先验和解码器统一在扩散框架下进行联合训练。与现有方法的本质区别是：不再将编码器噪声视为需要学习的自由变量，也不使用简单的高斯先验进行正则化，而是将编码器的输出噪声显式地与扩散先验的最小噪声级别对齐。具体来说，编码器预测确定性潜变量 $z_{\text{clean}}$，然后通过固定的前向噪声将其转换为 $z_0$，使用最终对数信噪比 $\lambda(0) = 5$。这一设计选择带来三个关键优势：（1）KL项简化为噪声水平上的加权MSE，提供对比特率的紧上界；（2）避免了学习灵活编码器分布带来的训练不稳定（如LSGM中报告的问题）；（3）通过损失因子 $c_{\text{lf}}$ 和sigmoid偏置 $b$ 两个简单超参数即可控制潜变量的信息量。这种'联合训练扩散先验'的方法，作者称之为Unified Latents（UL），与传统的'先训练自编码器再训练扩散模型'的两阶段范式形成鲜明对比。
核心方法
Unified Latents（UL）框架的核心思想是将编码器、扩散先验和扩散解码器统一在一个联合训练框架中。直觉上，想象潜变量需要在两个竞争目标之间找到平衡：一方面要保留足够的信息以便解码器精确重构原始图像，另一方面要足够'简单'以便扩散先验能够有效建模。传统方法通过手动调节KL权重来实现这种平衡，而UL通过一个优雅的设计解决了这个问题——让编码器输出的噪声级别精确匹配扩散先验能够建模的最小噪声水平。技术路线上，整个系统包含三个组件：（1）编码器 $E_\theta$ 将图像 $x$ 映射到确定性潜变量 $z_{\text{clean}}$；（2）扩散先验 $P_\theta$ 建模从纯噪声 $z_1$ 到略微带噪的潜变量 $z_0$ 的路径；（3）扩散解码器 $D_\theta$ 以 $z_0$ 为条件从噪声重构图像。训练时，编码器和先验通过ELBO损失联合优化，解码器使用加权ELBO损失，两者通过共同的潜变量 $z_0$ 连接。这种设计使得潜变量的信息含量可以通过简单的超参数（损失因子和sigmoid偏置）进行精确控制。
本文的核心创新是将编码器噪声与扩散先验的最小噪声级别显式对齐。与已有方法的本质区别体现在三个方面：首先，与LSGM（Vahdat et al., 2021）不同，UL不需要计算编码器分布的熵项 $\mathbb{E}_{q(z_0|x)} \log q(z_0|x)$，这个项在实践中会导致训练不稳定。UL使用确定性编码器配合固定噪声来避免这个问题，将编码器分布吸收到扩散前向过程中。其次，与标准Stable Diffusion不同，UL不依赖手动设置的KL权重，而是通过 $\lambda(0) = 5$ 的固定最终信噪比来隐式控制潜变量精度，对应 $\alpha_0 = \sqrt{\text{sigmoid}(+5)} \approx 1.0$ 和 $\sigma_0 = \sqrt{\text{sigmoid}(-5)} \approx 0.08$。第三，与使用通道瓶颈进行正则化的方法（如 $\epsilon$-VAE、SWYCC）不同，UL通过学习的扩散先验来正则化，提供可解释的比特率界。这种设计使得KL项简化为 $\mathbb{E}_t[-\frac{d\lambda_z(t)}{dt} \frac{\exp(\lambda_z(t))}{2} w(\lambda_z(t)) ||z_{\text{clean}} - \hat{z}(z_t, \theta)||^2] + \text{KL}[p(z_1|x) || \mathcal{N}(0, I)]$，其中先验的ELBO加权必须是 $w(\lambda_z(t)) = 1$ 以保证界的有效性。
方法步骤详情
Unified Latents的训练分为两个阶段。第一阶段（联合训练）：对于每个训练样本 $x$，首先通过编码器获得确定性潜变量 $z_{\text{clean}} = E(x, \theta)$。然后采样时间步 $t \sim U(0,1)$ 和噪声 $\epsilon \sim \mathcal{N}(0, I)$，计算带噪潜变量 $z_t = \alpha_z(t)z_{\text{clean}} + \sigma_z(t)\epsilon$。先验损失为 $\mathcal{L}_z(\theta) = -\frac{d\lambda_z(t)}{dt} \frac{\exp(\lambda_z(t))}{2} ||z_{\text{clean}} - \hat{z}(z_t, \theta)||^2 + \text{KL}[p(z_1|x) || \mathcal{N}(0, I)]$。对于解码器，采样新的 $t$、$\epsilon$ 和 $\epsilon_z$，计算 $z_0 = \alpha_z(0)z_{\text{clean}} + \sigma_z(0)\epsilon_z$ 和 $x_t = \alpha_x(t)x + \sigma_x(t)\epsilon$。解码器损失为 $\mathcal{L}_x(\theta) = \frac{d\lambda_x(t)}{dt} \frac{\exp(\lambda_x(t))}{2} w(\lambda_x(t)) ||x - \hat{x}(x_t, z_0, \theta)||^2$，其中 $w(\lambda_x(t)) = c_{\text{lf}} \cdot \text{sigmoid}(b - \lambda_x(t))$。总损失 $\mathcal{L}(\theta) = \mathcal{L}_z(\theta) + \mathcal{L}_x(\theta)$。第二阶段（基础模型训练）：冻结编码器和解码器，使用sigmoid加权重新训练先验模型作为基础模型，此时可以使用更大的模型和批量大小。采样时，先从 $\mathcal{N}(0,I)$ 采样 $z_1$，通过基础模型生成 $z_0$，再从 $\mathcal{N}(0,I)$ 采样 $x_1$，通过解码器生成最终图像 $x$。
技术新颖性
Unified Latents的技术新颖性体现在多个层面。首先，框架层面的创新是将'编码-正则化-解码'三个步骤统一在扩散框架下，消除了传统VAE中编码器分布建模的复杂性。确定性编码器加固定噪声的设计（$\lambda(0) = 5$）相比学习编码器方差（如LSGM）既简化了训练又提升了稳定性——消融实验表明，学习方差的设置在 $\lambda_z(0) = 5$ 时会迅速退化为零并导致训练崩溃。其次，损失设计的创新在于对先验和解码器使用不同的加权策略：先验必须使用无加权ELBO以保证比特率界的有效性，而解码器使用sigmoid加权配合损失因子 $c_{\text{lf}}$（1.3-1.7）来优化视觉质量。这种'双重加权'策略是前所未有的。第三，两阶段训练策略的创新在于第一阶段使用ELBO加权训练先验（保证界的有效性），第二阶段使用sigmoid加权重新训练（提升采样质量），解决了ELBO加权先验采样质量差的问题。第四，方法提供了对潜变量比特率的可解释上界，通过简单的两个超参数（$c_{\text{lf}}$ 和 $b$）即可精确控制信息量，这是现有方法无法提供的。实验表明，UL在ImageNet-512上以0.059 bits/pixel的低比特率实现了1.54的gFID，而同等比特率下传统VAE方法的gFID通常远高于此。
实验结果
本文在图像和视频生成任务上进行了全面实验，核心发现如下。在ImageNet-512图像生成任务上，Unified Latents在训练计算量与生成质量的权衡上超越了所有现有方法（Figure 4）。具体来说，UL small模型在约1.2 zettaflops训练成本下达到1.42 gFID，UL medium模型在约2.5 zettaflops下达到1.31 gFID，而同等计算预算下的Stable Diffusion潜变量基线（small SD约2.2 gFID，medium SD约1.6 gFID）明显落后。重构质量方面，UL在损失因子1.3时达到25.7 PSNR和0.79 rFID@50k，损失因子1.7时提升至28.9 PSNR和0.36 rFID@50k。在文生图任务上（Table 1），UL（LF=1.5）在无引导条件下达到4.1 gFID@30K和0.271 CLIP分数，显著优于像素扩散（5.0 gFID）和Stable Diffusion潜变量（6.8 gFID）。在Kinetics-600视频生成任务上（Figure 9），UL small达到1.7 FVD，UL medium达到1.3 FVD，创下新的state-of-the-art，超越了Video Diffusion、MAGVIT-v2和W.A.L.T.等方法。潜变量通道数实验（Table 3）表明UL对通道数不敏感，16-64通道均能获得良好性能，仅4-8通道时重构质量下降。空间下采样实验（Table 4）显示16x和8x下采样在重构质量上相近，但16x更易建模，gFID更低（1.63 vs 2.12）。消融实验（Table 6）验证了各组件的重要性：去除先验梯度（A）导致gFID从1.54恶化至7.80，过高精度设置（B，$\lambda(0)=10$）使先验失效，学习方差（D）导致gFID从1.54退化至1.81。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet-512 图像生成 | gFID@50k | 1.31 (medium), 1.42 (small) | Stable Diffusion latents: medium SD ~1.6, small SD ~2.2; EDM2-XXL ~1.5 | 同等计算预算下UL显著优于SD潜变量基线，small模型提升约35%，medium模型提升约18% |
| ImageNet-512 重构质量 | PSNR / rFID@50k | 27.6 PSNR / 0.47 rFID (LF=1.5) | DINO-based方法通常PSNR ≤ 20 | PSNR提升约38%以上，保持高质量重构的同时实现强生成性能 |
| 文生图生成质量 | gFID@30K (无引导) | 4.1 (UL, LF=1.5) | Pixel diffusion: 5.0; Stable Diffusion latents: 6.8 | 相比像素扩散提升18%，相比SD潜变量提升40% |
| Kinetics-600 视频生成 | FVD | 1.3 (medium), 1.7 (small) | Video Diffusion ~5; MAGVIT-v2 ~3; W.A.L.T. ~2 | FVD降低约35%以上，设置新的state-of-the-art |
| 训练效率 | 训练FLOPs (zettaflops) | 达到1.4 gFID仅需~1.5 zettaflops | EDM2-XXL达到~1.5 gFID需~3 zettaflops | 达到同等生成质量所需训练计算量减少约50% |
| 潜变量比特率控制 | bits/pixel | 0.035-0.116 bits/pixel (通过调节LF 1.3-2.1) | 传统VAE需要更高比特率才能达到同等重构质量 | 在极低比特率下实现高质量生成和重构 |
局限与改进
本文作者坦诚地指出了若干局限性。首先，关于潜变量信息量的权衡问题：现有文献和本文都表明信息量较低的潜变量更容易建模，但一个根本性问题是：较弱的潜变量是否只是将部分建模问题转移给了解码器？本文使用U-Net扩散解码器，而大多数先前工作使用基于GAN的解码器配合判别器损失但没有噪声输入。扩散解码器严格来说比这类GAN更强大，因为它预测的是分布而非单一图像，但GAN训练的模式崩溃特性可能有助于生成更好看的图像和更低的rFID分数。其次，不同潜变量扩散模型之间的比较因自编码器训练数据的差异而变得复杂。原始Stable Diffusion自编码器在大规模网络数据集上训练，而本文大部分实验仅使用ImageNet。半监督方法（如DINO）引入了在大型外部数据集上训练的编码器，使得直接比较更加困难。第三，扩散解码器的采样成本比基于GAN的解码器高出一个数量级。如果没有额外的解码器蒸馏步骤，使用Unified Latents的计算成本显著高于标准LDM。此外，论文未充分讨论大规模实验的详细设置（如文生图实验使用内部数据集），这可能影响结果的可复现性。方法在低损失因子下的文本对齐（CLIP分数）略有下降（Figure 6），暗示解码器可能也需要文本条件化。
独立分析的弱点
经过独立分析，Unified Latents存在以下几个弱点。第一，扩散解码器的高采样成本是一个实际部署的瓶颈。文中承认扩散解码器比GAN解码器贵一个数量级，但未提供具体的推理时间对比数据。改进方向包括：对解码器进行蒸馏（如一致性模型或渐进式蒸馏），或者探索混合架构（扩散先验+轻量GAN解码器）。第二，损失因子 $c_{\text{lf}}$ 和sigmoid偏置 $b$ 的选择仍需要经验调优。虽然论文展示了Figure 8中的参数扫描，但最优设置因数据集、模型大小和任务而异，缺乏自动化的选择策略。改进方向是建立类似扩散模型scaling law的理论框架来预测给定计算预算下的最优比特率。第三，两阶段训练增加了总训练成本和复杂性。第一阶段训练的先验因使用ELBO加权而采样质量不佳，必须进行第二阶段重训练。改进方向探索单阶段训练方案（论文Appendix B有初步讨论但性能略差）。第四，文生图实验使用内部数据集且未公开详细信息，限制了结果的可复现性和公平比较。第五，方法对文本条件的整合有限——CLIP分数在低损失因子下略有下降，暗示解码器可能需要额外的文本条件化来改善文生图的文本对齐质量。
未来方向
基于本文成果，可以延伸出多个有前景的研究方向。作者在讨论部分提出了建立Unified Latents scaling laws的方向：预测给定训练预算下的最优比特率，这类scaling law将依赖于数据集、评估指标和模型架构等实现细节，最好在生产级基础模型的背景下研究。这一方向的重要性在于，当前的超参数选择仍依赖经验调优，而理论指导将大幅降低开发成本。另一个重要方向是将Unified Latents扩展到离散数据。作者指出，通过使用离散（扩散）解码器，理论上文本等离散数据也可以用潜变量压缩。这可能为统一的多模态生成模型开辟新路径。此外，探索单阶段训练方案（避免两阶段的额外成本）、开发解码器蒸馏方法以降低推理成本、以及研究Unified Latents在条件生成（如类条件、文本条件）中的最优策略都是值得探索的方向。从更宏观的角度看，UL框架提供的可解释比特率控制可能启发新的潜空间操作技术，如潜空间插值、编辑和风格迁移，因为信息量的精确控制使得潜空间的结构更加可预测。
复现评估
从复现角度来看，本文的复现难度中等偏高。有利因素包括：方法描述清晰，三个关键设计（固定噪声、先验对齐、sigmoid加权）均有明确的数学公式；消融实验（Table 6）验证了各组件的必要性，为复现提供了指导；超参数范围明确（损失因子1.3-1.7，$\lambda(0) = 5$）。不利因素包括：论文未提供代码或模型权重；架构细节虽然给出（编码器ResNet [128,256,512,512]，先验ViT 8块1024通道，解码器UVit [128,256,512]），但某些实现细节（如优化器、学习率调度、训练迭代数）需要从上下文推断；文生图和文生视频实验使用内部数据集，无法直接复现这些大规模结果。数据方面，ImageNet-512和Kinetics-600是公开数据集，图像实验可复现。算力需求方面，ImageNet实验的计算量在1-3 zettaflops范围，对于学术实验室来说需要相当的GPU资源，但对工业实验室来说是可行的。总体而言，核心方法的ImageNet实验应该是可复现的，但大规模文生图/视频实验的精确复现较为困难。
论文图表
