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

原创 让你更懂AI的 2025-10-11 18:10 北京

只加1个Token，收敛快到飞

只需引入一个 class token，REG 就让 Diffusion Transformer 的训练速度飙升至 63 倍，几乎“零成本”实现了更快收敛与更优生成——这项来自 NeurIPS 2025 Oral 的工作，通过将视觉模型的高层语义表征与生成模型的低层 latent 深度纠缠，重塑了「理解与生成」的关系，让 Diffusion 不再只是会“画”，而是真正“懂得在画”。

总览：REG 是一种简单而有效的方法，仅通过引入一个 class token 便能大幅加速生成模型的训练收敛。其将基础视觉模型（如 DINOv2）的 class token 与 latent 在空间维度拼接后共同加噪训练，从而显著提升 Diffusion 的收敛速度与性能上限。

在 ImageNet 256×256 上，REG 相比 SiT 和 REPA 分别实现了 63 倍和 23 倍的收敛加速，如下图所示：

论文题目：

Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think

论文地址：

https://arxiv.org/abs/2507.01467v2

代码地址：

https://github.com/Martinser/REG

会议信息：

NeurIPS 2025 Oral

作者机构：

NKIARI, 深圳福田；VCIP, CS, 南开大学；极豪科技；哈佛大学；中国科学院大学

背景简述

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方法介绍

2.1 REPA方案回顾

REPA 在训练阶段（Figure 2a）通过将 SiT 的中间 noised latent 与预训练视觉模型 DINOv2 的 clean dense feature 进行对齐，从而引入判别式语义指导并加速收敛。

而在推理阶段（Figure 2b），REPA 无法再使用这一外部对齐机制，即无法直接获得和利用 DINOv2 的表征来提升生成效果。因为训练时其只是间接对齐 DINOv2 与 SiT，而并未将 DINOv2 的表征显式作为 SiT 的输入。

2.2 REG训练和推理方案

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

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

noises_cls = torch.randn_like(cls_token)
cls_target = d_alpha_t * cls_token + d_sigma_t * noises_cls
cls_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) 保持相同通道维度。

cls_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，进行空间维度的拼接。

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

4. 进行 SiT 的 forward。

5. 计算对应的 denoising_loss_cls。

denoising_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)；

z = 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。

ImageNet 256×256，REG 使用 CFG，480 epochs 的 1.40 FID 超越 REPA800 epochs1.42 FID，同时 800 epochs 更是达到 1.36 FID。

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

REG消融实验

4.1 开销对比

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

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

4.2 不同超参的对比

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

4.3 不同操作对比

作者评估了不同内容作为 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，在保持计算效率的同时，同步提升语义质量和输出质量。

4.4 单独插入class token效果

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

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

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 的理解，后层则聚焦于生成任务，从而实现“先理解、后生成”，最终带来优异的生成效果。

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

阅读原文

跳转微信打开