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

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

亮點直擊

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

背景介紹

在生成式預訓練中引入表征指導

Diffusion transformer在視覺生成領域展現出了有效的可擴展性和可遷移性。最近的研究發現高質量的內部表征能夠加速其生成訓練，并提升生成樣本的質量然而，現有方法在提供表征指導時，要么引入復雜訓練框架，要么依賴大規模預訓練模型，這限制了其應用場景和資源效率。

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

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

靈感和發現

靈感來源

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

發現

文章對普通的SiT和DiT進行了實證分析。我們首先發現其中的潛在特征隨著塊層數的增加和噪聲水平的降低而逐步細化，呈現出大致從粗糙到精細的過程。接下來，我們觀察到僅通過生成式預訓練的diffusion transformer已經能夠學習到有意義的判別性表征。

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

方法（SRA）

SRA的核心思想是利用擴散Transformer自身在不同訓練階段和噪聲條件下的表征差異，通過自我蒸餾的方式進行表征對齊。具體來說，SRA將早期層（高噪聲條件下）的潛在表征與后期層（低噪聲條件下）的潛在表征進行對齊，從而實現自我表征引導。這種方法無需額外的表征組件或外部預訓練模型，僅在diffusion transformer自身的生成訓練過程中引入指導。

實現細節

• 表征對齊：將學生模型（可訓練模型）早期層的潛在表征與教師模型（通過指數移動平均更新權重的模型）后期層的潛在表征進行對齊。學生模型的輸出通過輕量級可訓練MLP頭進行非線性變換后，與教師模型的輸出進行對齊。
• 損失函數：SRA通過最小化教師模型輸出和學生模型輸出變體之間的patch-wise距離來實現自我表征對齊。同時，將該目標函數與原始的擴散模型目標函數結合，進行聯合學習。
• EMA教師網絡：為了避免直接使用同模型輸出作為監督信號導致的訓練不穩定性，SRA采用指數移動平均（EMA）策略構建教師模型。教師模型的權重基于學生模型的權重進行加權移動平均更新，從而提供穩定的表征引導。
• 超參數設置：在實驗中，我們對超參數進行了細致的調整，包括表征對齊的塊層選擇、時間間隔、對齊目標函數以及EMA教師網絡的動量系數等。這些超參數的選擇對SRA的性能有著顯著的影響。

實驗

實驗設置

實驗中，文章采用了ImageNet 1K數據集，并遵循DiT和SiT的訓練設置，包括使用AdamW優化器、恒定學習率1e-4、批量大小256等。評估指標包括FID、sFID、IS、Precision和Recall，使用相同的評價框架和參考圖像以確保與先前方法的公平比較。

元素級分析

文章通過詳細的組件分析，研究了SRA中不同設計選擇對性能的影響。實驗結果表明：

• 塊層選擇：將教師模型的后期層與學生模型的早期層進行對齊能取得最佳性能，因為早期層需要更多的指導以捕捉語義豐富的表征。
• 時間間隔：使用教師模型比學生模型低噪聲輸入的特征能提升性能，且時間間隔值在0.1左右時效果最佳。動態時間間隔表現略優于固定時間間隔。
• 對齊目標函數：smooth-?1、?1和?2三種回歸訓練目標函數均能有效提升性能。
• 教師網絡更新策略：采用固定動量系數0.9999的EMA策略更新教師網絡效果最佳，其他更新策略并不適合SRA。
• 投影頭的作用：使用輕量級投影頭對學生的輸出進行后處理，能顯著提升性能，因為它有助于讓模型在進行非線性變換前捕捉更有效的隱藏表征以便對齊。

系統級比較

文章將SRA與最新擴散模型方法進行了系統級比較，結果表明：

• 性能提升：SRA在不同模型大小和類型的DiT和SiT上均能顯著提升性能。特別是在大型模型中，SRA的效果更為顯著，這可能是因為大型模型能提供更豐富的表征引導。
• 與現有方法對比：SRA在400個epoch內取得了優于原始SiT-XL模型的性能，并且在800個epoch時進一步提升。與依賴復雜表征訓練框架的方法（如MaskDiT）相比，SRA在性能上有顯著優勢，并且與依賴強大外部表征先驗的方法（如REPA）相當。

消融研究

由于SRA是隱式地引入表征監督，文章的消融研究旨在驗證表征在SRA中的重要性。實驗結果表明：

• 表征能力提升：SRA顯著提升了diffusion transformer的潛在表征質量，這通過linear probing結果在不同層和時間步上的表現得以驗證。
• 生成質量與表征引導的緊密關聯：隨著用于對齊的教師網絡層的變化，linear probing準確率與FID分數之間顯示出強烈的關聯性，證明了SRA中生成能力與自我表征引導機制的緊密聯系。

結論

研究的核心觀點是擴散Transformer自身具備提供表征引導的能力，無需引入額外的表征組件。通過提出SRA方法，證明了在僅進行生成訓練的過程中，通過自我表征對齊可以有效提升擴散Transformer的表征學習能力和生成性能。實驗結果表明，SRA在多個模型和數據集上均能帶來顯著的性能提升，并且在與依賴復雜表征訓練框架或強大外部表征先驗的方法對比中展現出優越性。

未來研究展望

盡管SRA取得了顯著的成果，但仍有一些開放問題值得進一步研究：

• 模型規模和數據擴展：研究SRA在更大模型規模和更多數據集上的可擴展性，特別是在文本到圖像等更復雜的生成任務中的應用潛力。
• 理論洞察：探索表征學習對生成任務有益的理論依據，以期為diffusion model的訓練提供更深入的理解和指導。
• 方法的普適性：驗證SRA在其他生成任務中的應用效果，探索其作為通用表征增強方法的潛力。

本文轉自AI生成未來 ，作者：AI生成未來

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