← 论文海报合集← Paper Notes|
cs.CV · arXiv 2511.14716

扩散作为自蒸馏:单模型端到端潜变量扩散Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model

Diffusion As Self-Distillation: End-to-End Latent Diffusion In One Model
Xiyuan Wang, Muhan Zhang  ·  2025年11月Xiyuan Wang, Muhan Zhang  ·  November 2025

本文提出将编码器(encoder)、解码器(decoder)与扩散网络(diffusion network)整合为一个单一可训练模型,彻底打破传统 Latent Diffusion Model(LDM)三阶段分步训练的桎梏。 核心贡献在于识别了联合训练失败的根本原因——"Latent Collapse"(潜空间坍塌),并提出 Diffusion as Self-Distillation(DSD) 框架加以解决。 DSD-B(205M 参数)仅需 50 个训练 epoch,在 ImageNet 256×256 上取得 FID=4.25,远优于参数量大 3–5 倍的基线方法。This paper merges the encoder, the decoder and the diffusion network into a single trainable model, breaking the three-stage sequential training pipeline that conventional Latent Diffusion Models (LDM) are locked into. Its core contribution is identifying the root cause of joint-training failure — "Latent Collapse" — and proposing the Diffusion as Self-Distillation (DSD) framework to resolve it. DSD-B (205M parameters) needs only 50 training epochs and reaches FID=4.25 on ImageNet 256×256, far ahead of baselines with 3–5 times more parameters.

ImageNet 256×256 生成模型 · 扩散模型Generative Models · Diffusion Models 8× NVIDIA A800 GPU 📄 arXiv:2511.14716 PDF 原文PDF
关键词Keywordslatent diffusion modelend-to-end trainingself-distillation潜空间坍塌latent collapseVAE联合训练joint VAE trainingrank differentiationVision Transformer图像生成image generationEMA momentum encoderFID ImageNet

01 动机Motivation

当前主流 Latent Diffusion Model(如 Stable Diffusion、DiT)均依赖"三件套"架构:独立的 VAE 编码器、VAE 解码器与扩散网络,三者依次独立训练。 这一设计带来三大痛点:无法统一到现代视觉基础模型(Vision Foundation Model)多阶段优化导致次优性能VAE 组件占用约 20% 参数量并增加推理延迟。 本文将此问题重新定义为一个无监督表征学习挑战,并借鉴 Self-Distillation 方法(如 DINO、SimSiam)避免表征坍塌的机制,提出端到端联合训练方案。Mainstream Latent Diffusion Models (Stable Diffusion, DiT and the like) all rely on a "three-piece" architecture: a standalone VAE encoder, a VAE decoder and a diffusion network, each trained one after another. This design brings three pain points: it cannot be unified with modern vision foundation models, multi-stage optimization yields sub-optimal performance, the VAE components take about 20% of the parameters and add inference latency. This paper recasts the problem as an unsupervised representation learning challenge and, borrowing the mechanism by which Self-Distillation methods (DINO, SimSiam) avoid representation collapse, proposes an end-to-end joint training scheme.

"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。Figure 1: The correspondence between the Self-Distillation and the Diffusion framework. Left is the classic Self-Distillation framework (DINO and the like): the online branch emits its output through a predictor and is aligned with the target branch (which uses stop-gradient). Right is DSD: the noising operation of the diffusion process corresponds to the data augmentation of the online branch, the denoiser corresponds to the predictor, and the clean latent corresponds to the target representation. Both share the same anti-collapse mechanism — Rank Differentiation.
4.25DSD-B FID(无 CFG)
ImageNet 256×256
DSD-B FID (no CFG)
ImageNet 256×256
3.35DSD-B FID(有 CFG)
50 epochs
DSD-B FID (with CFG)
50 epochs
205MDSD-B 全部参数
(含 encoder+decoder+diffusion)
All DSD-B parameters
(encoder + decoder + diffusion)
50训练 epochs
vs. 基线 800–1400 epochs
training epochs
vs. 800–1400 for baselines

为什么朴素联合训练会失败?Why does naive joint training fail?

