NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!

机器之心 2025-11-29 09:48

REG 是一种简单而有效的方法,仅通过引入一个 class token 便能大幅加速生成模型的训练收敛。其将基础视觉模型(如 DINOv2)的 class token 与 latent 在空间维度拼接后共同加噪训练,从而显著提升 Diffusion 的收敛速度与性能上限。在 ImageNet 256×256 上,REG 相比 SiT 和 REPA 分别实现了 63 倍和 23 倍的收敛加速,如下图所示:

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图2

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图3

 

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图4

 

背景简述

 

REPA 及其变体通过把 noised latent 和视觉基础模型判别式 clean image representation 进行对齐,来加速 Diffusion 训练收敛,但其推理时无法使用外部特征,限制了其性能上限。

 

为此作者提出 REG(Representation Entanglement for Generation):将 low-level latent 与预训练视觉模型的 high-level class token 纠缠,并在训练中对二者同时加噪、联合去噪优化,从而使模型具备直接从纯噪声生成图像 - 类别对的能力。REG 在显著提升生成质量的同时大幅加速训练收敛,仅需额外引入一个 token(计算开销 <0.5%),几乎不增加推理成本。此外,推理过程中 REG 能同步生成图像 latent 及其全局语义信息,并利用这些语义知识主动指导和增强图像生成。

 

在 ImageNet 256×256 上,SiT-XL/2+REG 收敛速度分别较 SiT-XL/2 与 SiT-XL/2+REPA 提升 63 倍与 23 倍;仅 400K 步的 SiT-L/2+REG 已优于 4M 步的 SiT-XL/2+REPA。

 

REG 方法介绍

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图5

 

REPA 方案回顾

 

REPA 在训练阶段(Figure 2a)通过将 SiT 的中间 noised latent 与预训练视觉模型 DINOv2 的 clean dense feature 进行对齐,从而引入判别式语义指导并加速收敛。而在推理阶段(Figure 2b),REPA 无法再使用这一外部对齐机制,即无法直接获得和利用 DINOv2 的表征来提升生成效果。因为训练时其只是间接对齐 DINOv2 与 SiT,而并未将 DINOv2 的表征显式作为 SiT 的输入。

 

REG 训练和推理方案

 

REG 训练方案非常简单,只需将 DINOv2 的 class token(携带全局语义信息)与原始 latent 一起加噪,然后拼接后输入 SiT 进行去噪训练。具体伪代码流程如下:

 

1. DINOv2 class token 按照标准流程进行加噪成为 noised class token (cls_input)。

 

Pythonnoises_cls = torch.randn_like(cls_token)cls_target = d_alpha_t * cls_token + d_sigma_t * noises_clscls_input = alpha_t.squeeze(-1).squeeze(-1) * cls_token + sigma_t.squeeze(-1).squeeze(-1) * noises_cls

 

2. noised class token 通过一个 linear 层,进行映射转换,和 noised latent (x) 保持相同通道维度。

 

Pythoncls_token = self.cls_projectors (cls_input)cls_token = self.norm (cls_token)cls_token = cls_token.unsqueeze (1)  

 

3. noised class token 和 noised latent 通过 torch.cat,进行空间维度的拼接。

Pythonx = torch.cat ((cls_token, x), dim=1)

 

4. 进行 SiT 的 forward。

 

5. 计算对应的 denoising_loss_cls。

 

Pythondenoising_loss_cls = mean_flat ((cls_output - cls_target) ** 2)

 

REG 的推理方案同样简洁:在原有 SiT 推理流程的基础上,只需将额外随机初始化的 class token 一起加噪并参与去噪即可,具体代码参考:https://github.com/Martinser/REG/blob/main/samplers.py

 

1. 随机初始化的 class token (cls_z) 和 latent (z)

 

Pythonz = torch.randn (n, model.in_channels, latent_size, latent_size, device=device)cls_z = torch.randn (n, args.cls, device=device)

 

2. class token 和 latent 一起进行联合推理

 

REG 具体效果

 

在 ImageNet 256×256 上,REG 在不使用 CFG 且不改造 SiT 与 VAE 的前提下,显著超越 REPA。SiT-XL/2+REG 相比 SiT-XL/2 与 SiT-XL/2+REPA 分别实现 63 倍和 23 倍的收敛加速,仅 400K 步的 SiT-L/2+REG 即优于 4M 步(10 倍时长)的 SiT-XL/2+REPA。而在 4M 步训练下,REG 的 FID 进一步达到 1.8。

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图6

 

ImageNet 256×256,REG 使用 CFG,480 epochs 的 1.40 FID 超越 REPA 800 epochs 1.42 FID,同时 800 epochs 更是达到 1.36 FID。

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图7

 

