CVPR2025最佳論文候選！收斂速度提升21倍！VA-VAE：重建vs生成，解決LDM的最佳化困境

2025-05-24 17:02 CVer

點選下方卡片，關注“CVer”公眾號

AI/CV重磅乾貨，第一時間送達
點選進入—>【頂會/頂刊】投稿交流群

新增微訊號：CVer2233，小助手會拉你進群！
掃描下方二維碼，加入CVer學術星球！可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料，及最前沿應用！發論文/搞科研/漲薪，強烈推薦！

轉載自：極市平臺 ｜ 作者：科技猛獸

導讀

在ImageNet 256×256 生成上實現了最佳 (SOTA) 效能，FID得分為1.35，同時在短短64個epoch內就達到了2.11的FID得分，展現了卓越的訓練效率——與原始DiT相比，收斂速度提高了21倍以上。

太長不看版

在訓練 VAE 的時候使用 REPA 方法。
本文是一篇研究 VAE 的文章。VAE 現在廣泛用在了 Latent diffusion model 架構裡面，用來去生成高保真影像。但是，最近一些研究發現這種兩階段設計存在最佳化困境：增加 tokenizer 的 feature dimension 可以提高重建 (Reconstruction) 質量，但會損害生成 (Generation) 能力，這就需要更大的擴散模型和更多的訓練才能達到相當的生成能力。
因為這個最佳化困境的問題，現在的系統通常要麼由於 tokenizer 內的資訊丟失而產生視覺偽影，要麼因為計算成本的問題而無法完全收斂。本文認為這種困境源於：VAE 在訓練的時候，如果把 feature dimension 開得比較高，那麼高維度 latent space 在最佳化的時候缺乏約束，因而學好這個特徵就比較困難。
本文提出的解決方案 VA-VAE 的做法是：在訓練 Tokenizer 的時候把特徵對齊 Vision foundation model。其實就是相當於把 REPA 用在了 VAE 的訓練上，使 DiT 在高維 latent space 收斂更快。
為了利用 VA-VAE 的全部潛力，作者構建了一個增強的 DiT baseline，改進了訓練策略和架構設計，稱為 LightningDiT。在 ImageNet 256×256 生成上實現了最先進的 (SOTA) 效能，FID 得分為 1.35。同時，僅 64 個 Epoch 就可以達到 2.11 的 FID 分數，展現出顯著的訓練效率提升：與原始 DiT 相比，表現出 21 倍收斂速度。

本文目錄

1 VA-VAE：解決 LDM 重建 vs. 生成的最佳化困境(來自華中科技大學)1 VA-VAE 論文解讀1.1 VA-VAE 研究背景1.2 VA-VAE 方法介紹1.3 邊際餘弦相似度損失函式1.4 邊際距離矩陣相似性損失函式1.5 自適應權重1.6 LightningDiT：改進的 Diffusion Transformer1.7 實驗設定1.8 實驗結果

1 VA-VAE：解決 LDM 重建 vs. 生成的最佳化困境

論文名稱：Reconstruction vs. Generation:Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025)
論文地址：
http://arxiv.org/pdf/2501.01423
程式碼連結：
http://github.com/hustvl/LightningDiT

1.1 VA-VAE 研究背景

Latent diffusion model (LDM) 通常是利用一個連續的 variational autoencoder (VAE) 或者 visual tokenizer 來壓縮視覺訊號，從而減少高解析度影像生成的計算需求。visual tokenizer 的效能，特別是它們的壓縮和重建能力，對於整體系統的有效性起至關重要的作用。
一種直接增強重建能力的方法是增加 visual tokenizer 的 feature dimension，可以有效地擴大 latent 表徵資訊的能力。
隨著研究的先進，latent diffusion model 的重建和生成效能之間出現了最佳化困境：雖然增加 feature dimension 可以提高 tokenizer 的重建精度，但它顯著降低了生成效能，如圖 1 所示。目前，存在兩個常見的策略來解決這個問題：
  1. 增加模型引數，比如 Stable Diffusion 3。這表明高維 tokenizer 可以以更大的模型容量實現更強的生成效能。但是，這種方法需要更多的訓練計算，這使得它在大多數實際應用中非常昂貴。
  2. 故意限制分詞器的重構能力，例如 Sana、W.A.L.T，為了更快地收斂擴散模型訓練。但是，這種受損的重建質量會固有地限制了生成效能的上限，導致生成結果中的視覺細節不完善。
