REG 是一種簡(jiǎn)單而有效的方法，僅通過(guò)引入一個(gè) class token 便能大幅加速生成模型的訓(xùn)練收斂。其將基礎(chǔ)視覺(jué)模型（如 DINOv2）的 class token 與 latent 在空間維度拼接后共同加噪訓(xùn)練，從而顯著提升 Diffusion 的收斂速度與性能上限。在 ImageNet 256×256 上，REG 相比 SiT 和 REPA 分別實(shí)現(xiàn)了 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 及其變體通過(guò)把 noised latent 和視覺(jué)基礎(chǔ)模型判別式 clean image representation 進(jìn)行對(duì)齊，來(lái)加速 Diffusion 訓(xùn)練收斂，但其推理時(shí)無(wú)法使用外部特征，限制了其性能上限。

為此作者提出 REG（Representation Entanglement for Generation）：將low-level latent 與預(yù)訓(xùn)練視覺(jué)模型的high-level class token糾纏，并在訓(xùn)練中對(duì)二者同時(shí)加噪、聯(lián)合去噪優(yōu)化，從而使模型具備直接從純?cè)肼暽蓤D像-類(lèi)別對(duì)的能力。

REG 在顯著提升生成質(zhì)量的同時(shí)大幅加速訓(xùn)練收斂，僅需額外引入一個(gè) token（計(jì)算開(kāi)銷(xiāo) <0.5%），幾乎不增加推理成本。此外，推理過(guò)程中 REG 能同步生成圖像 latent 及其全局語(yǔ)義信息，并利用這些語(yǔ)義知識(shí)主動(dòng)指導(dǎo)和增強(qiáng)圖像生成。

在 ImageNet 256×256 上，SiT-XL/2+REG 收斂速度分別較 SiT-XL/2 與 SiT-XL/2+REPA 提升 63 倍與 23 倍；僅 400K 步的 SiT-L/2+REG 已優(yōu)于 4M 步的 SiT-XL/2+REPA。

二、REG方法介紹

2.1 REPA方案回顧

REPA 在訓(xùn)練階段（Figure 2a）通過(guò)將 SiT 的中間 noised latent 與預(yù)訓(xùn)練視覺(jué)模型 DINOv2 的 clean dense feature 進(jìn)行對(duì)齊，從而引入判別式語(yǔ)義指導(dǎo)并加速收斂。

而在推理階段（Figure 2b），REPA 無(wú)法再使用這一外部對(duì)齊機(jī)制，即無(wú)法直接獲得和利用 DINOv2 的表征來(lái)提升生成效果。因?yàn)橛?xùn)練時(shí)其只是間接對(duì)齊 DINOv2 與 SiT，而并未將 DINOv2 的表征顯式作為 SiT 的輸入。

2.2 REG訓(xùn)練和推理方案

REG 訓(xùn)練方案非常簡(jiǎn)單，只需將 DINOv2 的 class token（攜帶全局語(yǔ)義信息）與原始 latent 一起加噪，然后拼接后輸入 SiT 進(jìn)行去噪訓(xùn)練。具體偽代碼流程如下：

1. DINOv2 class token 按照標(biāo)準(zhǔn)流程進(jìn)行加噪成為 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 通過(guò)一個(gè) linear 層，進(jìn)行映射轉(zhuǎn)換，和 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 通過(guò) torch.cat，進(jìn)行空間維度的拼接。

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

1. 進(jìn)行 SiT 的 forward。

2. 計(jì)算對(duì)應(yīng)的 denoising_loss_cls。

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

REG 的推理方案同樣簡(jiǎn)潔：在原有 SiT 推理流程的基礎(chǔ)上，只需將額外隨機(jī)初始化的 class token 一起加噪并參與去噪即可，具體代碼參考：

https://github.com/Martinser/REG/blob/main/samplers.py

1. 隨機(jī)初始化的 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 一起進(jìn)行聯(lián)合推理。

在 ImageNet 256×256 上，REG 在不使用 CFG 且不改造 SiT 與 VAE 的前提下，顯著超越 REPA。

