文章鏈接：https://arxiv.org/pdf/2411.11505

亮點直擊

• 高效處理高維視覺數據
• 創新的聯合擴散Transformer
• 強大的任務泛化能力
• 顯著性能與效率提升

總結速覽

解決的問題

當前的大型視覺模型大多直接從自然語言處理架構中改造而來，使用低效的自回歸技術，這種方法破壞了視覺數據中至關重要的空間關系，并且在高維度視覺數據的生成任務中表現有限。

提出的方案

• 空間-時間變分自編碼器 (spatial-temporal variational autoencoder)
  編碼高維視覺數據到連續的潛在空間，降低數據維度。
• 聯合擴散Transformer (joint diffusion transformer)
  通過逐步擴散生成視覺輸出，優化生成建模能力。
• 基于上下文學習的統一多任務訓練 (in-context learning)
  利用輸入-目標對作為任務上下文，引導擴散Transformer在潛在空間中完成特定任務。推理時，通過任務上下文和測試數據實現跨任務的泛化，無需微調。

應用的技術

• 變分自編碼器 (VAE) 技術用于高效編碼視覺數據。
• 擴散模型與Transformer的結合，用于逐步生成高質量視覺結果。
• 上下文學習（in-context learning）機制，支持多任務統一訓練和推理。
• 可擴展性：模型參數從 0.1B 擴展至 3.4B，充分驗證了其可擴展性能。

達到的效果

• 性能提升：支持超過20+視覺任務，多項SOTA。
• 統一框架：首次在生成框架中實現大規模、多任務視覺建模。
• 高效推理：通過上下文學習，在無需微調的情況下實現任務泛化。
• 開源貢獻：計劃開放代碼和模型，為后續研究提供支持。

方法

LaVin-DiT 模塊

ST-VAE

直接在原始像素空間處理視覺數據計算量巨大。為了解決這個問題，本文提出使用空間-時間變分自編碼器（ST-VAE）。ST-VAE 能有效壓縮空間和時間信息，將其從像素空間編碼到緊湊的潛在空間。如上圖 2(b) 所示，ST-VAE 使用因果 3D 卷積和反卷積來壓縮和重建視覺數據。它總體上包含一個編碼器、一個解碼器和一個潛在正則化層。這些組件被結構化為四個對稱階段，交替進行2X下采樣和上采樣。前兩個階段作用于空間和時間維度，而最后階段僅作用于空間維度，實現了4X8X8的壓縮，有效減少了計算量。此外，應用了 Kullback-Leibler (KL) 約束來正則化高斯潛在空間。

為防止未來信息泄漏及其對時間預測的不利影響，在時間卷積空間開始時對所有位置進行填充。此外，為了支持圖像和視頻處理，將輸入視頻的第一幀獨立對待，僅對其進行空間壓縮以保持時間獨立性。后續幀同時在空間和時間維度上進行壓縮。

ST-VAE 的編碼器將輸入壓縮到低維潛在空間，通過解碼過程實現重建。ST-VAE 的訓練分為兩個階段：首先單獨對圖像進行訓練，然后聯合圖像和視頻進行訓練。在每個階段，我們使用均方誤差、感知損失和對抗性損失的組合來優化模型。

J-DiT

Diffusion Transformer（DiT）已成為生成建模的強大方法。聯合Diffusion Transformer（J-DiT）基于 DiT 構建，并引入修改以支持任務條件生成。與原始 DiT 的關鍵區別在于我們考慮了兩種概念上不同的潛在表示。條件潛在表示是干凈的，而目標潛在表示被高斯噪聲擾動，這導致兩者可能具有不同的值范圍。為了處理這種差異并改善任務特定信息與視覺信息之間的對齊，為條件和目標潛在表示分別構建了獨立的分塊嵌入層。每個嵌入層使用2X2的塊大小，能夠針對每種潛在類型定制表示。如前面圖 2 所示，采樣的時間步t以及條件和目標序列被輸入到一系列擴散Transformer層中。基于 MM-DiT 的架構，在 AdaRN 層中引入了針對條件和目標的自適應 RMS 正則化（AdaRN），分別調節每種表示空間。這通過為條件和目標分別生成的時間步嵌入實現。

全序列聯合注意力
全序列聯合注意力是Transformer層的核心，能夠將條件和噪聲目標序列一起處理，以增強任務特定的對齊。如前面圖 2(c) 所示，條件和目標序列被線性投影后連接在一起，并通過雙向注意力模塊處理，從而使兩者在各自的空間中操作時也能互相考慮。

為提高速度和內存效率，用分組查詢注意力（grouped-query attention）替代多頭注意力。這種方法將查詢頭分組，以共享一組鍵值頭，從而減少參數量，同時保持表達能力，與標準多頭注意力性能接近。此外，為了在處理更大的模型和更長的序列時穩定訓練，我們在查詢-鍵點積之前添加了 QK-Norm 來控制注意力熵的增長。按照 [53] 的方法，還在每個注意力和前饋層后應用了夾心正則化（sandwich normalization），以在殘差連接中保持激活幅度。