直接将 VAE 与扩散网络进行端到端联合训练会导致 Latent Collapse——潜变量的有效秩(effective rank)骤降至接近 1,表征退化,生成质量崩溃。 论文识别出两个根本原因:Training the VAE and the diffusion network jointly end-to-end leads straight to Latent Collapse — the effective rank of the latents plunges to nearly 1, the representation degenerates and generation quality falls apart. The paper identifies two root causes:

原因一:潜变量方差压制Cause 1: suppression of the latent variance

L2 扩散损失隐式包含方差惩罚项,迫使编码器将所有潜变量压缩至均值附近,导致表征空间坍塌。The L2 diffusion loss implicitly carries a variance penalty term that forces the encoder to squeeze every latent towards the mean, collapsing the representation space.

原因二:秩差异条件违反Cause 2: violation of the rank differentiation condition

Self-Distillation 理论要求 erank(z₂) > erank(P(z₁,t,ε))。标准速度预测(velocity prediction)输出全秩噪声,打破此条件,稳定机制失效。Self-Distillation theory requires erank(z₂) > erank(P(z₁,t,ε)). Standard velocity prediction outputs full-rank noise, which breaks this condition and disables the stabilizing mechanism.

02 方法Method

DSD 通过两项关键设计解决 Latent Collapse,并在单一 ViT 主干上构建统一架构,用三个轻量任务头分别处理编解码与扩散。DSD resolves Latent Collapse with two key designs and builds a unified architecture on a single ViT backbone, where three lightweight task heads take care of coding/decoding and diffusion respectively.

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 增强表征能力。Figure 4: The unified DSD model architecture. A single Vision Transformer (ViT, patch size 16×16) backbone carries three tasks at once: (1) the Image Head (two convolution layers) handles encoding/decoding; (2) the Diffusion Head (one DiT block) handles velocity prediction and clean-latent prediction; (3) the Classification Head (auxiliary loss). The online branch takes noised latents with strong data augmentation, the target branch (EMA update, decay 0.99) takes the original latents, and both branches share the ViT backbone weights. Latent space: 256 tokens × 16 dimensions, with 4 register tokens to strengthen the representation.

关键设计一:Stop-Gradient 解耦(Decoupling)Key design 1: Stop-Gradient decoupling

对目标潜变量 z₂ 施加 stop-gradient 算子,切断方差惩罚的梯度路径,消除 L2 损失对编码器方差的隐式约束:Applying a stop-gradient operator to the target latent z₂ cuts the gradient path of the variance penalty and removes the implicit constraint the L2 loss puts on the encoder variance:

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

同时引入 Detached Velocity Loss,为扩散头提供完整梯度信号:A Detached Velocity Loss is introduced at the same time, giving the diffusion head a complete gradient signal: ℒ_velo = 𝔼t ‖v(z_t, t) − sg(z − ε)‖²

关键设计二:损失变换(Loss Transformation)Key design 2: Loss Transformation

论文通过数学推导证明,速度预测与干净潜变量预测在期望意义下等价:The paper proves by derivation that velocity prediction and clean-latent prediction are equivalent in expectation:

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

此等价性迫使 predictor 在学习过程中充当"低通滤波器"(low-pass filter),输出有效秩低于输入的表征,从而满足 Self-Distillation 的 Rank Differentiation 稳定条件。This equivalence forces the predictor to act as a "low-pass filter" during learning, emitting representations whose effective rank is lower than that of the input, which satisfies the Rank Differentiation stability condition of Self-Distillation.

损失变换的几何解释
图 3:损失变换的几何解释。 速度预测目标(v = z − ε)与干净潜变量目标(z)在几何上形成特定角度关系,论文证明优化速度预测损失等价于对干净潜变量的加权预测,确保 predictor 输出有效秩低于输入,激活 Rank Differentiation 稳定机制。Figure 3: A geometric reading of the loss transformation. The velocity-prediction target (v = z − ε) and the clean-latent target (z) form a specific angular relation; the paper proves that optimizing the velocity-prediction loss is equivalent to a weighted prediction of the clean latent, which keeps the effective rank of the predictor output below that of its input and activates the Rank Differentiation stabilizing mechanism.

关键设计三:EMA 目标编码器 + 数据增强Key design 3: EMA target encoder + data augmentation

