扩散作为自蒸馏：单模型端到端潜变量扩散

01 动机

当前主流 Latent Diffusion Model（如 Stable Diffusion、DiT）均依赖"三件套"架构：独立的 VAE 编码器、VAE 解码器与扩散网络，三者依次独立训练。这一设计带来三大痛点：无法统一到现代视觉基础模型（Vision Foundation Model）、多阶段优化导致次优性能、VAE 组件占用约 20% 参数量并增加推理延迟。本文将此问题重新定义为一个无监督表征学习挑战，并借鉴 Self-Distillation 方法（如 DINO、SimSiam）避免表征坍塌的机制，提出端到端联合训练方案。

"We unify these three components into a single trainable network, enabling end-to-end optimization and integration with modern vision foundation models."

Self-distillation 与 Diffusion 的类比 — **图 1：Self-Distillation 与 Diffusion 框架的对应类比。** 左侧为经典 Self-Distillation 框架（如 DINO）：在线分支（online branch）通过 predictor 输出，与目标分支（target branch，使用 stop-gradient）对齐。右侧为 DSD：扩散过程的加噪操作对应 online 分支的数据增强，去噪网络（denoiser）对应 predictor，干净潜变量对应 target 表征。两者共享相同的防坍塌机制——Rank Differentiation。

4.25DSD-B FID（无 CFG）
ImageNet 256×256

3.35DSD-B FID（有 CFG）
50 epochs

205MDSD-B 全部参数
（含 encoder+decoder+diffusion）

50训练 epochs
vs. 基线 800–1400 epochs

为什么朴素联合训练会失败？

直接将 VAE 与扩散网络进行端到端联合训练会导致 Latent Collapse——潜变量的有效秩（effective rank）骤降至接近 1，表征退化，生成质量崩溃。论文识别出两个根本原因：

原因一：潜变量方差压制

L2 扩散损失隐式包含方差惩罚项，迫使编码器将所有潜变量压缩至均值附近，导致表征空间坍塌。

原因二：秩差异条件违反

Self-Distillation 理论要求 erank(z₂) > erank(P(z₁,t,ε))。标准速度预测（velocity prediction）输出全秩噪声，打破此条件，稳定机制失效。

02 方法

DSD 通过两项关键设计解决 Latent Collapse，并在单一 ViT 主干上构建统一架构，用三个轻量任务头分别处理编解码与扩散。

DSD 统一架构 — **图 4：DSD 统一模型架构。** 单一 Vision Transformer（ViT, patch size 16×16）主干同时承担三项任务： (1) *Image Head*（两层卷积）负责编码/解码； (2) *Diffusion Head*（单层 DiT block）负责速度预测与干净潜变量预测； (3) *Classification Head*（辅助损失）。 Online 分支接受带噪潜变量与强数据增强输入，Target 分支（EMA 更新，衰减率 0.99）接受原始潜变量，两分支共享 ViT 主干权重。潜变量空间：256 tokens × 16 维，4 个 register tokens 增强表征能力。

关键设计一：Stop-Gradient 解耦（Decoupling）

对目标潜变量 z₂ 施加 stop-gradient 算子，切断方差惩罚的梯度路径，消除 L2 损失对编码器方差的隐式约束：

ℒ_DSD = 𝔼_t ‖ṽ(z_t, t) − sg(z₂)‖²

同时引入 Detached Velocity Loss，为扩散头提供完整梯度信号： ℒ_velo = 𝔼_t ‖v(z_t, t) − sg(z − ε)‖²

关键设计二：损失变换（Loss Transformation）

论文通过数学推导证明，速度预测与干净潜变量预测在期望意义下等价：

ℒ_z = 𝔼_t w_t 𝔼 ‖ṽ(z_t, t) − z‖²

此等价性迫使 predictor 在学习过程中充当"低通滤波器"（low-pass filter），输出有效秩低于输入的表征，从而满足 Self-Distillation 的 Rank Differentiation 稳定条件。

**图 3：损失变换的几何解释。** 速度预测目标（v = z − ε）与干净潜变量目标（z）在几何上形成特定角度关系，论文证明优化速度预测损失等价于对干净潜变量的加权预测，确保 predictor 输出有效秩低于输入，激活 Rank Differentiation 稳定机制。

关键设计三：EMA 目标编码器 + 数据增强

借鉴 DINO/BYOL 中的 Momentum Encoder：Target 分支权重通过指数移动平均（EMA，衰减率 0.99）平滑更新，减少训练波动。 Online 分支施加强增强：75% 比例的随机 Masking、Gaussian Blur、Color Jittering、Solarization，进一步强化 Rank Differentiation。