SiT-XL/2+REG 相比 SiT-XL/2與SiT-XL/2+REPA 分別實(shí)現(xiàn) 63 倍和 23 倍的收斂加速，僅 400K 步的 SiT-L/2+REG 即優(yōu)于 4M 步（10 倍時(shí)長(zhǎng)）的 SiT-XL/2+REPA。而在 4M 步訓(xùn)練下，REG 的 FID 進(jìn)一步達(dá)到 1.8。

ImageNet 256×256，REG 使用 CFG，480 epochs 的 1.40 FID 超越 REPA 800 epochs 1.42 FID，同時(shí) 800 epochs 更是達(dá)到 1.36 FID。

ImageNet 512×512，REG 的表現(xiàn)也非常不錯(cuò)，REG 使用 CFG，80 epochs FID 達(dá)到 1.68，超越 REPA 200 epochs 和 SiT 600 epochs 結(jié)果。

四、REG消融實(shí)驗(yàn)4.1 開(kāi)銷(xiāo)對(duì)比

訓(xùn)練開(kāi)銷(xiāo)對(duì)比。REG 和 SiT-XL/2 相比，達(dá)到相似的 FID，REG 訓(xùn)練時(shí)間減少了 97.90%。同時(shí)和 REG 相比，也達(dá)到相似的 FID，REG 訓(xùn)練時(shí)間減少了 95.72%。

推理開(kāi)銷(xiāo)對(duì)比。可以看到 REG 整體的開(kāi)銷(xiāo)幾乎沒(méi)有，Params，F(xiàn)LOPs 和 Latency 增加的開(kāi)銷(xiāo)小于 0.5%，但是 FID 卻比 SiT-XL/2 + REPA 大幅提高 56.46%。

4.2 不同超參的對(duì)比

作者對(duì)不同對(duì)齊目標(biāo)，不同深度，還有 class token 的去噪 loss 的權(quán)重 (velocity prediction loss)，進(jìn)行廣泛的消融實(shí)驗(yàn)，證明 REG 有效性。

4.3 不同操作對(duì)比

作者評(píng)估了不同內(nèi)容作為 SiT 輸入的影響，具體對(duì)比如下：

1. one learnable token，把 DINOv2 的 class token，換成一個(gè)可學(xué)習(xí)的 token；

2. avg (latent feature)，改進(jìn)有限，缺乏判別性語(yǔ)義；

3. avg (DINOv2 feature)，顯著提升；

4. DINOv2 class token，效果最佳。

表明 high-level 全局判別信息（avg (DINOv2 feature) 和 class token）能大幅提升生成質(zhì)量，能有效規(guī)范生成的 latent，在保持計(jì)算效率的同時(shí)，同步提升語(yǔ)義質(zhì)量和輸出質(zhì)量。

4.4 單獨(dú)插入class token效果

研究 SiT 只插入 class token，不使用 REPA 效果。發(fā)現(xiàn)插入 class token，并進(jìn)行聯(lián)合加噪訓(xùn)練，依舊有非常不錯(cuò)的效果。

4.5 增強(qiáng)生成模型的判別語(yǔ)義學(xué)習(xí)

PS：CKNNA 是一個(gè)類(lèi)似相似度的指標(biāo)，REG 遵循 REPA 的設(shè)置，計(jì)算 REG 和 DINOv2 之間的 CKNNA，CKNNA 數(shù)值越大，代表 REG 和 DINOv2 越相似，REG 學(xué)習(xí)到了更多判別式語(yǔ)義。

Figure 3 系統(tǒng)分析了 REG 是否學(xué)習(xí)到了判別式語(yǔ)義，發(fā)現(xiàn) REG 和判別式語(yǔ)義相關(guān)的指標(biāo) CKNNA，能在不同 block，layer，timesteps 都比 REPA 更高，具體情況如下：

1. 不同 Training steps (Figure 3(a))：隨訓(xùn)練時(shí)間增加，F(xiàn)ID 和 CKNNA 均提升，且“更高 CKNNA → 更好生成”。REG 在各訓(xùn)練階段均優(yōu)于 REPA，說(shuō)明引入 DINOv2 class token 增強(qiáng)了判別語(yǔ)義。