借鉴 DINO/BYOL 中的 Momentum Encoder:Target 分支权重通过指数移动平均(EMA,衰减率 0.99)平滑更新,减少训练波动。 Online 分支施加强增强:75% 比例的随机 Masking、Gaussian Blur、Color Jittering、Solarization,进一步强化 Rank Differentiation。Following the Momentum Encoder of DINO/BYOL, the target branch weights are updated smoothly by an exponential moving average (EMA, decay 0.99), which damps training fluctuation. The online branch receives strong augmentation: random masking at a 75% ratio, Gaussian Blur, Color Jittering and Solarization, further strengthening Rank Differentiation.

各设计对有效秩和重建损失的影响
图 2:消融实验——各设计组件对有效秩(erank)与重建损失(rec. loss)的影响。 Case 1(朴素联合训练):erank 骤降至 1,重建损失发散,彻底坍塌。 Case 2(仅 Decoupling):坍塌变慢但仍发生(秩差异条件未满足)。 Case 3(+Loss Transformation):erank 稳定,重建损失持续下降。 Case 4(+EMA 编码器):收敛更平滑。 Case 5(+数据增强):erank 稳定在高位,训练最稳定。Figure 2: Ablation — how each design component affects the effective rank (erank) and the reconstruction loss (rec. loss). Case 1 (naive joint training): erank plunges to 1, the reconstruction loss diverges, collapse is total. Case 2 (Decoupling only): collapse is slower but still happens (the rank differentiation condition is unmet). Case 3 (+Loss Transformation): erank is stable and the reconstruction loss keeps falling. Case 4 (+EMA encoder): convergence is smoother. Case 5 (+data augmentation): erank settles high and training is at its most stable.

完整损失函数The full loss function

DSD 的总损失由以下部分组成:The total DSD loss is made of the following parts:

03 实验Experiments

在 ImageNet 256×256 类别条件生成任务上评估,使用 gFID、sFID、Inception Score(IS)作为主要指标。 采用 Euler 采样器,250 步扩散步数(使用 CFG 时额外指定引导强度 1.5)。 训练使用 Muon 优化器,学习率 1e-4,梯度裁剪范数 3.0,在 8 张 NVIDIA A800 GPU 上训练。Evaluated on class-conditional generation on ImageNet 256×256, with gFID, sFID and Inception Score (IS) as the main metrics. An Euler sampler is used with 250 diffusion steps (when CFG is on, the guidance scale is additionally set to 1.5). Training uses the Muon optimizer, learning rate 1e-4, gradient-clipping norm 3.0, on 8 NVIDIA A800 GPUs.

与主流方法的全面对比A full comparison against mainstream methods

方法Method VAE 参数VAE params 主模型参数Main model params 总参数Total params Epochs gFID↓ (w/o CFG) sFID↓ (w/o CFG) gFID↓ (w/ CFG)
DiT-XL/284M675M759M1400 9.626.852.27
SiT-XL/284M675M759M1400 8.616.322.06
MAR-L66M945M1011M800 2.351.55
REPA-E (800ep)70M675M745M800 1.694.171.12
LightningDiT (800ep)392M675M1067M800 2.054.371.25
DSD-S(本文)DSD-S (ours)42M50 13.4411.747.89
DSD-M(本文)DSD-M (ours)118M50 6.389.794.38
DSD-B(本文)DSD-B (ours)205M50 4.258.963.35

注:所有数值均直接引用论文原文。绿色为本文方法结果,橙色为当前同类最优(SOTA)。 DSD-B(205M)在仅 50 epochs 下的表现优于 DiT-XL/2(759M,1400 epochs),参数量节省 73%,训练量节省 96%。 REPA-E 在更大参数与更多 epochs 下仍保持更优 FID,但其 VAE 与扩散网络独立训练。Note: every number is quoted directly from the paper. Green marks the results of this paper, amber marks the current best in class (SOTA). DSD-B (205M) at only 50 epochs beats DiT-XL/2 (759M, 1400 epochs), saving 73% of the parameters and 96% of the training budget. REPA-E still holds a better FID with more parameters and more epochs, but its VAE and diffusion network are trained separately.

定性生成结果Qualitative generation results

