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

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

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

1. 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)

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

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

1. 进行 SiT 的 forward。

2. 计算对应的 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)

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

在 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 超越 REPA 800 epochs 1.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】

llustration From IconScout By IconScout Store

-The End-

本周上新！

扫码观看！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线700+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

michellechang@thejiangmen.com

或添加工作人员微信（michelle333_）投稿，沟通投稿详情

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