文章鏈接: https://arxiv.org/pdf/2504.10483 
項目鏈接：https://end2end-diffusion.github.io/ 
Git鏈接：https://github.com/End2End-Diffusion/REPA-E 
模型鏈接：https://huggingface.co/REPA-E 

亮點直擊

• 端到端聯合優化的突破首次實現VAE與擴散模型的端到端聯合訓練，通過REPA Loss替代傳統擴散損失，解決兩階段訓練目標不一致問題，使隱空間與生成任務高度適配。
• 訓練效率革命性提升REPA-E僅需傳統方法1/45的訓練步數即可收斂，且生成質量顯著超越現有方法（如FID從5.9降至4.07），大幅降低計算成本。
• 雙向性能增益不僅提升擴散模型性能，還通過反向傳播優化VAE的隱空間結構，使其成為通用型模塊，可遷移至其他任務（如替換SD-VAE后下游任務性能提升）。

總結速覽

解決的問題

現有隱空間擴散模型（LDM）采用兩階段訓練（先訓練VAE，再固定VAE訓練擴散模型），導致兩個階段的優化目標不一致，限制了生成性能。直接端到端聯合訓練VAE和擴散模型時，傳統擴散損失（Diffusion Loss）失效，甚至導致性能下降。

提出的方案

提出REPA-E訓練框架，通過表示對齊損失（REPA Loss）實現VAE與擴散模型的端到端聯合優化。REPA Loss通過對齊隱空間表示的結構，協調兩者的訓練目標，替代傳統擴散損失的直接優化。

應用的技術

• 表示對齊損失（REPA Loss）：在擴散模型的去噪過程中，對齊隱空間表示的分布，確保VAE生成的隱空間編碼更適配擴散模型的訓練目標。
• 端到端梯度傳播：通過聯合優化框架，將擴散模型的梯度反向傳播至VAE，動態調整其隱空間結構。
• 自適應隱空間優化：根據擴散模型的訓練需求，自動平衡VAE的重建能力與隱空間的可學習性。

達到的效果

• 訓練加速：相比傳統兩階段訓練（4M步），REPA-E僅需400K步即達到更優性能，訓練速度提升45倍；相比單階段REPA優化（17倍加速）。
• 生成質量SOTA：在ImageNet 256×256上，FID達到1.26（使用分類器引導）和1.83（無引導），刷新當前最佳記錄。
• 隱空間改善：對不同初始結構的VAE（如高頻噪聲的SD-VAE、過平滑的IN-VAE），端到端訓練自動優化其隱空間，提升生成細節

REPA-E：解鎖VAE的聯合訓練

概述。 給定一個變分自編碼器（VAE）和隱空間diffusion transformer （例如SiT），本文希望以端到端的方式聯合優化VAE的隱空間表示和擴散模型的特征，以實現最佳的最終生成性能。首先提出三個關鍵見解：1）樸素的端到端調優——直接反向傳播擴散損失到VAE是無效的。擴散損失鼓勵學習更簡單的隱空間結構（下圖3a），這雖然更容易最小化去噪目標，但會降低最終生成性能。接著分析了最近提出的表示對齊損失，發現：2）更高的表示對齊分數與改進的生成性能相關（圖3b）。這為使用表示對齊分數作為代理來提升最終生成性能提供了另一種途徑。3）使用樸素REPA方法可達到的最大對齊分數受限于VAE隱空間特征的瓶頸。進一步表明，在訓練過程中將REPA損失反向傳播到VAE有助于解決這一限制，顯著提高最終的表示對齊分數（圖3c）。

基于上述見解本文提出了REPA-E；一種用于聯合優化VAE和LDM特征的端到端調優方法。我們的核心思想很簡單：不直接使用擴散損失進行端到端調優，而是使用表示對齊分數作為最終生成性能的代理。這促使本文提出最終方法，即不使用擴散損失，而是使用表示對齊損失進行端到端訓練。通過REPA損失的端到端訓練有助于更好地提高最終的表示對齊分數（圖3b），從而提升最終生成性能。

用REPA推動端到端訓練的動機

單純端到端訓練對擴散損失的影響。

更高的表示對齊分數與更好的生成性能相關。 本文還使用CKNNA分數在不同模型規模和訓練迭代中測量表示對齊。如前圖3b所示，訓練過程中更高的表示對齊分數會帶來更好的生成性能。這表明可以通過使用表示對齊目標（而非擴散損失）進行端到端訓練來提升生成性能。

表示對齊受限于VAE特征。 圖3c顯示，雖然樸素應用REPA損失可以提高表示對齊（CKNNA）分數，但可達到的最大對齊分數仍受限于VAE特征，飽和值約為0.4（最大值為1）。此外，我們發現將表示對齊損失反向傳播到VAE有助于解決這一限制；允許端到端優化VAE特征以最好地支持表示對齊目標。

