← 返回 2026-05-11

扩散友好的潜在流形由什么决定?面向潜在扩散模型的先验对齐自编码器 What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion

Zhengrong Yue, Taihang Hu, Mengting Chen, Haiyu Zhang, Zihao Pan, Tao Liu, Zikang Wang, Jinsong Lan, Xiaoyong Zhu, Bo Zheng, Yali Wang 📅 2026-05-08 👍 9 2026-07-13 08:36
图像生成 流形学习 潜在扩散模型 自编码器 视觉基础模型

提出先验对齐自编码器PAE,显式塑造扩散友好流形,三项正则协同提升生成质量与收敛。

前置知识

潜在扩散模型(LDM)

在压缩后的潜在空间中进行扩散过程以降低计算开销的一类生成模型。自编码器把图像压缩到潜在码,扩散模型学习潜在码的分布,再由解码器还原图像。SD、DiT 都属于此类。

本文研究的就是 LDM 中第一阶段自编码器的潜在空间设计,直接关系到下游扩散模型的训练效率和生成质量。

视觉基础模型(VFM)

在大规模图像数据上预训练得到的强大特征提取器,例如 DINOv2、SigLIP、MAE 等。本文用冻结的 VFM 提供语义先验,避免重新从零学习语义结构。

PAE 的核心是用 VFM 特征作为对齐目标来塑造潜在流形,理解 VFM 才能看懂 SSR/MCR/SCR 三个正则化的设计动机。

流形假设(Manifold Assumption)

高维数据实际上分布在一个低维流形附近。本文把潜在空间也看成流形,研究其局部连续性、全局语义组织等几何性质如何影响扩散学习。

这是 PAE 区别于传统 VAE 的视角:不再只看重建误差,而是把潜在流形的几何结构作为显式优化目标。

rFID/gFID

rFID 衡量自编码器重建图像与原图在 Inception 特征空间的距离;gFID 衡量扩散生成图像与真实图像的分布距离,二者都是生成模型的常用评估指标。

本文反复强调 rFID 与 gFID 不强相关,PAE 的关键卖点就是在同等 rFID 下大幅降低 gFID,并加速收敛。

Gram 矩阵对齐

Gram 矩阵 $\mathbf{G} = \mathbf{Z}^\top\mathbf{Z}$ 编码特征 token 两两之间的内积关系,反映样本内部相对结构。在自监督表征学习中常用来保持结构信息。

PAE 的空间结构正则化(SSR)就是对齐潜在码与 VFM 特征的 Gram 矩阵,是其最核心的实例级结构保持手段。

研究动机

现有潜在扩散模型的第一阶段自编码器主要优化像素级重建损失(如 $\mathcal{L}_{\ell_1}$ + LPIPS + GAN),导致重建质量与下游扩散生成质量出现显著失配。论文的受控实验(图2a)显示,在固定潜在通道预算下,单纯提升 rFID 并不必然带来 gFID 的下降,甚至会出现 rFID 降低而 gFID 升高的情况。这种重建-生成鸿沟在多种主流 tokenizer(如 SD-VAE、RAE、SVG)上都被观察到。Vanilla VAE 仅靠 KL 正则虽然能稳定训练,但学到的潜在空间缺乏实例级空间结构、局部邻域不连续、全局语义相互纠缠,给下游扩散模型增加了额外的优化负担。

本文的目标是本文旨在回答"什么样的潜在空间才是扩散友好的"这一根本问题,并据此设计一个能够显式塑造扩散友好流形的 tokenizer 框架 PAE。具体目标包括:(1)通过受控实验识别与下游 gFID 强相关的潜在流形几何性质;(2)把这些性质转化为可微的训练正则项;(3)在 ImageNet 256×256 上同时提升最终生成质量与收敛速度,达到 rFID=0.26、800 epoch gFID=1.03 的新 SOTA。

