跳转至

Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?

会议: ECCV2026
论文: ECCV 原文
代码: https://github.com/csslc/Self-Transcendence
领域: 图像生成
关键词: 扩散 Transformer、VAE 结构引导、特征级 CFG、表示对齐、训练加速

一句话总结

Self-Transcendence 先用干净 VAE 潜变量培养内部教师,再通过特征级无分类器引导增强其语义,用冻结教师指导从头训练的扩散模型;在不使用采样 CFG 的 ImageNet 256 实验中,SiT-XL/2 训练 80 epoch 的 FID 从 17.63 降到 7.51,教师预热成本另计。

研究背景与动机

扩散 Transformer 的最终去噪损失需要经过很长的梯度路径才能影响浅层,浅层因而较慢学会从带噪输入中分离有用结构。REPA 用 DINO 特征直接约束中间表示,让模型更早获得语义与空间组织,但同时引入了专门的外部视觉编码器及其预训练依赖。

SRA 和 LayerSync 尝试用模型自身的深层表示教浅层,问题是训练早期的深层特征同样不成熟,持续变化的监督目标难以提供可靠起点。本文将有效监督拆成两个条件:结构足够干净,才能区分信号与噪声;语义足够可分,才能让浅层理解条件对应的内容。VAE 潜变量满足前者,却不能单独替代 DINO 的语义能力。

作者因此不直接让一个尚未成熟的模型边学边教,而是先培养一个结构较好的内部教师,再增强它的条件语义。核心 idea:先用原有 VAE 的干净潜变量完成结构预热,再把 CFG 从预测输出搬到中间特征,用固定的内部教师替代外部表示监督。

方法详解

整体框架

输入是标准潜空间扩散训练所用的带噪潜变量及类别或文本条件,输出仍是原骨干的扩散预测。图 3 的第一阶段训练一个与目标骨干同架构的引导模型,在扩散损失之外,用干净 VAE 潜变量监督其中间层。

预热后冻结该模型,提取同一输入、同一层的有条件与无条件特征,并用特征级 CFG 构造监督目标。第二阶段重新初始化学生,而不是继续微调教师;学生浅层通过 MLP 与该目标对齐,同时优化扩散损失。表示对齐只在早期启用,之后移除教师分支,继续标准扩散训练。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["带噪潜变量与条件"] --> Structure["VAE 结构引导"]
    Clean["干净 VAE 潜变量"] --> Structure
    Structure --> Frozen["预热后冻结教师"]
    Frozen --> Feature["特征级 CFG"]
    Feature --> Align["早停表示对齐"]
    Student["重新初始化学生"] --> Align
    Align --> Output["仅扩散损失继续训练"]

关键设计

1. VAE 结构引导:先让教师学会识别干净结构

第一阶段取模型某个中间层的表示,经轻量 MLP 投影到 VAE 潜变量空间,用 L2 损失对齐当前训练图像的干净潜变量。监督来自潜空间扩散本来就需要的 VAE,而不是再接一个 DINO;它给浅层提供较短的训练信号路径,帮助模型从带噪输入中恢复空间结构。

原文式 (1) 的缓存数学排版有缺失。依据相邻正文,下式用常见平方 L2 记法重述该目标;原文未在可读文本中明确展示归约方式,应以原 PDF 或实现为准:

\[ \mathcal{L}_{\mathrm{VAE-guide}}=\|\mathrm{MLP}(\mathbf{f}_n)-\mathbf{z}\|_2^2. \]

其中 \(\mathbf{z}\) 是干净 VAE 潜变量,而非带噪训练输入;\(\mathbf{f}_n\) 是第 \(n\) 层表示。原扩散损失仍保留,辅助监督并不把整个任务改成单纯的潜变量重建。

为什么不直接用 VAE 特征指导整个最终模型?它们结构清晰,但类别语义较弱。表 4 的 V1 证明仅保留这一机制确有收益,却明显不及完整方法。因此第一阶段主要是在为第二阶段准备合适教师,而非最终语义监督的终点。

教师默认预热 40 epoch,随后参数固定。监督特征仍会随输入和条件变化,但不再因教师参数持续更新而漂移,这一点不同于在线 EMA 教师或同步层间学习。

2. 特征级 CFG:放大教师内部的条件语义

教师面对同一带噪输入,分别提取有条件特征 \(\mathbf{f}_c\) 和无条件特征 \(\mathbf{f}_u\)。本文将通常用于输出预测的 CFG 组合移到中间层,通过放大两者之差,使目标特征更突出类别或文本条件相关的信息。