基于REPA的端到端訓練

REPA-E——一種用于聯合訓練VAE和LDM特征的端到端調優方案。建議使用表示對齊損失而非直接使用擴散損失來進行端到端訓練。通過REPA損失實現的端到端訓練能夠更好地提升最終表示對齊分數（圖3c），從而改善最終生成性能。

VAE隱空間歸一化的批歸一化層

為了實現端到端訓練，我們首先在VAE和隱空間擴散模型之間引入批歸一化層。典型的LDM訓練需要使用預計算的隱空間統計量（例如SD-VAE的std=1/0.1825）對VAE特征進行歸一化。這有助于將VAE隱空間輸出歸一化為零均值和單位方差，從而提高擴散模型的訓練效率。然而，在端到端訓練中，每當VAE模型更新時都需要重新計算統計量——這代價高昂。為解決這個問題，我們提出使用批歸一化層，該層使用指數移動平均（EMA）均值和方差作為數據集級統計量的替代。因此，批歸一化層充當了可微分歸一化算子，無需在每次優化步驟后重新計算數據集級統計量。

端到端表示對齊損失

整體訓練

最終的整體訓練以端到端方式使用以下損失函數進行：

其中θ、φ、ω分別表示潛擴散模型（LDM）、變分自編碼器（VAE）以及可訓練的REPA投影層的參數。

實驗

本文通過廣泛實驗驗證REPA-E的性能及所提組件的影響，主要探究三個關鍵問題：

1. REPA-E能否顯著提升生成性能與訓練速度？（下表2、前圖1、下圖4）
2. REPA-E是否適用于不同訓練設置（模型規模、架構、REPA編碼器等）？（下表3-8）
3. 分析端到端調優（REPA-E）對VAE隱空間結構及下游生成性能的影響。（圖6、表9-10）

實驗設置

實現細節

評估指標

圖像生成評估嚴格遵循ADM：在5萬張生成圖像上報告Fréchet Inception距離（gFID）、結構FID（sFID）、Inception分數（IS）、精確度（Prec.）與召回率（Rec.）。采樣采用SiT和REPA的SDE Euler-Maruyama方法（250步）。VAE評估在ImageNet驗證集5萬張256×256圖像上測量重建FID（rFID）。

訓練性能與速度的影響

首先分析REPA-E對隱空間diffusion transformer 訓練性能與速度的提升。

定量評估

下表2比較了不同隱空間擴散模型（LDM）基線。在ImageNet 256×256生成任務中評估相似參數量（～675M）的模型，結果均未使用分類器無關引導。關鍵發現：

1. 端到端調優加速訓練：相比REPA，gFID從19.40→12.83（20輪）、11.10→7.17（40輪）、7.90→4.07（80輪）持續提升；
2. 端到端訓練提升最終性能：REPA-E在80輪時gFID=4.07，優于FasterDiT（400輪，gFID=7.91）、MaskDiT、DiT和SiT（均訓練1400輪以上）。REPA-E僅需40萬步即超越REPA 400萬步的結果（gFID=5.9。

定性評估

下圖4對比REPA 與REPA-E在5萬、10萬、40萬步的生成效果。REPA-E在訓練早期即生成結構更合理的圖像，且整體質量更優。

REPA-E的泛化性與可擴展性

進一步分析REPA-E在不同訓練設置下的適應性（模型規模、分詞器架構、表示編碼器、對齊深度等）。默認使用SiT-L 為生成模型，SD-VAE為VAE，DINOv2-B為REPA損失的預訓練視覺模型，對齊深度為8。各變體訓練10萬步，結果均未使用分類器無關引導。

模型規模的影響

下表3比較SiT-B、SiT-L和SiT-XL：

• REPA-E在所有配置下均優于REPA基線，gFID從49.5→34.8（SiT-B）、24.1→16.3（SiT-L）、19.4→12.8（SiT-XL）；
• 增益隨模型規模增大：SiT-B提升29.6%，SiT-L提升32.3%，SiT-XL提升34.0%，表明REPA-E對大模型更具擴展性。

表示編碼器的選擇

表4顯示不同感知編碼器（CLIP-L、I-JEPA-H、DINOv2-B/DINOv2-L）下REPA-E均一致提升性能。例如DINOv2-B下gFID從24.1→16.3，DINOv2-L下從23.3→16.0。

VAE的變體。 下表5評估了不同VAE對REPA-E性能的影響。我們報告了使用三種不同VAE的結果：1）SD-VAE，2）VA-VAE，以及3）IN-VAE（一個在ImageNet上訓練的16倍下采樣、32通道的VAE，使用[39]中的官方訓練代碼）。在所有變體中，REPA-E始終在性能上優于REPA基線。REPA-E將gFID從24.1降至16.3（SD-VAE），從22.7降至12.7（IN-VAE），從12.8降至11.1（VA-VAE）。結果表明，REPA-E在VAE的架構、預訓練數據集和訓練設置多樣性的情況下，始終能穩健提升生成質量。