与已有工作不同的是,已有工作要么沿用重建导向 VAE(如 SD-VAE),要么直接用冻结 VFM 特征作为潜在空间(如 RAE,前者难以生成高频细节,后者难以做精细编辑),还有的通过蒸馏把 VFM 特征塞进潜在空间(如 AlignTok、GAE、VTP)。这些方法都默认把潜在空间设计当成"重建"或"特征继承"的副产品,没有把潜在流形的几何组织本身作为主目标。PAE 的独特切入点是:放弃被动继承,明确把"实例级空间结构 + 局部流形连续性 + 全局语义组织"这三项流形属性作为显式的训练目标,并配套设计瓶颈匹配的 VFM 先验精炼策略。

核心方法

PAE 的整体思路可以一句话概括:用冻结 VFM 充当语义锚点,让自编码器在重建像素的同时,把潜在码的几何组织对齐到 VFM 流形上。架构上,输入图像先经冻结 VFM 编码得到高层语义特征 $H_{vfm}$,再经过一个轻量的 Detail-Aware Modulator (DAM) 注入像素级细节,得到调制后表征 $H_z$,由 Projector 投影到低维球形流形上的潜在码 $z \in \mathbb{R}^{B\times d\times H'\times W'}$($d=32$,相比 VFM 的 $d=1024/1152$ 大幅压缩)。解码路径由 Deprojector 与像素 Decoder 组成。整个训练在标准重建损失 $\mathcal{L}_{recon} = \mathcal{L}_{\ell_1} + \lambda_{lpips}\mathcal{L}_{LPIPS} + \lambda_{gan}\mathcal{L}_{GAN}$ 之外,额外加入三项先验对齐正则化:SSR、MCR、SCR,分别对应空间结构、局部连续、全局语义。

PAE 与已有 Representation-Native DiT(如 GAE、VTP、RAE)的本质区别在于:(1)显式流形塑造 vs 隐式特征继承:把三项扩散友好的几何性质变成训练目标,而不是让它们从重建或蒸馏中自发涌现;(2)级联式 MCR:用渐进扰动 $\alpha_l > \alpha_m > 0$ 让大扰动重建匹配中扰动、中扰动重建匹配原图,避免直接重建扰动目标带来的保真度损失;(3)VFM 先验精炼:先用一个轻量 prior projector 把原始高维 VFM 特征压到瓶颈维度并做空间低通精炼,让对齐目标在通道和空间上都与 tokenizer 解码器分辨率匹配;(4)零初始化 DAM:用零初始化的 scale-shift 融合 $\gamma_p, \beta_p = \text{split}(W\Delta H)$ 让训练从 $H_z = H_{vfm}$ 开始渐进注入细节,避免简单残差拼接造成的不可控混合。

方法步骤详情

PAE 训练分四步。冻结 VFM 提特征 $H_{vfm}=\mathcal{E}(x)$,$E$ 不训练;DAM 用 $K=6$ Transformer 让 patch token 与 $H_{vfm}$ 交互得 $\Delta H$,零初始化 scale-shift 注入 $H_z$,初期等价 $H_z=H_{vfm}$,避免残差拼接不可控;Projector 投到球形流形潜在码 $z$($d=32$);解码并算三项正则:SSR 用 Gram 矩阵对齐保实例结构,MCR 级联扰动让大扰动重建匹配中扰动、中扰动匹配原图避免保真度损失,SCR 用余弦对齐保 pooled 与 patch 语义;总损失 $\mathcal{L}=\mathcal{L}_{recon}+\lambda_{ssr}\mathcal{L}_{SSR}+\lambda_{mcr}\mathcal{L}_{MCR}+\lambda_{scr}\mathcal{L}_{SCR}$。VFM 精炼作预步骤:用 prior projector 压缩 VFM 特征并约束重建原表征,做空间上采样→低通→下采样,得瓶颈匹配目标。

技术新颖性

PAE 的技术新颖性主要体现在四个层面。第一,首次系统地把"扩散友好流形"拆解为可量化的三项几何指标 SSC/LPC/GSQ,并通过受控实验证明它们与 gFID 的相关性显著高于 rFID。第二,提出级联式 MCR 而非单一扰动重建,避免了大扰动损害重建保真度、小扰动正则太弱的权衡难题,把局部 Lipschitz 连续性以非破坏性方式注入训练目标。第三,瓶颈匹配的 VFM 精炼策略,区别于之前方法直接蒸馏高维 VFM 特征造成的语义-容量失配;其同时在通道(压缩到 tokenizer 同维度)和空间(低通精炼)两侧做匹配。第四,零初始化 DAM 在保留 VFM 主干结构的同时精准注入像素细节,相比直接 finetune VFM 或简单残差 concat 都更稳定。