2. 不同的 layers，F(xiàn)igure 3(b) 發(fā)現(xiàn)各模型在第 8 層語(yǔ)義分?jǐn)?shù)達(dá)峰值后下降（此處進(jìn)行 SiT 和 DINOv2 dense feature 的對(duì)齊，并計(jì)算對(duì)齊損失），但 REG 在所有層始終高于 REPA 和 SiT。

   這得益于其將 DINOv2 low-level latent 與 class token 融合，通過(guò) attention 傳遞判別語(yǔ)義，使前中層專(zhuān)注理解 latent，后層專(zhuān)注生成細(xì)節(jié)。

3. 不同的 timesteps，F(xiàn)igure 3(c) 呈現(xiàn) REG 在所有時(shí)間步保持顯著優(yōu)勢(shì)，驗(yàn)證了其在全噪聲范圍內(nèi)穩(wěn)定的語(yǔ)義引導(dǎo)能力。

整體而言，REG 是一種極其簡(jiǎn)潔而高效的方案，首次提出了 high-level 與 low-level token 混合糾纏去噪的訓(xùn)練范式，在顯著提升訓(xùn)練效率與收斂速度的同時(shí)完全不增加推理開(kāi)銷(xiāo)。

其核心在于促進(jìn)生成模型的“理解—生成”解耦：class token 指引 SiT 前層更專(zhuān)注于 noise latent 的理解，后層則聚焦于生成任務(wù)，從而實(shí)現(xiàn)“先理解、后生成”，最終帶來(lái)優(yōu)異的生成效果。

來(lái)源：公眾號(hào)【PaperWeekly】

llustration From IconScout By IconScout Store

-The End-

本周上新！

掃碼觀看！

“AI技術(shù)流”原創(chuàng)投稿計(jì)劃

TechBeat是由將門(mén)創(chuàng)投建立的AI學(xué)習(xí)社區(qū)（www.techbeat.net）。社區(qū)上線(xiàn)700+期talk視頻，3000+篇技術(shù)干貨文章，方向覆蓋CV/NLP/ML/Robotis等；每月定期舉辦頂會(huì)及其他線(xiàn)上交流活動(dòng)，不定期舉辦技術(shù)人線(xiàn)下聚會(huì)交流活動(dòng)。我們正在努力成為AI人才喜愛(ài)的高質(zhì)量、知識(shí)型交流平臺(tái)，希望為AI人才打造更專(zhuān)業(yè)的服務(wù)和體驗(yàn)，加速并陪伴其成長(zhǎng)。

投稿內(nèi)容

// 最新技術(shù)解讀/系統(tǒng)性知識(shí)分享 //

// 前沿資訊解說(shuō)/心得經(jīng)歷講述 //

投稿須知

稿件需要為原創(chuàng)文章，并標(biāo)明作者信息。

我們會(huì)選擇部分在深度技術(shù)解析及科研心得方向，對(duì)用戶(hù)啟發(fā)更大的文章，做原創(chuàng)性?xún)?nèi)容獎(jiǎng)勵(lì)

投稿方式

發(fā)送郵件到

michellechang@thejiangmen.com

或添加工作人員微信（michelle333_）投稿，溝通投稿詳情

關(guān)于我“門(mén)”

將門(mén)是一家以專(zhuān)注于數(shù)智核心科技領(lǐng)域的新型創(chuàng)投機(jī)構(gòu)，也是北京市標(biāo)桿型孵化器。 公司致力于通過(guò)連接技術(shù)與商業(yè)，發(fā)掘和培育具有全球影響力的科技創(chuàng)新企業(yè)，推動(dòng)企業(yè)創(chuàng)新發(fā)展與產(chǎn)業(yè)升級(jí)。

將門(mén)成立于2015年底，創(chuàng)始團(tuán)隊(duì)由微軟創(chuàng)投在中國(guó)的創(chuàng)始團(tuán)隊(duì)原班人馬構(gòu)建而成，曾為微軟優(yōu)選和深度孵化了126家創(chuàng)新的技術(shù)型創(chuàng)業(yè)公司。

如果您是技術(shù)領(lǐng)域的初創(chuàng)企業(yè)，不僅想獲得投資，還希望獲得一系列持續(xù)性、有價(jià)值的投后服務(wù)，歡迎發(fā)送或者推薦項(xiàng)目給我“門(mén)”:

bp@thejiangmen.com

點(diǎn)擊右上角，把文章分享到朋友圈