各设计对有效秩和重建损失的影响 — **图 2：消融实验——各设计组件对有效秩（erank）与重建损失（rec. loss）的影响。** Case 1（朴素联合训练）：erank 骤降至 1，重建损失发散，彻底坍塌。 Case 2（仅 Decoupling）：坍塌变慢但仍发生（秩差异条件未满足）。 Case 3（+Loss Transformation）：erank 稳定，重建损失持续下降。 Case 4（+EMA 编码器）：收敛更平滑。 Case 5（+数据增强）：erank 稳定在高位，训练最稳定。

完整损失函数

DSD 的总损失由以下部分组成：

Self-Distillation Loss（ℒ_DSD）：主要防坍塌目标
Detached Velocity Loss（ℒ_velo）：为扩散头提供完整梯度
Reconstruction Loss：L2 + LPIPS + GAN 对抗损失，保证像素级重建质量
辅助损失：ViT 中间层对齐损失、中间表征的 Self-Distillation 损失、分类损失

03 实验

在 ImageNet 256×256 类别条件生成任务上评估，使用 gFID、sFID、Inception Score（IS）作为主要指标。采用 Euler 采样器，250 步扩散步数（使用 CFG 时额外指定引导强度 1.5）。训练使用 Muon 优化器，学习率 1e-4，梯度裁剪范数 3.0，在 8 张 NVIDIA A800 GPU 上训练。

与主流方法的全面对比

方法	VAE 参数	主模型参数	总参数	Epochs	gFID↓ (w/o CFG)	sFID↓ (w/o CFG)	gFID↓ (w/ CFG)
DiT-XL/2	84M	675M	759M	1400	9.62	6.85	2.27
SiT-XL/2	84M	675M	759M	1400	8.61	6.32	2.06
MAR-L	66M	945M	1011M	800	2.35	—	1.55
REPA-E (800ep)	70M	675M	745M	800	1.69	4.17	1.12
LightningDiT (800ep)	392M	675M	1067M	800	2.05	4.37	1.25
DSD-S（本文）	—		42M	50	13.44	11.74	7.89
DSD-M（本文）	—		118M	50	6.38	9.79	4.38
DSD-B（本文）	—		205M	50	4.25	8.96	3.35

注：所有数值均直接引用论文原文。绿色为本文方法结果，橙色为当前同类最优（SOTA）。 DSD-B（205M）在仅 50 epochs 下的表现优于 DiT-XL/2（759M，1400 epochs），参数量节省 73%，训练量节省 96%。 REPA-E 在更大参数与更多 epochs 下仍保持更优 FID，但其 VAE 与扩散网络独立训练。

定性生成结果

消融实验结论

论文通过逐步添加各设计组件的消融实验（如图 2 所示）验证了每个模块的必要性：

Decoupling（Stop-Gradient）：单独使用可延缓坍塌但无法完全阻止，因秩差异条件尚未满足。
Loss Transformation：与 Decoupling 结合后训练稳定，erank 维持高位，重建损失持续降低。
EMA 目标编码器：显著平滑训练曲线，减少振荡。
强数据增强（masking ratio 0.75）：进一步强化 Rank Differentiation，是最终稳定训练的关键。

训练效率

在 8 张 NVIDIA A800 GPU 上，每个 epoch 的墙钟时间：DSD-S 约 14.0 分钟，DSD-M 约 19.7 分钟，DSD-B 约 27.9 分钟。 50 epochs 总训练时间 DSD-B 约 23 小时，大幅低于基线方法所需资源。

04 局限性

说明：以下局限性 1–2 为作者在论文中明确陈述（stated）的内容；局限性 3–4 为基于方法设计的推断（inferred）。

模型规模受计算资源限制（stated）

论文原文："Due to computation resource constraint, we do not scale DSD to larger model sizes aligned with our baselines." 目前最大模型 DSD-B 仅有 205M 参数，远小于 DiT-XL/2 的 675M 主模型，更未达到 MAR 的 945M 规模。大规模模型下 DSD 的性能上界及扩展规律尚未验证。

未验证无监督学习能力（stated）

论文原文："We do not conduct experiments for verifying the effectiveness of our DSD as an unsupervised learning method." DSD 框架理论上兼容无监督表征学习，但论文中仅评估了类别条件图像生成任务，未在下游迁移学习或线性探针等任务上验证表征质量。

sFID 指标相对较弱（inferred）

DSD-B 的 sFID（空间 FID）为 8.96（无 CFG），而 REPA-E 仅为 4.17，LightningDiT 为 4.37。 sFID 对局部空间结构更敏感，较弱的 sFID 表明 DSD 在细粒度空间一致性上仍有提升空间，可能与单一 ViT 共享特征导致编解码与扩散之间存在任务冲突有关。

仅在 ImageNet 上验证，泛化性待确认（inferred）

所有实验均在 ImageNet 256×256 类别条件生成任务上进行，未验证文本条件生成、更高分辨率（如 512×512、1024×1024）或其他领域（医学图像、视频生成等）的有效性。端到端联合训练在数据分布差异较大时是否仍能防止 Latent Collapse 尚不明确。