3D旋轉位置編碼與 [4] 不同，將視覺數據建模為一維序列是不理想的，因為一維位置嵌入在捕捉精確的時空位置方面具有局限性。相反，通過將多組圖像-標注對或視頻片段視為一個連續序列，可以使用三維旋轉位置編碼（3D RoPE）簡潔地表示時空關系。這樣，視頻中的每個位置都可以用三維坐標表示。引入 3D RoPE 后，為各種視覺任務提供了一種統一且精確的時空位置編碼表示。

LaVin-DiT 推理

完成 LaVin-DiT 的訓練后，該模型具有多功能性，可以應用于多種下游任務。具體而言，對于任何選定任務，在給定查詢（例如，一張圖像或一段視頻）時，我們隨機采樣一組定義任務的輸入-目標對。這些對與視覺輸入以及高斯噪聲組件一起被輸入到聯合 Diffusion Transformer（J-DiT）。在 J-DiT 中，這些元素被處理以生成一個潛在表示。最終，該潛在表示被傳遞到 ST-VAE 解碼器中，將其轉換為原始像素空間，從而生成所需的預測結果。

實驗

設置

訓練數據
為了統一多個計算機視覺任務，構建了一個大規模的多任務數據集，涵蓋了室內和室外環境，跨越了現實世界和合成領域。該數據集包含約320萬張獨特圖像和60萬段獨特視頻，涵蓋了超過20個任務：

• 基于圖像的任務：物體檢測、實例分割、全景分割、姿態估計、邊緣提取、深度估計、表面法線估計、圖像修復、圖像著色、圖像恢復任務（如去雨、去玻璃模糊、去運動模糊）、深度到圖像生成、法線到圖像生成。
• 基于視頻的任務：幀預測、視頻深度估計、視頻表面法線估計、視頻光流估計、視頻實例分割、深度到視頻生成、法線到視頻生成。

為了克服深度和表面法線估計的大規模標注限制，利用Depth-anything V2  和Stable-Normal（turbo）在ImageNet-1K 上生成了偽深度圖和法線圖。

實現細節
分兩個階段進行訓練，逐步增加圖像分辨率。在第一階段，以256×256分辨率訓練100,000步，利用DeepSpeed ZeRO-2優化和梯度檢查點技術來管理內存和計算效率。使用640的全局批量大小，并采用AdamW優化器，學習率為0.0001，betas設置為0.9和0.95，權重衰減為0.01。該設置提供了穩定的訓練，無需熱身或額外的正則化技術。在第二階段，將分辨率上調到512×512，并繼續訓練20,000步，同時將學習率調整為0.00005，其他超參數保持不變。該兩階段策略能夠高效擴展，確保不同分辨率下的最佳性能。默認情況下，在推理過程中使用20個時間步（N=20）。所有實驗均在64個NVIDIA A100-80G GPU上進行。

評估協議
在一系列涵蓋圖像和視頻領域的計算機視覺任務上評估我們的模型。按照已建立的協議，報告每個任務的標準評估指標。

主要結果

定量分析
為了評估本文提出方法的有效性，在一系列計算機視覺任務上進行了廣泛的實驗，并報告了3.4B模型的結果，如下表1和表2所示。本文的方法在多個任務上始終優于現有基準方法，包括一些具有挑戰性的案例，如未見過的前景分割和單物體檢測，展示了模型在各種場景中的卓越泛化能力和適應性。除非另有說明，否則告的是LaVin-DiT (3.4B)的性能。

如上表1所示，報告了在不同分割上的前景分割和單物體檢測任務的性能。LaVin-DiT在所有分割上都取得了顯著的提升。在前景分割任務中，在四個分割上分別達到了67.87%、75.80%、66.98%和66.90%的mIoU，始終大幅超過了LVM和MAE-VQGAN等先前方法。

此外，對于單物體檢測，本文的模型在所有分割上都表現出色，特別是在第4分割中，我們獲得了68.88%的mIoU，相較于表現最好的基準LVM，提升了19.96%。這些顯著的提升凸顯了我們模型在多種場景中有效地進行物體分割和檢測的能力，尤其是在面對訓練中未見過的任務時。根據先前的工作 [4, 7]，我們還在著色任務中評估了我們的模型，其中較低的LPIPS和MSE值表示更優的性能。

如前面表1所示，本文的方法取得了0.26的LPIPS和0.24的MSE，顯著優于所有基準方法。這些結果凸顯了我們模型從灰度圖像中生成逼真自然顏色的能力，這對于恢復和藝術領域至關重要。

