文章：??https://arxiv.org/pdf/2505.02831?? 

代碼：https://github.com/vvvvvjdy/SRA  

亮點(diǎn)直擊

• 分析發(fā)現(xiàn)：文章深入分析diffusion transformer的潛在表征趨勢(shì)，發(fā)現(xiàn)當(dāng)只執(zhí)行生成訓(xùn)練時(shí)，其會(huì)擁有一個(gè)大致從粗糙到細(xì)致，從差到好的表征（判別）過(guò)程。
• 方法提出：文章引入SRA方法，通過(guò)將早期層高噪聲條件下的潛在表征與后期層低噪聲條件下的潛在表征對(duì)齊，實(shí)現(xiàn)自我表征引導(dǎo)，提升diffusion transforme的訓(xùn)練速度和生成表現(xiàn)。
• 實(shí)驗(yàn)驗(yàn)證：文章通過(guò)一系列實(shí)驗(yàn)驗(yàn)證了SRA對(duì)DiTs（diffusion-based）和SiTs（flow-based）有顯著性能提升，并與依賴復(fù)雜表征訓(xùn)練框架或強(qiáng)大外部表征先驗(yàn)的方法進(jìn)行對(duì)比，證明了SRA的簡(jiǎn)單有效性。

背景介紹

在生成式預(yù)訓(xùn)練中引入表征指導(dǎo)

Diffusion transformer在視覺生成領(lǐng)域展現(xiàn)出了有效的可擴(kuò)展性和可遷移性。最近的研究發(fā)現(xiàn)高質(zhì)量的內(nèi)部表征能夠加速其生成訓(xùn)練，并提升生成樣本的質(zhì)量然而，現(xiàn)有方法在提供表征指導(dǎo)時(shí)，要么引入復(fù)雜訓(xùn)練框架，要么依賴大規(guī)模預(yù)訓(xùn)練模型，這限制了其應(yīng)用場(chǎng)景和資源效率。

先前表征引入方法的局限性

• 復(fù)雜訓(xùn)練框架：如MaskDiT和SD-DiT采用MAE或IBOT的訓(xùn)練范式，需要設(shè)計(jì)額外的訓(xùn)練任務(wù)和框架，增加了訓(xùn)練的復(fù)雜度和資源消耗。
• 外部表征先驗(yàn)依賴：如REPA利用DINOv2，CLIP等大規(guī)模預(yù)訓(xùn)練模型作為表征指導(dǎo)模型，雖然能大幅提升性能，但依賴于大量數(shù)據(jù)和計(jì)算資源訓(xùn)練得到的外部先驗(yàn)，限制了方法的靈活性和普適性。

靈感和發(fā)現(xiàn)

靈感來(lái)源

不同于將干凈的圖像作為輸入的表征模型然后輸出語(yǔ)義豐富的特征，擴(kuò)散模型通常以latent噪聲作為輸入并一步一步清理出干凈圖像。換句話說(shuō)，擴(kuò)散模型的生成機(jī)制操作宏觀上可以被認(rèn)為是一個(gè)coarse-to-fine 過(guò)程。因此，我們推測(cè)，擴(kuò)散模型中的表征也遵循這一趨勢(shì)。

發(fā)現(xiàn)

文章對(duì)普通的SiT和DiT進(jìn)行了實(shí)證分析。我們首先發(fā)現(xiàn)其中的潛在特征隨著塊層數(shù)的增加和噪聲水平的降低而逐步細(xì)化，呈現(xiàn)出大致從粗糙到精細(xì)的過(guò)程。接下來(lái)，我們觀察到僅通過(guò)生成式預(yù)訓(xùn)練的diffusion transformer已經(jīng)能夠?qū)W習(xí)到有意義的判別性表征。

同時(shí)，盡管在大約第20層達(dá)到峰值后，由于模型需要轉(zhuǎn)移注意力以生成高頻細(xì)節(jié)的圖像，準(zhǔn)確率有所下降，但通過(guò)增加塊層數(shù)和降低噪聲水平， 表征的質(zhì)量總體上是從差到好逐漸轉(zhuǎn)變的。