這兩種方法都涉及固有的權衡，並沒有提供對底層最佳化困境的有效控制。

圖1：latent diffusion model 中的最佳化困境。在 LDM 中，增加視覺 tokenizer 的維度可以增強重建能力，但顯著地降低了生成質量。(f 和 d 分別代表下采樣率和維度。在擴散模型訓練期間，所有結果都是在具有固定計算預算的 ImageNet 256×256 資料集上評估)
本文提出一種簡單而有效的方法來解決這種最佳化困境。
作者從 Auto-Regressive (AR) 生成中汲取靈感，即：增加 discrete VAE 的 codebook 大小會導致 codebook 利用率變低。透過視覺化不同 feature dimension 的 latent 空間分佈 (圖 1)，作者觀察到高維 tokenizer 的分佈視覺化中更集中於高強度區域。這個分析表明，最佳化困境源自學習沒有約束的高維 latent 空間是比較困難的。
為了解決這個問題，本文在 LDM 中為連續 VAE 開發了一種 vision model 引導的最佳化策略。結果表明，在保留其原始重建能力的同時，學習視覺 vision model 引導的 latent 表徵顯著提高高維 tokenizer 的生成能力 (圖 2)。

圖2：LDM 的 Reconstruction-generation frontier。VA-VAE 改善了高維 latent 的特徵分佈。透過與視覺基礎模型對齊，VA-VAE 擴充套件了 LDM 重建和生成之間的邊界

1.2 VA-VAE 方法介紹

VA-VAE 是一種透過 vision foundation model 訓練視覺 tokenizer 的方法，具體是透過 REPA 的表徵對齊策略。VA-VAE 利用基礎模型的特徵空間來約束 tokenizer 的 latent 空間，來增強其對生成任務的適用性。
如圖 3 所示，VA-VAE 的架構和訓練過程主要遵循 LDM，使用具有連續 latent 空間的 VQGAN 架構，受 KL Loss 約束。VA-VAE 的主要貢獻是設計了 Vision Foundation model 對齊的 Loss：VF Loss，在不改變模型架構或 training pipeline 的情況下大大優化了 latent 空間，高效地解決了最佳化困境。

圖3：Vision foundation model Aligned VAE (VA-VAE) 方法。vision foundation model 指導高維視覺 tokenizer 訓練，有效地減輕最佳化困境並提高生成效能
VF 損失由兩部分組成：邊際餘弦相似度損失函式 (marginal cosine similarity loss) 和邊際距離矩陣相似性損失函式 (marginal distance matrix similarity loss)。這些元件是一個可以即插即用的模組，與 VAE 架構解耦。

1.3 邊際餘弦相似度損失函式

在訓練期間，給定的影像  由視覺 tokenizer 的 Encoder 和凍結的 vision foundation model 處理。Encoder 輸出影像 latent  ，視覺基礎模型輸出視覺表徵  。如式 1 所示，使用線性變換投影  來匹配  的維度，其中  ，得到  ：

如式 2 所示，損失函式  力求各空間位置  的特徵矩陣  和  對應特徵  和  之間的相似度差距最小。對於每一對，計算餘弦相似度  並減去一個 margin  函式確保只有相似度低於  的才會對損失函式起作用。這裡的意思是希望損失函式聚焦在相似度低於  的項。最終損失在  的所有位置進行平均。

1.4 邊際距離矩陣相似性損失函式

損失函式  聚焦的是 Encoder 的輸出  和視覺表徵  點對點的相似度。作者還期待特徵內的相對分佈距離矩陣儘可能相似。邊際距離矩陣相似性損失函式就是為了實現這一目的。
如式 3 所示，邊際距離矩陣相似性損失函式對齊特徵矩陣  和  的內部分佈。式中，  表示特徵圖中的元素總數。對於每個對  ，計算特徵矩陣  和  中對應向量之間的餘弦相似度差的絕對值，從而促進其相對結構的更緊密對齊。類似地，減去 margin  來放鬆約束。ReLU 函式確保只有差異超過 margin  的項才有助於損失。

1.5 自適應權重