為了驗證本文模型理解3D場景幾何結構的能力，在NYU-v2深度估計和表面法線估計任務上進行了評估，如表2所示。由于Bai等人未在他們的論文中報告相關結果，我們使用他們官方的7B模型進行了評估。對于深度估計，本文的模型達到了6.2的AbsRel和96.1%的δ1閾值準確度，表現與Marigold和DPT等專家模型相當。在表面法線估計任務中，本文的方法取得了15.901的MAE和58.382的<11.25°閾值準確度，超越了強大的專家模型StableNormal。這一表現突顯了我們模型在準確估計表面方向方面的能力，增強了其在需要精確幾何理解的任務中的應用，如增強現實和3D重建。這些結果反映了我們模型在復雜環境中精確理解3D場景幾何結構的能力，這對現實世界的應用如3D場景重建和空間感知至關重要。

此外，還在圖像修復任務上將LaVin-DiT與LVM進行了比較。通過使用從ImageNet-1K驗證集中隨機選擇的2,500張圖像，本文的模型在FID上取得了1.65的成績，相較于LVM獲得的4.05，取得了顯著的改進。

定性分析如下圖3所示，展示了在多種基于圖像和基于視頻的任務中的定性結果。我們的模型始終遵循任務上下文，并準確地生成相應的預測。此外，給定帶有任務上下文的連續幀，模型能夠生成接下來的12幀預測，展示了其有效處理時間一致性和場景動態的能力。

可擴展性

為了研究所提出的LaVin-DiT的可擴展性，進行了三個模型大?。?.1B、1.0B和3.4B參數）的實驗，訓練了這三個模型100,000步。下圖4展示了訓練損失曲線，結果顯示，較大的模型在訓練過程中始終能夠取得較低的損失值。此外，3.4B模型收斂速度更快，在更少的訓練步驟內達到了較小的損失值。這一加速收斂表明，較大的模型更適合捕捉復雜的數據模式，從而提高了學習效率。觀察到的訓練動態強調了在復雜視覺任務中擴展模型容量的優勢，較大的模型能夠更有效地捕捉多樣的數據特征。

除了訓練動態外，模型大小對下游任務性能也有顯著影響。這在著色和深度估計任務中得到了體現，這些任務因其在捕捉顏色保真度和空間結構上的不同要求而被選中。如下圖5所示，隨著模型規模的增加，性能持續提高。具體而言，在著色任務中，3.4B模型的MSE為0.273，顯著優于1.0B和0.1B模型的MSE，分別為0.311和0.609。同樣，在深度估計中，3.4B模型的AbsRel為6.2，而1.0B和0.1B模型分別為6.5和7.6。這些結果表明，較大的模型在多個任務中確實能提供更強的性能，證明了LaVin-DiT是一個可擴展且適應性強的高性能視覺應用框架。

推理延遲分析

如下圖6所示，比較了LaVin-DiT和LVM（兩者均為7B模型）在不同分辨率下的推理延遲，結果表明本文的方法始終更高效。在256分辨率下，LaVin-DiT每個樣本僅需4.67秒，而LVM需要8.1秒；在更高分辨率下（例如512分辨率時，LaVin-DiT為20.1秒，LVM為47.2秒），這一效率差距進一步擴大。這個差異突顯了擴散模型在視覺任務中的一個關鍵優勢：與逐步處理tokens的自回歸模型不同，擴散模型可以并行處理tokens，從而在面對較大輸入時能夠更有效地擴展。這種并行性使得我們的LaVin-DiT在大規模視覺應用中成為更合適的選擇。

任務上下文長度的影響

上下文學習使得模型能夠通過少量示例適應新任務，隨著提供更多示例，性能通常會提高。我們通過評估任務上下文長度對十個下游任務的影響來進行研究。如下圖7所示，模型在任務上下文變長的情況下，性能持續提升，取得了顯著的性能增益。例如，隨著更多輸入-目標對的引入，LaVin-DiT在深度到圖像生成任務中取得了更低的FID，并在去除運動模糊任務中實現了更高的PSNR。這些結果表明，LaVin-DiT能夠有效利用擴展的任務上下文，突出其利用額外信息來增強任務適應性和準確性的能力。

結論

本文提出了LaVin-DiT，一個可擴展且統一的計算機視覺基礎模型，集成了空間-時間變分自編碼器和擴散Transformer，能夠高效處理高維視覺數據，同時保持空間和視覺一致性。通過上下文學習，LaVin-DiT能夠在不進行微調的情況下有效適應廣泛的任務，展現了其顯著的多功能性和適應性。廣泛的實驗驗證了LaVin-DiT的可擴展性和性能，確立了它作為開發通用視覺模型的有前景框架。

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

原文鏈接:??https://mp.weixin.qq.com/s/J46LIX-LcUrv3z-vI7LbXQ??