方法（SRA）

SRA的核心思想是利用擴(kuò)散Transformer自身在不同訓(xùn)練階段和噪聲條件下的表征差異，通過(guò)自我蒸餾的方式進(jìn)行表征對(duì)齊。具體來(lái)說(shuō)，SRA將早期層（高噪聲條件下）的潛在表征與后期層（低噪聲條件下）的潛在表征進(jìn)行對(duì)齊，從而實(shí)現(xiàn)自我表征引導(dǎo)。這種方法無(wú)需額外的表征組件或外部預(yù)訓(xùn)練模型，僅在diffusion transformer自身的生成訓(xùn)練過(guò)程中引入指導(dǎo)。

實(shí)現(xiàn)細(xì)節(jié)

• 表征對(duì)齊：將學(xué)生模型（可訓(xùn)練模型）早期層的潛在表征與教師模型（通過(guò)指數(shù)移動(dòng)平均更新權(quán)重的模型）后期層的潛在表征進(jìn)行對(duì)齊。學(xué)生模型的輸出通過(guò)輕量級(jí)可訓(xùn)練MLP頭進(jìn)行非線性變換后，與教師模型的輸出進(jìn)行對(duì)齊。
• 損失函數(shù)：SRA通過(guò)最小化教師模型輸出和學(xué)生模型輸出變體之間的patch-wise距離來(lái)實(shí)現(xiàn)自我表征對(duì)齊。同時(shí)，將該目標(biāo)函數(shù)與原始的擴(kuò)散模型目標(biāo)函數(shù)結(jié)合，進(jìn)行聯(lián)合學(xué)習(xí)。
• EMA教師網(wǎng)絡(luò)：為了避免直接使用同模型輸出作為監(jiān)督信號(hào)導(dǎo)致的訓(xùn)練不穩(wěn)定性，SRA采用指數(shù)移動(dòng)平均（EMA）策略構(gòu)建教師模型。教師模型的權(quán)重基于學(xué)生模型的權(quán)重進(jìn)行加權(quán)移動(dòng)平均更新，從而提供穩(wěn)定的表征引導(dǎo)。
• 超參數(shù)設(shè)置：在實(shí)驗(yàn)中，我們對(duì)超參數(shù)進(jìn)行了細(xì)致的調(diào)整，包括表征對(duì)齊的塊層選擇、時(shí)間間隔、對(duì)齊目標(biāo)函數(shù)以及EMA教師網(wǎng)絡(luò)的動(dòng)量系數(shù)等。這些超參數(shù)的選擇對(duì)SRA的性能有著顯著的影響。

實(shí)驗(yàn)

實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)中，文章采用了ImageNet 1K數(shù)據(jù)集，并遵循DiT和SiT的訓(xùn)練設(shè)置，包括使用AdamW優(yōu)化器、恒定學(xué)習(xí)率1e-4、批量大小256等。評(píng)估指標(biāo)包括FID、sFID、IS、Precision和Recall，使用相同的評(píng)價(jià)框架和參考圖像以確保與先前方法的公平比較。

元素級(jí)分析

文章通過(guò)詳細(xì)的組件分析，研究了SRA中不同設(shè)計(jì)選擇對(duì)性能的影響。實(shí)驗(yàn)結(jié)果表明：

• 塊層選擇：將教師模型的后期層與學(xué)生模型的早期層進(jìn)行對(duì)齊能取得最佳性能，因?yàn)樵缙趯有枰嗟闹笇?dǎo)以捕捉語(yǔ)義豐富的表征。
• 時(shí)間間隔：使用教師模型比學(xué)生模型低噪聲輸入的特征能提升性能，且時(shí)間間隔值在0.1左右時(shí)效果最佳。動(dòng)態(tài)時(shí)間間隔表現(xiàn)略優(yōu)于固定時(shí)間間隔。
• 對(duì)齊目標(biāo)函數(shù)：smooth-?1、?1和?2三種回歸訓(xùn)練目標(biāo)函數(shù)均能有效提升性能。
• 教師網(wǎng)絡(luò)更新策略：采用固定動(dòng)量系數(shù)0.9999的EMA策略更新教師網(wǎng)絡(luò)效果最佳，其他更新策略并不適合SRA。
• 投影頭的作用：使用輕量級(jí)投影頭對(duì)學(xué)生的輸出進(jìn)行后處理，能顯著提升性能，因?yàn)樗兄谧屇Ｐ驮谶M(jìn)行非線性變換前捕捉更有效的隱藏表征以便對(duì)齊。