Overview of the PAE framework
Figure 3: Overview of the PAE framework
Refined VFM priors provide better-matched alignment targets for PAE
Figure 4: Refined VFM priors provide better-matched alignment targets for PAE

实验结果

核心实验在 ImageNet 256×256 类条件生成,下游 LightningDiT-XL(675M)。表1显示 PAE(DINOv2) 80 epoch guided gFID=1.27,超过 VTP 1.44、GAE 1.48;800 epoch 把 gFID 推到 1.03(新 SOTA),unguided gFID=1.43、IS=244.8;收敛比 RAE 快约 13×,45 步采样 gFID 仍 1.05。图6 表明 PAE 在保真度-可学习性平面处 Pareto 前沿,SSC/LPC/GSQ/eRank 四指标最均衡。表2(a) 消融显示基线 7.18 gFID 单加 SSR/MCR/SCR 后分别降至 2.74/2.53/2.63,全加后 gFID=1.86、IS=210.8,每项正则最显著提升其对应几何维度,证实"显式优化→几何→质量"因果链。表2(b) 显示 VFM 精炼把 SSC 0.33→0.35、GSQ 0.48→0.50、gFID 1.95→1.86。图8 显示 PAE 对潜在维度 16-128、DAM 深度 2-8、四种 VFM 均稳定。

Generation performance on ImageNet 256×256
Table 1: Generation performance on ImageNet 256×256
Ablation study on prior alignment
Table 2: Ablation study on prior alignment
Ablation study on core design choices based on PAE (DINOv2)
Table 3: Ablation study on core design choices based on PAE (DINOv2)
Class-conditional samples by PAE with LightningDiT-XL/1
Figure 5: Class-conditional samples by PAE with LightningDiT-XL/1
Understanding PAE's fidelity-learnability advantage
Figure 6: Understanding PAE's fidelity-learnability advantage
Qualitative Comparison
Figure 7: Qualitative Comparison
Generalization and design analysis
Figure 8: Generalization and design analysis
查看结构化数据
任务指标本文基线提升
ImageNet 256×256 类条件生成(80 epoch 收敛) guided gFID ↓ 1.27 (PAE-DINOv2) VTP 1.44、GAE 1.48、RAE(DiTDH-XL) 报告 80 epoch 结果 在更少训练 epoch 下比同期 Representation-Native DiT 最佳结果再降 ~12%
ImageNet 256×256 类条件生成(800 epoch 长期训练) guided gFID ↓ 1.03 (PAE-DINOv2) RAE(DiTDH-XL) 1.13、FAE 1.29、AlignTok 1.37 新 SOTA,相对 RAE 降低 0.10,相对 AlignTok 降低 0.34
ImageNet 256×256 类条件生成(800 epoch) unguided gFID ↓ 1.43 (PAE-DINOv2) FAE 1.48、RAE(DiTDH-XL) 1.51 无 guidance 条件下仍优于 FAE/RAE
少步采样(45 步) gFID ↓ 1.05 全步采样 ~1.03 (800 epoch) 步数从数百降至 45 时 gFID 仅升 0.02
收敛效率(达到 RAE 同等质量所需 epoch) 训练 epoch 比 约 13× 加速 RAE 收敛所需 epoch 13× 训练加速
Tokenizer 重建保真度 rFID ↓ 0.26 (PAE-DINOv2) AlignTok 0.26、SFD 0.26、VA-VAE 0.28 与 SOTA 重建方法相当
几何指标综合表现 SSC↑/LPC↓/GSQ↑ 0.35 / 0.170 / 0.50 无 Lp 基线 0.18 / 0.320 / 0.19 三项几何指标全面大幅改善

局限与改进