對齊深度的變體。 下表6研究了在擴散模型不同層應用對齊損失的效果。觀察到REPA-E在不同對齊深度選擇下均能持續提升生成質量，相較REPA基線，gFID分別從23.0降至16.4（第6層）、24.1降至16.3（第8層）、23.7降至16.2（第10層）。

組件設計的消融實驗。 本文還進行了消融研究，分析各組件的重要性，結果如表7所示。觀察到每個組件對REPA-E的最終性能都起到了關鍵作用。特別地，觀察到對擴散損失應用stop-grad操作有助于防止隱空間結構的退化。同樣地，批歸一化（batch norm）通過自適應地規范化潛變量統計信息，提升了gFID從18.09至16.3。同樣地，正則化損失對保持微調后VAE的重建性能起到了關鍵作用，從而將gFID從19.07提升至16.3。

端到端從頭訓練。 分析了VAE初始化對端到端訓練的影響。如表8所示，我們發現雖然使用預訓練權重初始化VAE能略微提升性能，但REPA-E也可以在從頭訓練VAE和LDM的情況下使用，并仍然在性能上優于REPA，后者在技術上需要一個VAE訓練階段以及LDM訓練階段。例如，REPA在400萬次迭代后達到FID 5.90，而REPA-E在完全從頭訓練的情況下（同時訓練VAE和LDM）在僅40萬次迭代內就達到了更快更優的生成FID 4.34。

端到端微調對VAE的影響

接下來分析了端到端微調對VAE的影響。首先展示端到端微調能改善隱空間結構（下圖6）。然后展示一旦使用REPA-E進行微調，微調后的VAE可以作為原始VAE的直接替代品，顯著提升生成性能。

端到端訓練提升隱空間結構。 結果如圖6所示。本文使用主成分分析（PCA）將隱空間結構可視化為RGB著色的三個通道。我們考慮三種不同的VAE：1）SD-VAE，2）IN-VAE（一個在ImageNet上訓練的16倍下采樣、32通道的VAE），3）最新VA-VAE。觀察到使用REPA-E進行端到端微調自動改善了原始VAE的隱空間結構。例如，與同時期工作的發現一致，我們觀察到SD-VAE的隱空間存在高噪聲成分。應用端到端訓練后可自動幫助調整隱空間以減少噪聲。相比之下，其他VAE如最新提出的VA-VAE的隱空間則表現為過度平滑。使用REPA-E進行端到端微調可自動學習更具細節的隱空間結構，以更好地支持生成性能。

端到端訓練提升VAE性能。 接下來評估端到端微調對VAE下游生成性能的影響。首先對最近提出的VA-VAE進行端到端微調。然后我們使用該微調后的VAE（命名為E2E-VAE），將其下游生成性能與當前最先進的VAE進行比較，包括SD-VAE和VA-VAE。本文進行了傳統的潛擴散模型訓練（不使用REPA-E），即僅更新生成器網絡，同時保持VAE凍結。表9展示了在不同訓練設置下的VAE下游生成性能比較。端到端微調后的VAE在各種LDM架構和訓練設置下的下游生成任務中始終優于其原始版本。有趣的是，觀察到使用SiT-XL進行微調的VAE即使在使用不同LDM架構（如DiT-XL）時仍能帶來性能提升，進一步展示了本文方法的穩健性。

結論

本文探討了一個基本問題：“我們是否能夠實現基于隱空間擴散 Transformer 的端到端訓練，從而釋放 VAE 的潛力？”具體來說，觀察到，直接將擴散損失反向傳播到 VAE 是無效的，甚至會降低最終的生成性能。盡管擴散損失無效，但可以使用最近提出的表示對齊損失進行端到端訓練。所提出的端到端訓練方案（REPA-E）顯著改善了隱空間結構，并展現出卓越的性能：相較于 REPA 和傳統訓練方案，擴散模型訓練速度分別提升了超過 17× 和 45×。

REPA-E 不僅在不同訓練設置下表現出一致的改進，還改善了多種 VAE 架構下原始的隱空間結構。總體而言，本文的方法達到了新的SOTA水平，在使用和不使用 classifier-free guidance 的情況下，分別取得了 1.26 和 1.83 的生成 FID 分數。希望本工作能夠推動進一步的研究，推動隱空間擴散 Transformer 的端到端訓練發展。

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

原文鏈接:??https://mp.weixin.qq.com/s/AawnxSR6jf3oJ3Xp-960rg??