來源：媒礦工廠本文共5000字，建議閱讀10分鐘
本文着重於潛在擴散系統的優化困境。
題目： Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models
作者： Jingfeng Yao, Bin Yang, Xinggang Wang
論文鏈接： http://arxiv.org/abs/2501.01423
來源：CVPR 2025
內容整理： 周楚駸具有 Transformer 架構的潛在擴散模型擅長生成高保真圖像。然而，最近的研究揭示了這種兩階段設計中的優化困境：雖然在視覺分詞器中增加每個標記的特徵維度可以提高重建質量，但它需要更大的擴散模型和更多的訓練迭代才能實現可比的生成性能。因此，現有系統經常滿足於次優解決方案，要麼由於分詞器中的信息丟失而產生視覺偽影，要麼由於昂貴的計算成本而無法完全收斂。本文認為，這種困境源於學習無約束高維潛在空間的固有困難。為了解決這個問題，本文建議在訓練視覺分詞器時，將潛在空間與預先訓練的視覺基礎模型對齊。本文提出的 VA-VAE（視覺基礎模型對齊變分自動編碼器）顯著擴展了潛在擴散模型的重建生成前沿，使擴散變壓器 （DiT） 在高維潛在空間中的收斂速度更快。為了充分利用 VA-VAE 的潛力，本文通過改進的訓練策略和架構設計構建了增強的 DiT 基線，稱為 LightningDiT。該集成系統在 的ImageNet 上實現了SOTA 性能，FID 得分為 1.35，同時在短短 64 個時期內就達到了 2.11 的 FID 得分，展示了卓越的訓練效率——與原始 DiT 相比，收斂速度提高了21倍以上。
目錄
- 簡介
- 使 VAE 與 Vision Foundation 模型保持一致
- Marginal Cosine Similarity Loss
- Marginal Distance Matrix Similarity Loss
- 自適應權重
- Improved Diffusion Transformer
- 實驗
- 實現細節
- 基礎模型改善收斂性
- 基礎模型提高可擴展性
- 消融和討論
- 生成友好的VA-VAE
- 視覺基礎模型
- 總結
簡介
圖1 潛在擴散模型中的優化困境
潛在擴散模型利用連續值變分自動編碼器 （VAE）或視覺分詞器來壓縮視覺信號，從而減少高分辨率圖像生成的計算需求。這些視覺分詞器的性能，特別是它們的壓縮和重建能力，在決定整體系統的有效性方面起着至關重要的作用。增強重建能力的直接方法是增加視覺標記的特徵維度，從而有效地擴展 Latent 表示的信息容量。最近，一些有影響力的文本到圖像作品與 Stable Diffusion 中廣泛採用的 VAE 相比，探索了更高維的分詞器，因為這些分詞器提供了改進的細節重建，實現了更精細的生成質量。
然而，隨着研究的深入，在潛在擴散模型中，重建和生成性能之間出現了一個優化困境。具體來說，雖然增加標記特徵維度可以提高標記器的重建精度，但它會顯著降低生成性能（見上圖）。目前，有兩種常見的策略可以解決這個問題：第一種涉及擴大模型參數，如 Stable Diffusion 3，它表明更高維的分詞器可以通過明顯更大的模型容量實現更強的生成性能——然而，這種方法需要更多的訓練計算，這使得它對於大多數實際應用來說非常昂貴。第二種策略是故意限制分詞器的重建能力，例如 Sana、W.A.L.T，以更快地收斂擴散模型訓練。然而，這種受損的重建質量本身就限制了生成性能的上限，導致生成結果中的視覺細節不完美。這兩種方法都涉及固有的權衡，並且無法有效控制潛在的優化困境。
本文提出了一種簡單而有效的方法來解決這種優化困境。從自回歸 （AR） 生成中汲取靈感，其中增加離散值 VAE 的碼簿大小會導致碼簿利用率低。通過可視化不同特徵維度的潛在空間分佈（見上圖），本文觀察到高維分詞器以不那麼分散的方式學習潛在表示，分佈可視化中更集中的高強度區域證明了這一點。該分析表明，優化困境源於從頭開始學習無約束高維潛在空間的固有困難。為了解決這個問題，本文為潛在擴散模型中的連續VAEs開發了一種視覺基礎模型指導的優化策略。本文的主要發現表明，由視覺基礎模型指導的學習潛在表徵顯著提高了高維分詞器的生成性能，同時保留了其原始的重建能力（如下圖所示）。
圖2 潛在擴散模型的重建生成能力
本文的主要技術貢獻是 Vision Foundation 模型對齊損失 （VF Loss），這是一個即插即用的模塊，可在分詞器訓練期間將潛在表示與預先訓練的視覺基礎模型保持一致。而事實證明，使用預先訓練的 Vision Foundation 模型天真地初始化 VAE 編碼器是無效的——可能是因為潛在表示會迅速偏離其初始狀態以優化重建——本文發現精心設計的關節重建和對齊損失至關重要。本文的對齊損失是專門為規範高維潛在空間而設計的，而不會過度限制它們的容量。首先，本文強制執行元素級和成對相似性，以確保特徵空間中全局和局部結構的全面正則化。其次，本文在相似性成本中引入了一個邊際，以提供受控的比對靈活性，從而防止過度正則化。此外，本文還研究了不同視力基礎模型的影響。
為了評估生成性能，本文將擬議的 Vision 基礎模型對齊 VAE （VA-VAE） 與擴散 Transformer （DiT）耦合起來，以創建潛在擴散模型。為了充分利用 VA-VAE 的潛力，本文通過先進的擴散訓練策略和 Transformer 架構改進構建了一個增強的 DiT 框架，本文將其命名為 LightningDiT。本恩的貢獻實現了以下重要的里程碑：
- 所提出的 VF Loss 有效解決了潛在擴散模型中的優化困境，使用高維分詞器使 DiT 訓練速度提高了 2.5× 以上;
- 集成系統僅用 64 個訓練 epoch 就達到了 2.11 的 FID，與原始 DiT 相比，收斂速度提高了 21× 以上;
- 該集成系統在 ImageNet-256 圖像生成方面實現了 1.35 的SOTA FID 分數。
使 VAE 與 Vision Foundation 模型保持一致
本節介紹了 VA-VAE，這是一種通過視覺基礎模型對齊訓練的視覺分詞器。關鍵方法包括通過利用基礎模型的特徵空間來限制分詞器的潛在空間，從而增強其對生成任務的適用性。
圖3 提出的VA-VAE
如上圖所示，本文的架構和訓練過程主要遵循 LDM，採用具有連續潛在空間的 VQGAN 模型架構，受 KL 損失的約束。本文的主要貢獻在於 Vision Foundation 模型對齊損失、VF 損失的設計，它有效地在不改變模型架構或訓練管道的情況下，優化了潛在空間解決上文提到的優化困境。
VF 損失由兩個部分組成：marginal cosine similarity loss 和 marginal distance matrix similarity loss。這些組件經過精心設計，是一個簡單明了的即插即用模塊，與 VAE 架構分離。
Marginal Cosine Similarity Loss
在訓練過程中，給定的圖像 均由視覺令牌的編碼器和冷凍視覺基礎模型處理，從而導致圖像潛在的 和基礎視覺表示。,投影 以使用線性變換匹配 的維度，其中，來產生。
損失函數 最小化相應特徵 和 之間的相似性差距。對於每對，計算餘弦相似性，並減去邊緣 。 ReLU可確保只有低於 的成對有助於損失，從而集中在不太相似的對上。最終損失是在 的特徵網格中的所有位置上平均的。
Marginal Distance Matrix Similarity Loss
補充 （強制點對點絕對比對），本文還旨在使特徵中的相對分佈距離矩陣儘可能相似。為此提出了邊緣距離矩陣相似性損失。
距離矩陣相似性損耗使特徵矩陣 和 的內部分佈對齊。在這裡， 表示每個扁平特徵圖中元素的總數。對於每對，計算特徵矩陣 和 中相應向量之間餘弦相似性差的絕對值，從而促進了其相對結構的更緊密比對。同樣，減去邊界 以放鬆約束。 ReLU功能可確保只有超過 的差異對損失的差異。
自適應權重
原始的重建損失和KL損失都是總和損失，這使VF損失完全不同量級，這使調整穩定訓練的重量變得具有挑戰性。受GAN損失的啟發，採用了一種自適應加權機制。在反向傳播之前，在編碼器的最後一個卷積層上計算和的梯度。自適應加權設置為這兩個梯度的比率，以確保和對模型優化具有相似的影響。這種比對大大減少了VF損耗的調整範圍。
然後，本文將通過自適應加權獲得VF損失。自適應加權的目的是快速對不同的VAE訓練管道進行損失量表。在此基礎上，本文仍然可以使用手動調整的超參數來進一步提高性能。
本文將評估VF損失在即將進行的實驗中重建和產生的潛在擴散帕累托前沿中的重要作用。
Improved Diffusion Transformer
本節介紹了本文的LightningDiT。DiT 作為文本-圖像和文本-視頻任務的基礎模型取得了巨大成功。但是，其收斂速度顯著的慢，導致了高實驗迭代成本。以前有影響力的工作Dinov2，Convnext 和Eva 展示了如何融合高級設計策略可以振興經典方法。在本文的工作中，旨在擴大DiT體系結構的潛力，並探索DiT可以走多遠的邊界。雖然本文沒有聲稱任何個人優化細節是最初貢獻，但本文認為，開源、快速收斂的DiT訓練管道將極大地支持社區對DiT的持續研究。
表1 DiT的表現
本文利用以 f8d4 為visual tokenizer的SD-VAE ，並使用DiT-XL/2作為實驗模型。在表1中顯示了優化例程。每個模型均已訓練為80個epoch，並用dopri5 integrator進行採樣，該模型的 NFE 比原始 DiT 少用於快速推斷。為了確保進行公平的比較，不使用諸如CFG間隔和時間段偏移之類的樣本優化方法。採用三類優化策略。在計算級別，實施torch.compile和bfloat16加速訓練。此外，將批量的大小增加，並將 AdamW 的 降低到0.95，從先前的 Auraflow 中汲取了靈感。為了進行擴散優化，結合了Rectified Flow，logit正態分佈（LogNorm）採樣和速度方向損失。在模型體系結構級別，我們應用常用的 Transformer 優化，包括 RMSNorm，SWiGLU和RoPE。在訓練期間，觀察到某些加速策略不是正交的。例如，單獨使用時梯度剪輯是有效的，但在對數點採樣和速度方向損失後組合時傾向於降低性能。
本文優化的模型LightningDiT在Imagenet類條件生成上達到了7.13（CFG = 1）的FID，僅用80個epoch，僅是原始DiT和SiT所需的1400個epoch訓練量的6％。以前的偉大工作MDT 或REPA 在蒙版圖像建模（MIM）和表示對齊的幫助下達到了類似的收斂性能。我們的結果表明，即使沒有任何複雜的培訓管道，簡單的DiT仍然可以取得非常具競爭力的表現。這種優化的體系結構在接下來的快速實驗驗證方面有很大幫助。
實驗
本節主要目標是通過利用提出的VF損失來實現潛在擴散系統內重建和生成的重建和生成前沿。引入的LightningDit中，證明了VF損失如何有效地解決優化的障礙，從而從收斂性，可伸縮性和整體系統性能的角度來看有效地解決了優化的障礙。
實現細節
詳細介紹了潛在擴散系統。對於visual tokenizer，採用主要遵循LDM的體系結構和培訓策略。具體而言，利用 VQGAN 網絡結構，省略量化並應用KL損失來調節連續的潛在空間。為了實現多節點訓練，在MAR的設置上分別將學習率和全局批量規模擴展到1e-4和256。訓練三個不同的 f16 tokenizer：一種沒有VF損失，一種使用VF損失（MAE），另一種使用VF損失（Dinov2）。這裡f表示下採樣率，D表示潛在尺寸。根據經驗，。對於生成模型，採用Lightningdit，進一步完善了概述的設計技術。從tokenizer中提取所有潛在特徵，並以80或160個epoch的分辨率為256上的ImageNet上的LightningDiT的各種版本的LightningDiT。將DiT的patch size設置為1，以確保整個系統的下採樣率為16。所有壓縮步驟均由VAE處理。除非另有說明，否則我們模型的其他架構參數與DiT的架構參數一致。
基礎模型改善收斂性
表2 VF loss提升性能
上表列出了對八種不同的象徵器的重建和產生的評估，所有生成模型均經過160個 epoch（LightningDit-B）或80個 epoch（LightningDit-L＆LightningDit-XL）的訓練。提供以下發現：
結果突出了潛在擴散系統中的優化困境。表中以藍色突出顯示的結果說明了重建性能（rFID）和相應的生成性能（FID）。可以觀察到，隨着tokenizer牌尺寸的增加，其rFID減小，而相應的生成FID增加。
VF損失可以有效地增強高維tokenizer的生成性能。在表的F16D32和F16D64部分中，VF損失（Dinov2）和VF損失（MAE）都顯着改善了不同尺度上DiT模型的生成性能。這使得以更高的重建性能和更高的生成性能（即引言中提到的重建生成前沿）實現系統。但是，值得注意的是，對於低維數字的VF損失是不必要的，例如通常使用的 f16d16。這與圖1中的潛在分佈觀察一致。我們認為這是因為較低維空間可以學習更多合理的分佈，而無需其他額外的監督信號。
圖4 VF Loss加速收斂
此外，上圖中提出了FID的收斂圖。在f16d32和d16f64上，使用VF損失的使用分別按2.54和2.76的倍數加速。這些還表明，VF損失顯著提高了高維tokenizer的生成性能和收斂速度。
基礎模型提高可擴展性
增加模型參數計數是提高高維tokenizer生成性能的一種方法。使用尺寸為0.1b至1.6B的LightningDiT模型來評估3種不同的tokenizer的生成性能。為了促進縮放中冪律的觀察，使用軸的對數尺度。注意到，隨着參數數量的增加，藍色和綠色線之間存在輕微的收斂趨勢，但仍然存在很大的差距。這意味着，高維f16d32 tokenizer 對產生的負面影響即使在1.6B處也沒有完全緩解，這是一個已經被認為是在Imagenet上被認為很重要的參數大小。發現VF損失有效地彌合了這一差距。低於0.6B，橙色和藍線的性能相似。但是，隨着模型尺度超過1B，f16d32 VF DINOV2逐漸與f16d16差距變大，表明可擴展性更強。
消融和討論
在本節中，對VF損失的設計進行消融實驗，以評估各種基礎模型和損失公式的影響。然後們對VF損失的潛在機制進行更深入的分析，從而提供可能有所幫助的其他見解。
生成友好的VA-VAE
與patch size為2的SD-VAE相比，patch size為1的VA-VAE具有出色的生成性能。用VA-VAE代替了SD-VAE [34]，導致FID-50K從7.13降低到7.13至4.29。這種改進可以歸因於兩個主要原因。首先，觀察到使用f16訓練的DiT與使用f8的DIT相比，使用f16且patch size為1的DiT訓練比 f8 patch size為2的DIT更容易收斂。
視覺基礎模型
表3 基礎模型消融
我們使用三種類型的基礎模型訓練VA-VAE：具有掩蓋圖像模型的自我監督模型，圖像文本對比度學習模型CLIP和SAM。為了加速收斂，分別將學習率和全球批量大小調整為1e-4和256。與以前的設置相反，每個tokenizer在ImageNet 256×256上進行了50個epoch的訓練。對於每個tokenizer，在相應的潛在空間中訓練LightningDit-B，以160個epoch的速度訓練。上表總結了我們的發現，表明所有這些視覺基礎模型都增強了擴散模型的生成性能。其中，自我監管的預訓練的模型DinoV2取得了卓越的生成結果。
總結
本文着重於潛在擴散系統的優化困境。為了解決這個問題，提出了與視覺基礎模型保持一致的VA-VAE，以及一項優化的DIT，其中包括了高級設計策略。在VA-VAEe中，VF損失結合邊緣餘弦的相似性和距離矩陣損失 - 與視覺模型保持一致，從而導致更均勻的特徵分佈和高達2.8倍更快的收斂性。通過LightningDiT，整合了先進的訓練技術和架構改進，以實現更快的DIiT收斂。將VA-VAE（rFID = 0.28）的高重建能力與LightningDit的快速收斂結合在一起，本文的方法在Imagenet 256上實現了1.35的最新FID。此外，本文的方法可實現2.11 FID，僅使用64個epoch，證明了21.8倍提速。