在圖 3 中，Reconstruction Loss 和 KL Loss 都是加和性質的 Loss，這樣一來，VF Loss 就變成了不同量級的值，使得調整權重以進行穩定訓練具有挑戰性。作者採用了自適應加權機制。在反向傳播之前，計算  和  對 Encoder 最後一個卷積層的梯度，如式 4 所示。自適應加權設定為這兩個梯度的比例，確保  和  對模型最佳化的影響相似。這種對齊顯著減少了 VF 損失的調整範圍。

具有自適應加權的 VF Loss 如式 5 所示。自適應加權的目的是快速對齊不同 Loss 的尺度。在此基礎上，仍然可以使用手動調整的超引數來進一步提高效能。

1.6 LightningDiT：改進的 Diffusion Transformer

在這一節中，作者探索了 DiT 架構和訓練策略的最佳化，希望看看 DiT 能力的邊界。作者使用 SD-VAE (f8d4)，DiT-XL/2 模型作為驗證。實驗結果如圖 4 所示。
對於 Diffusion 最佳化，作者使用 Rectified Flow、logit normal distribution (lognorm) 取樣和 velocity direction loss。在模型架構級別，應用了常見的 Transformer 最佳化，包括 RMSNorm、SwiGLU 和 RoPE。在訓練期間，觀察到一些加速策略不是正交的。例如，梯度裁剪僅在單獨使用時有效，但在使用 lognorm sampling 和 velocity direction loss 後會損失效能。

圖4：LightningDiT 的效能。使用 SD-VAE，LightningDiT 在 ImageNet class-conditional 生成上實現了 FID-50k=7.13，與原始 DiT 相比，訓練樣本減少了 94%
最佳化後的模型 LightningDiT 在 ImageNet class-conditional 生成上達到了 7.13 (cfg=1)，SD-VAE 在 80 Epoch 左右，僅為原始 DiT 和 SiT 模型 6% 的訓練時間 (1400 Epoch)。這些結果表明，即使沒有任何複雜的 training pipeline，DiT 仍然可以實現非常有競爭力的效能。

1.7 實驗設定

對於 visual tokenizer，本文主要遵循 LDM 的架構和策略。具體來說，作者利用 VQGAN 網路結構，省略量化並應用 KL Loss 來調節連續的 latent 空間。作者訓練了 3 個不同的  tokenizer：一個沒有 VF Loss，一個使用 VF Loss（MAE），另一個使用 VF Loss（DINOv2）。這裡  表示下采樣率， 表示 latent 維度。根據經驗，設定  。生成模型直接採用 LightningDiT。
從 tokenizer 中提取所有 latent 特徵，並在 ImageNet 上訓練各種版本的 LightningDiT，解析度為 256。將 DiT 的 patch size 設定為 1，確保整個系統的下采樣率為 16。這種方法與 DC-AE 中提出的策略一致，即所有壓縮步驟都由 VAE 處理。

1.8 實驗結果

基礎模型改善收斂

圖 5 顯示了對8種不同 tokenizer 的重建和生成的評估，所有生成模型在 ImageNet 上訓練了 160 Epoch (LightningDiT-B) 或 80 Epoch (LightningDiT-L 和 LightningDiT-XL) 。有以下發現：
結果突出了 LDM 的最佳化困境。藍色突出顯示的結果說明了重建效能 (rFID) 和相應的生成效能 (FID)。可以觀察到，隨著 tokenizer 維度的增加，rFID 減小，而相應的生成 FID 增加。

圖5：VF Loss 提高了生成效能。隨著維度的增加，觀察：1) 高維提高了重建，但降低了生成能力，突出了 LDM 的最佳化困境。2) VF Loss 在高維 tokenizer 中顯著提高了生成效能，而對重建的影響最小
VF Loss 可以有效提高高維 tokenizer 的生成效能。在 f16d32 和 f16d64 部分中，VF Loss (DINOv2) 和 VF Loss (MAE) 都顯著提高了 DiT 在不同尺度上的生成效能。這使得實現具有更高重建效能和更高生成效能的系統成為可能。然而，值得注意的是，低維 tokenizer (比如 f16d16) 不需要 VF Loss。作者認為這是因為低維空間可以在不需要額外的監督訊號的情況下學習更合理的分佈。
此外，作者在圖 6 (a) 和 (b) 中展示了 FID 相對於訓練時間的收斂圖。在 f16d32 和 f16d64 上，VF Loss 分別加速了 2.54 和 2.76 倍的收斂，表明 VF Loss 可以顯著提高高維 tokenizer 的生成效能和收斂速度。