按式 (3) 及其文字说明,特征组合为:

\[ \mathbf{f}_g=\mathbf{f}_u+\omega(\mathbf{f}_c-\mathbf{f}_u). \]

这里 \(\omega=1.0\) 对应普通有条件特征,而不是完全没有教师。更大的权重增强条件方向,但不意味着越大越好:表 5 中从 30.0 增到 60.0 后,FID 从 22.91 恶化到 23.57。

SiT 默认使用 \(\omega=30.0\),LightningDiT 使用 \(\omega=10.0\)。这些是训练监督目标的特征权重,不能误写成图像采样时的 CFG 权重;例如图 4、图 6 的采样可视化使用的是 4.0。

特征取自教师较深层,但并非越靠近输出越好。设骨干共有 \(N\) 个 Transformer 块,默认用教师第 \(2N/3\) 层指导学生第 \(N/2\) 层;保留适度的抽象差距,同时避免过深表示与学生浅层需求不匹配。

3. 早停表示对齐:只在学生最需要监督时施加约束

新学生的指定层经过 MLP 后,以 L2 损失匹配冻结教师产生的 \(\mathbf{f}_g\)。这同时利用了第一阶段获得的结构先验和第二个设计增强的条件语义,而不是仅对齐未经处理的教师特征。

式 (4) 的缓存排版同样受损。以下仍以平方 L2 形式重述正文描述;它表达训练机制,不声称补齐了实现中的全部归约细节:

\[ \mathcal{L}_{\mathrm{guide}}=\|\mathrm{MLP}(\mathbf{f}_m)-\mathbf{f}_g\|_2^2, \qquad \mathcal{L}=\mathcal{L}_{\mathrm{diff}}+\lambda_{\mathrm{guide}}\mathcal{L}_{\mathrm{guide}}. \]

教师冻结使监督目标不跟随学生共同优化。MLP 负责将学生中间表示映射到教师特征空间;正文把具体 MLP 架构留给补充材料,因此不能据当前缓存指定层数或隐藏维度。

作者观察到,持续强制浅层语义对齐会令后续深层训练不稳定,所以只在早期启用辅助损失。停止对齐后,学生仍继续训练,并非整个训练提前结束;推理也不需要教师或这条辅助监督分支。

损失函数 / 训练策略

默认 \(\lambda_{\mathrm{guide}}=0.5\),扩散损失权重为 1.0。第 4.1 节规定 base 模型在前 20 epoch 使用表示对齐,更大模型使用前 10 epoch,之后只保留扩散损失。

实验采用 patch size 为 2 的 SiT 和 patch size 为 1 的 LightningDiT,并沿用各骨干原有训练、推理配置;教师与学生架构相同。潜空间覆盖 SD-VAE 与 VAVAE,不能将“无新增外部表示教师”理解为“不使用预训练 VAE”。

表 5 的 SiT-B/2 默认层配对是教师第 8 层到学生第 6 层,教师训练 200K iteration,即文中默认的 40 epoch。较久预热并非稳定更好:300K iteration 的 FID 为 23.05,略逊于 200K 的 22.91。

当前缓存只有主文与参考文献。正文提到的单阶段比较、MLP 细节、512 分辨率设置及更完整文本到图像协议均需补充材料核验,本笔记不补造这些参数。

实验关键数据

主实验

下表整理原文表 1、表 2、表 3 的关键结果,FID 越低越好;不同任务与采样设置分别列明。表 1 和表 3 使用 50,000 张生成样本,列出的训练长度是学生预算,不含 40 epoch 教师预热。

原表 / 任务 骨干与设置 基线或 REPA Self-Transcendence 比较边界
表 1 / ImageNet 256 SiT-B/2,无采样 CFG,80 epoch 基线 36.14;REPA 24.40 FID 20.49 130M 参数
表 1 / ImageNet 256 SiT-L/2,无采样 CFG,80 epoch 基线 21.41;REPA 9.70 FID 8.74 458M 参数
表 1 / ImageNet 256 SiT-XL/2,无采样 CFG,80 epoch 基线 17.63;REPA 7.90 FID 7.51 675M 参数
表 1 / ImageNet 256 LightningDiT-XL/1,无采样 CFG,64 epoch 基线 5.30;REPA 4.09 FID 3.55 使用 VAVAE
表 2 / 文本到图像 MMDiT,150K iteration REPA:FID 4.90,IS 32.55 FID 4.56,IS 33.08 正文称沿用 REPA 协议
表 3 / ImageNet 256 SiT-XL/2,使用采样 CFG REPA:800 epoch,FID 1.42 400 epoch,FID 1.44 接近而非更低 FID
表 3 / ImageNet 256 LightningDiT-XL/1,使用采样 CFG 基线:800 epoch,FID 1.35 400 epoch,FID 1.25 学生预算不同