ImageNet 512×512,REG 的表现也非常不错,REG 使用 CFG,80 epochs FID 达到 1.68,超越 REPA 200 epochs 和 SiT 600 epochs 结果。

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图8

 

REG 消融实验

 

开销对比

 

训练开销对比。REG 和 SiT-XL/2 相比,达到相似的 FID,REG 训练时间减少了 97.90%。同时和 REG 相比,也达到相似的 FID,REG 训练时间减少了 95.72%。

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图9

 

推理开销对比。可以看到 REG 整体的开销几乎没有,Params,FLOPs 和 Latency 增加的开销小于 0.5%,但是 FID 却比 SiT-XL/2 + REPA 大幅提高 56.46%。

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图10

 

不同超参的对比

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图11

 

作者对不同对齐目标,不同深度,还有 class token 的去噪 loss 的权重 (velocity prediction loss),进行广泛的消融实验,证明 REG 有效性。

 

不同操作对比

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图12

 

作者评估了不同内容作为 SiT 输入的影响,具体对比如下:

 

1. one learnable token,把 DINOv2 的 class token,换成一个可学习的 token

2. avg (latent feature),改进有限,缺乏判别性语义

3. avg (DINOv2 feature),显著提升

4. DINOv2 class token,效果最佳

 

表明 high-level 全局判别信息(avg (DINOv2 feature) 和 class token)能大幅提升生成质量,能有效规范生成的 latent,在保持计算效率的同时,同步提升语义质量和输出质量。

 

单独插入 class token 效果

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图13

 

研究 SiT 只插入 class token,不使用 REPA 效果。发现插入 class token,并进行联合加噪训练,依旧有非常不错的效果。

 

增强生成模型的判别语义学习

 

NeurIPS 2025 Oral | 1个Token零成本,REG让Diffusion训练收敛快20倍!图14

 

PS:CKNNA 是一个类似相似度的指标,REG 遵循 REPA 的设置,计算 REG 和 DINOv2 之间的 CKNNA,CKNNA 数值越大,代表 REG 和 DINOv2 越相似,REG 学习到了更多判别式语义。

 

Figure 3 系统分析了 REG 是否学习到了判别式语义,发现 REG 和判别式语义相关的指标 CKNNA,能在不同 block,layer,timesteps 都比 REPA 更高,具体情况如下:

 

1. 不同 Training steps (Figure 3 (a)): 随训练时间增加,FID 和 CKNNA 均提升,且 “更高 CKNNA → 更好生成”。REG 在各训练阶段均优于 REPA,说明引入 DINOv2 class token 增强了判别语义。

 

2. 不同的 layers,Figure 3 (b) 发现各模型在第 8 层语义分数达峰值后下降(此处进行 SiT 和 DINOv2 dense feature 的对齐,并计算对齐损失),但 REG 在所有层始终高于 REPA 和 SiT。这得益于其将 DINOv2 low-level latent 与 class token 融合,通过 attention 传递判别语义,使前中层专注理解 latent,后层专注生成细节。

 

3. 不同的 timesteps,Figure 3 (c) 呈现 REG 在所有时间步保持显著优势,验证了其在全噪声范围内稳定的语义引导能力。

 

总结

 

整体而言,REG 是一种极其简洁而高效的方案,首次提出了 high-level 与 low-level token 混合纠缠去噪的训练范式,在显著提升训练效率与收敛速度的同时完全不增加推理开销。其核心在于促进生成模型的 “理解 — 生成” 解耦:class token 指引 SiT 前层更专注于 noise latent 的理解,后层则聚焦于生成任务,从而实现 “先理解、后生成”,最终带来优异的生成效果。

© THE END 

声明:内容取材于网络,仅代表作者观点,如有内容违规问题,请联系处理。 
IP
more
Pipistrel的卖身,终得“北美”市场
苹果iPhone 17e全新曝光,主打卖点意想不到
NeurIPS 2025 | DynaAct:DeepSeek R1之外,探索大模型推理的另一条道路
iPhone可以检测针孔摄像头?官方回应
DeepSeek罗福莉正式亮相小米AI;「周杰伦概念股」联手宇树造 IP 机器人;乘用车百公里加速不少于5秒 | 极客早知道
宇树完成上市辅导,拟境内IPO;传索尼开发Labubu电影;2699元!原道X小岛秀夫耳机限量发售 | 极客早知道
砸 43 亿!长江存储是大股东的武汉新芯启动IPO
专为端侧大模型而生:解密AI计算加速“秘密武器”Arm China周易X3 NPU IP
苹果2nm A20 性能太强了!曝:苹果将跳过 19 直推 iPhone 20
杭州芯片“小巨人”,要IPO了
Copyright © 2025 成都区角科技有限公司
蜀ICP备2025143415号-1
  
川公网安备51015602001305号