DSD 在 ImageNet 256×256 上的定性生成结果
图 5:DSD-B 在 ImageNet 256×256 上的条件生成样本。 图中展示多个 ImageNet 类别的生成结果,DSD 能够生成高质量、多样化的图像,且在类别条件引导下保持良好的语义一致性,证明单模型架构可以有效学习高质量的图像生成能力。Figure 5: Conditional samples from DSD-B on ImageNet 256×256. The figure shows generations for a number of ImageNet classes: DSD produces high-quality, diverse images and keeps good semantic consistency under class-conditional guidance, showing that a single-model architecture can effectively learn high-quality image generation.

消融实验结论Ablation findings

论文通过逐步添加各设计组件的消融实验(如图 2 所示)验证了每个模块的必要性:Ablations that add the design components one by one (as in Figure 2) confirm that every module is necessary:

训练效率Training efficiency

在 8 张 NVIDIA A800 GPU 上,每个 epoch 的墙钟时间:DSD-S 约 14.0 分钟,DSD-M 约 19.7 分钟,DSD-B 约 27.9 分钟。 50 epochs 总训练时间 DSD-B 约 23 小时,大幅低于基线方法所需资源。On 8 NVIDIA A800 GPUs the wall-clock time per epoch is about 14.0 minutes for DSD-S, 19.7 minutes for DSD-M and 27.9 minutes for DSD-B. The full 50 epochs take DSD-B about 23 hours, far below the resources the baseline methods demand.

04 局限性Limitations

说明:以下局限性 1–2 为作者在论文中明确陈述(stated)的内容;局限性 3–4 为基于方法设计的推断(inferred)Note: limitations 1–2 are stated explicitly by the authors in the paper; limitations 3–4 are inferred from the design of the method.
模型规模受计算资源限制(stated)Model scale limited by compute resources (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 的性能上界及扩展规律尚未验证。From the paper: "Due to computation resource constraint, we do not scale DSD to larger model sizes aligned with our baselines." The largest model so far, DSD-B, has only 205M parameters, far below the 675M main model of DiT-XL/2 and well short of the 945M scale of MAR. The performance ceiling and the scaling behaviour of DSD at large model sizes remain unverified.

未验证无监督学习能力(stated)Unsupervised learning ability not verified (stated)

论文原文:"We do not conduct experiments for verifying the effectiveness of our DSD as an unsupervised learning method." DSD 框架理论上兼容无监督表征学习,但论文中仅评估了类别条件图像生成任务,未在下游迁移学习或线性探针等任务上验证表征质量。From the paper: "We do not conduct experiments for verifying the effectiveness of our DSD as an unsupervised learning method." The DSD framework is in principle compatible with unsupervised representation learning, yet the paper only evaluates class-conditional image generation and never checks representation quality on downstream transfer or linear-probing tasks.

sFID 指标相对较弱(inferred)Relatively weak sFID (inferred)

DSD-B 的 sFID(空间 FID)为 8.96(无 CFG),而 REPA-E 仅为 4.17,LightningDiT 为 4.37。 sFID 对局部空间结构更敏感,较弱的 sFID 表明 DSD 在细粒度空间一致性上仍有提升空间, 可能与单一 ViT 共享特征导致编解码与扩散之间存在任务冲突有关。The sFID (spatial FID) of DSD-B is 8.96 (without CFG), whereas REPA-E reaches 4.17 and LightningDiT 4.37. sFID is more sensitive to local spatial structure, so the weaker sFID suggests DSD still has room to improve on fine-grained spatial consistency, possibly because sharing a single ViT creates a task conflict between coding/decoding and diffusion.

仅在 ImageNet 上验证,泛化性待确认(inferred)Validated on ImageNet only; generalization yet to be confirmed (inferred)

所有实验均在 ImageNet 256×256 类别条件生成任务上进行,未验证文本条件生成、更高分辨率(如 512×512、1024×1024)或其他领域(医学图像、视频生成等)的有效性。 端到端联合训练在数据分布差异较大时是否仍能防止 Latent Collapse 尚不明确。All experiments are run on class-conditional generation on ImageNet 256×256; text-conditional generation, higher resolutions (512×512, 1024×1024) and other domains (medical imaging, video generation and so on) are untested. Whether end-to-end joint training can still prevent Latent Collapse under a markedly different data distribution is unclear.