系統(tǒng)級(jí)比較

文章將SRA與最新擴(kuò)散模型方法進(jìn)行了系統(tǒng)級(jí)比較，結(jié)果表明：

• 性能提升：SRA在不同模型大小和類型的DiT和SiT上均能顯著提升性能。特別是在大型模型中，SRA的效果更為顯著，這可能是因?yàn)榇笮湍Ｐ湍芴峁└S富的表征引導(dǎo)。
• 與現(xiàn)有方法對(duì)比：SRA在400個(gè)epoch內(nèi)取得了優(yōu)于原始SiT-XL模型的性能，并且在800個(gè)epoch時(shí)進(jìn)一步提升。與依賴復(fù)雜表征訓(xùn)練框架的方法（如MaskDiT）相比，SRA在性能上有顯著優(yōu)勢(shì)，并且與依賴強(qiáng)大外部表征先驗(yàn)的方法（如REPA）相當(dāng)。

消融研究

由于SRA是隱式地引入表征監(jiān)督，文章的消融研究旨在驗(yàn)證表征在SRA中的重要性。實(shí)驗(yàn)結(jié)果表明：

• 表征能力提升：SRA顯著提升了diffusion transformer的潛在表征質(zhì)量，這通過(guò)linear probing結(jié)果在不同層和時(shí)間步上的表現(xiàn)得以驗(yàn)證。
• 生成質(zhì)量與表征引導(dǎo)的緊密關(guān)聯(lián)：隨著用于對(duì)齊的教師網(wǎng)絡(luò)層的變化，linear probing準(zhǔn)確率與FID分?jǐn)?shù)之間顯示出強(qiáng)烈的關(guān)聯(lián)性，證明了SRA中生成能力與自我表征引導(dǎo)機(jī)制的緊密聯(lián)系。

結(jié)論

研究的核心觀點(diǎn)是擴(kuò)散Transformer自身具備提供表征引導(dǎo)的能力，無(wú)需引入額外的表征組件。通過(guò)提出SRA方法，證明了在僅進(jìn)行生成訓(xùn)練的過(guò)程中，通過(guò)自我表征對(duì)齊可以有效提升擴(kuò)散Transformer的表征學(xué)習(xí)能力和生成性能。實(shí)驗(yàn)結(jié)果表明，SRA在多個(gè)模型和數(shù)據(jù)集上均能帶來(lái)顯著的性能提升，并且在與依賴復(fù)雜表征訓(xùn)練框架或強(qiáng)大外部表征先驗(yàn)的方法對(duì)比中展現(xiàn)出優(yōu)越性。

未來(lái)研究展望

盡管SRA取得了顯著的成果，但仍有一些開放問(wèn)題值得進(jìn)一步研究：

• 模型規(guī)模和數(shù)據(jù)擴(kuò)展：研究SRA在更大模型規(guī)模和更多數(shù)據(jù)集上的可擴(kuò)展性，特別是在文本到圖像等更復(fù)雜的生成任務(wù)中的應(yīng)用潛力。
• 理論洞察：探索表征學(xué)習(xí)對(duì)生成任務(wù)有益的理論依據(jù)，以期為diffusion model的訓(xùn)練提供更深入的理解和指導(dǎo)。
• 方法的普適性：驗(yàn)證SRA在其他生成任務(wù)中的應(yīng)用效果，探索其作為通用表征增強(qiáng)方法的潛力。

本文轉(zhuǎn)自AI生成未來(lái) ，作者：AI生成未來(lái)

原文鏈接:??https://mp.weixin.qq.com/s/nxZEU9t0gaqsxyOVOzXCgw??