消融实验

原文表 4 使用 SiT-B/2、80 epoch,并仅用 10,000 张样本评估。其 FID 22.91 与主表的 20.49 不是同一评估样本数,不能据此判为互相矛盾。

配置 VAE 结构引导 自引导表示 FID IS
基线 无 无 38.60 41.95
V1 有 无 32.20 52.38
V2 无 有,教师仅用扩散损失训练 25.21 63.83
完整方法 有 有 22.91 70.37

关键发现

  • 表 4 中,仅结构监督把 FID 从 38.60 降到 32.20;完整方法相较 V2 再降低 2.30,表明教师预热质量与后续自引导互补。
  • 表 5 中,特征 CFG 权重为 1.0 时 FID 是 29.30,默认 30.0 时为 22.91。普通有条件特征不足以复现完整收益。
  • 同一表中,教师第 8 层指导学生第 8 层的 FID 为 24.29,逊于指导第 6 层的 22.91;选择监督位置本身是重要超参数。
  • 表 6 在 8 张 A800、batch size 256、学生训练 400 epoch 的口径下,报告基线 58.87 小时、REPA 65.44 小时、本文 60.45 小时。教师另需 6.39 小时,若从零计入则合计 66.84 小时,不能据学生耗时直接宣称总成本优于 REPA。
  • 表 6 的前 50K iteration 峰值显存是 12.34 GB/卡,高于 REPA 的 10.39;早停后恢复 9.03。该成本实验按前 50K、后 1950K 分段报告,不应强行改写成所有设置统一的早停时间。

亮点与洞察

  • 结构与语义不是同一种监督。干净 VAE 潜变量先解决“哪里是有效信号”,特征 CFG 再增强“哪些信号与条件相关”,使两阶段的必要性可以用 V1、V2 分别检验。
  • CFG 不只服务最终采样。把条件与无条件的差分放大后用于表示监督,为其他条件生成模型提供了一个可检验的内部教师构造方式。
  • 教师分类能力并不是生成收益的唯一代理。表 7 在噪声时刻 0.5 做线性探测,第 8 层本文为 41.18%,REPA 为 51.87%,但本文仍能在部分生成设置中获得更低 FID。

局限与展望

  • 作者明确承认,内部监督质量仍受骨干自身容量限制,并引入引导权重、层配对、预热长度等额外超参数;跨任务迁移不能保证直接沿用默认值。
  • 已验证范围是图像生成,包括类别条件与文本条件任务。文生视频、文生 3D 尚未验证,不能由图像结果外推。
  • 两阶段需要额外教师预热,短训练预算或教师只使用一次时,这部分固定成本更显著;完整成本比较应同时报告教师和学生训练。
  • LightningDiT 实验使用 VAVAE,而主文相关工作说明 VA-VAE 本身涉及视觉基础模型对齐。因此“self-contained”宜限定为本方法没有新增外部特征教师,不是整个已有骨干与 tokenizer 的预训练都完全不依赖外部表示。
  • 本缓存未提供补充材料和多随机种子误差条。应进一步核验复现细节,并评估 7.90 与 7.51 这类小差距的统计稳定性。

相关工作与启发

  • vs REPA:REPA 对齐外部 DINO 表示;本文对齐 VAE 预热、特征 CFG 增强后的内部教师。表 1 支持部分同预算优势,但表 3 的 SiT FID 1.44 并未优于 REPA 的 1.42。
  • vs SRA / LayerSync:两者在训练过程中利用自身深层特征;本文先预热再冻结教师,并从头训练学生,用额外阶段换取稳定且更有条件判别力的监督。
  • vs Dispersive Loss:后者鼓励内部表示分散而不需要预训练教师;本文增加显式目标特征与对齐分支,取得更好生成结果,但训练流程和调参更复杂。

评分

  • 新颖性: 4/5,将 VAE 结构预热与特征级 CFG 结合成明确的内部教师方案。
  • 实验充分度: 4/5,覆盖多规模骨干、两类生成任务和组件消融,但完整成本与补充设置需谨慎核对。
  • 写作质量: 4/5,结构与语义的动机清晰,成本口径和部分实现细节仍需要额外说明。
  • 价值: 4/5,为减少新增外部表示依赖提供了有实证支持的训练路径。