作者在论文中相对克制地讨论局限性,主要局限包括:(1)依赖冻结 VFM 提供语义锚点,因此 VFM 的偏置(如 DINOv2 偏实例级、SigLIP 偏语义对齐)会直接影响 PAE 学到的流形,从图6(c) 看 SigLIP 和 MAE 版本的几何均衡性弱于 DINO 变体;(2)级联 MCR 中扰动尺度 $\alpha_m$、$\alpha_l$ 与方向 $\Delta$ 是额外超参,需在训练前选定,可能在不同数据集/分辨率下需要重新调优;(3)DAM 深度 K=6 带来可观的额外训练成本,论文未给出 decoder 单独推理延迟的数据;(4)实验主要在 ImageNet 256×256 类条件场景,缺乏在文本到图像、更高分辨率、非类别条件生成上的验证。本人额外观察到:PAE 表1中 SFD (1.30 gFID@80ep) 是带 AutoGuidance 的强基线,但 PAE 仅用 CFG 即超越,对照公平性需进一步细看;GRAM 矩阵对齐 SSR 与 MAE/对比学习中的常见做法类似,但 PAE 未讨论其与对比损失的可替代性。

独立分析的弱点

独立分析 PAE 的弱点与改进方向。第一,SSR 仅对齐 Gram 矩阵是"二阶统计量"对齐,对潜在码绝对位置和幅度没有显式约束,可能在不同 batch 间尺度漂移;可考虑结合 prototype anchoring 进一步收紧。第二,MCR 的扰动方向 $\Delta$ 是随机各向同性采样,未利用 VFM 语义方向;如果让扰动沿 VFM 主成分方向,可能让"局部连续"更贴近"语义连续",对条件生成更友好。第三,DAM 使用零初始化 scale-shift 虽然稳定但收敛偏慢,可考虑前期更激进 warm-up 加速训练。第四,所有评估集中在 ImageNet 256×256 类条件生成,缺乏文本条件/高分辨率/多分辨率扩展验证,对 SDXL 类系统迁移价值尚不明确。第五,论文未讨论 tokenizer 单独推理时的 FLOPS/延迟,DAM 引入的开销对工业部署的友好度未量化。第六,三项正则的权重 $\lambda_{ssr}$、$\lambda_{mcr}$、$\lambda_{scr}$ 缺乏 sensitivity 分析,读者难以判断超参鲁棒性。

未来方向

作者明确指出的方向与可合理延伸的方向包括:(1)把 PAE 框架扩展到更高分辨率(如 512×1024)以及文本到图像场景,验证其在 SDXL 类系统中的可移植性;(2)探索更便宜或更快的扰动机制,例如用 VFM 主成分方向替代随机 $\Delta$,把 MCR 与对比学习目标结合;(3)将三项正则化原则(空间结构 + 局部连续 + 全局语义)推广到其他表征驱动的生成任务,如视频潜在扩散、3D 生成;(4)作者暗示但未展开的方向是"流形组织原则"在自监督预训练中的应用,不局限于 LDM;(5)基于"瓶颈匹配"的 VFM 精炼思想,可延伸为通用的 teacher-student 蒸馏对齐方案,用于其他多模态任务;(6)把 SSC/LPC/GSQ 三个几何指标作为扩散友好性的统一评估基准,纳入未来 tokenizer 设计的标准报告项。

复现评估

复现评估整体良好。论文明确声明代码将在 GitHub、HuggingFace、ModelScope 三个平台发布,附有项目链接占位,但截至目前并未在公开数据中给出具体 URL。实验设置详细:使用 ImageNet 256×256,冻结 VFM 候选包括 DINOv2-L、SigLIP2-SO400M、DINOv3-L、MAE-L,latent 尺寸默认 $16\times16\times32$,DAM $K=6$,tokenizer 训练 50 epoch,下游 LightningDiT-XL 设置与 VA-VAE 一致。算力方面训练在 NVIDIA A100 GPU 上完成,但未提供具体卡数与训练时长。潜在的不复现风险包括:(1)VFM 精炼 prior projector 的实现细节在论文中相对模糊,需要参考附录 C.2;(2)MCR 扰动尺度与级联策略缺乏伪代码,依赖仔细复现;(3)下游 LightningDiT 的训练 schedule 与官方版本是否完全一致需要确认。综合来看,复现难度中等偏上,需要较强的工程能力对齐 VFM 特征与 tokenizer 训练管线。