圖6：VF Loss 可以改善收斂。使用不同的 tokenizer 在 ImageNet 上訓練 LightningDiT-B 160 Epoch，解析度為 256。VF 損失顯著加速了收斂，最大加速高達 2.7 倍

基礎模型改善縮放效能

增加模型引數量是提高高維 tokenizer 生成效能的方法之一。
作者用大小從 0.1B 到 1.6B 的 LightningDiT 模型來評估 3 個不同的 tokenizer 的生成效能。採用對數標度。
可以觀察到到隨著引數量的增加，藍線和綠線之間仍然存在顯著差距。這意味著即使在 1.6B 的引數量級，高維 f16d32 tokenizer 對生成的負面影響也沒有完全緩解。而 1.6B 的引數量對於 ImageNet 而言已經很大。作者發現 VF Loss 可有效彌合了這一差距。當模型規模超過 1B 時，f16d32 + VF DINOv2 的效能逐漸與 f16d16 拉開差距，展示出可擴充套件性。

圖7：VF Loss 提高了縮放性質。VF Loss 減少了高維 tokenizer 生成模型中對大引數的需求，實現了更好的可擴充套件性

收斂速度

作者發現 VF Loss (DINOv2) 在生成效能方面帶來了最顯著的改進。因此，作者擴充套件了 tokenizer 的訓練時間，並採用漸進式訓練策略來訓練 LDM VF Loss (DINOv2) 125 Epoch。訓練 LightningDiT-XL 800 Epoch。具體來說，在 480 Epoch 時停用 lognorm 引數，以使接近收斂的網路能夠在所有噪聲區間中更有效地學習。在取樣過程中，使用 250-step Euler integrator。為了提高取樣效能，採用類似於 FLUX 的 cfg interval 和 timestep shift。

圖8：ImageNet 256×256 效能對比結果
實驗結果如圖 9 所示，本文模型訓練了 800 Epoch 之後，模型實現了最先進的效能，FID 為 1.35。此外，模型在不使用 cfg 生成方面表現出卓越的效能，實現了 2.17 的 FID，超過了許多使用 cfg 的方法的結果。
本文也展示出了很快的收斂速度：在 64 Epoch 時，實現了 2.11 的 FID，與原始 DiT 相比，實現了超過 21 倍的加速。

何愷明在MIT授課的課件PPT下載
在CVer公眾號後臺回覆：何愷明，即可下載本課程的所有566頁課件PPT！趕緊學起來！

CVPR 2025 論文和程式碼下載

在CVer公眾號後臺回覆：CVPR2025，即可下載CVPR 2025論文和程式碼開源的論文合集

ECCV 2024 論文和程式碼下載

在CVer公眾號後臺回覆：ECCV2024，即可下載ECCV 2024論文和程式碼開源的論文合集
CV垂直方向和論文投稿交流群成立
掃描下方二維碼，或者新增微訊號：CVer2233，即可新增CVer小助手微信，便可申請加入CVer-垂直方向和論文投稿微信交流群。另外其他垂直方向已涵蓋：目標檢測、影像分割、目標跟蹤、人臉檢測&識別、OCR、姿態估計、超解析度、SLAM、醫療影像、Re-ID、GAN、NAS、深度估計、自動駕駛、強化學習、車道線檢測、模型剪枝&壓縮、去噪、去霧、去雨、風格遷移、遙感影像、行為識別、影片理解、影像融合、影像檢索、論文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要備註：研究方向+地點+學校/公司+暱稱（如Mamba、多模態學習或者論文投稿+上海+上交+卡卡），根據格式備註，可更快被透過且邀請進群

▲掃碼或加微訊號: CVer2233，進交流群
CVer計算機視覺（知識星球）人數破萬！如果你想要了解最新最快最好的CV/DL/AI論文、實戰專案、行業前沿、從入門到精通學習教程等資料，一定要掃描下方二維碼，加入CVer知識星球！最強助力你的科研和工作！

▲掃碼加入星球學習
▲點選上方卡片，關注CVer公眾號
整理不易，請點贊和在看


相關文章