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

亮點(diǎn)直擊

• 高效處理高維視覺數(shù)據(jù)
• 創(chuàng)新的聯(lián)合擴(kuò)散Transformer
• 強(qiáng)大的任務(wù)泛化能力
• 顯著性能與效率提升

總結(jié)速覽

解決的問題

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

提出的方案

• 空間-時間變分自編碼器 (spatial-temporal variational autoencoder)
  編碼高維視覺數(shù)據(jù)到連續(xù)的潛在空間，降低數(shù)據(jù)維度。
• 聯(lián)合擴(kuò)散Transformer (joint diffusion transformer)
  通過逐步擴(kuò)散生成視覺輸出，優(yōu)化生成建模能力。
• 基于上下文學(xué)習(xí)的統(tǒng)一多任務(wù)訓(xùn)練 (in-context learning)
  利用輸入-目標(biāo)對作為任務(wù)上下文，引導(dǎo)擴(kuò)散Transformer在潛在空間中完成特定任務(wù)。推理時，通過任務(wù)上下文和測試數(shù)據(jù)實(shí)現(xiàn)跨任務(wù)的泛化，無需微調(diào)。

應(yīng)用的技術(shù)

• 變分自編碼器 (VAE) 技術(shù)用于高效編碼視覺數(shù)據(jù)。
• 擴(kuò)散模型與Transformer的結(jié)合，用于逐步生成高質(zhì)量視覺結(jié)果。
• 上下文學(xué)習(xí)（in-context learning）機(jī)制，支持多任務(wù)統(tǒng)一訓(xùn)練和推理。
• 可擴(kuò)展性：模型參數(shù)從 0.1B 擴(kuò)展至 3.4B，充分驗(yàn)證了其可擴(kuò)展性能。

達(dá)到的效果

• 性能提升：支持超過20+視覺任務(wù)，多項(xiàng)SOTA。
• 統(tǒng)一框架：首次在生成框架中實(shí)現(xiàn)大規(guī)模、多任務(wù)視覺建模。
• 高效推理：通過上下文學(xué)習(xí)，在無需微調(diào)的情況下實(shí)現(xiàn)任務(wù)泛化。
• 開源貢獻(xiàn)：計劃開放代碼和模型，為后續(xù)研究提供支持。

方法

LaVin-DiT 模塊

ST-VAE

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

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

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

J-DiT

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

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

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

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

LaVin-DiT 推理

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

實(shí)驗(yàn)

設(shè)置

訓(xùn)練數(shù)據(jù)
為了統(tǒng)一多個計算機(jī)視覺任務(wù)，構(gòu)建了一個大規(guī)模的多任務(wù)數(shù)據(jù)集，涵蓋了室內(nèi)和室外環(huán)境，跨越了現(xiàn)實(shí)世界和合成領(lǐng)域。該數(shù)據(jù)集包含約320萬張獨(dú)特圖像和60萬段獨(dú)特視頻，涵蓋了超過20個任務(wù)：

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

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

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

評估協(xié)議
在一系列涵蓋圖像和視頻領(lǐng)域的計算機(jī)視覺任務(wù)上評估我們的模型。按照已建立的協(xié)議，報告每個任務(wù)的標(biāo)準(zhǔn)評估指標(biāo)。

主要結(jié)果

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

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

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

如前面表1所示，本文的方法取得了0.26的LPIPS和0.24的MSE，顯著優(yōu)于所有基準(zhǔn)方法。這些結(jié)果凸顯了我們模型從灰度圖像中生成逼真自然顏色的能力，這對于恢復(fù)和藝術(shù)領(lǐng)域至關(guān)重要。

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

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

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

可擴(kuò)展性

為了研究所提出的LaVin-DiT的可擴(kuò)展性，進(jìn)行了三個模型大小（0.1B、1.0B和3.4B參數(shù)）的實(shí)驗(yàn)，訓(xùn)練了這三個模型100,000步。下圖4展示了訓(xùn)練損失曲線，結(jié)果顯示，較大的模型在訓(xùn)練過程中始終能夠取得較低的損失值。此外，3.4B模型收斂速度更快，在更少的訓(xùn)練步驟內(nèi)達(dá)到了較小的損失值。這一加速收斂表明，較大的模型更適合捕捉復(fù)雜的數(shù)據(jù)模式，從而提高了學(xué)習(xí)效率。觀察到的訓(xùn)練動態(tài)強(qiáng)調(diào)了在復(fù)雜視覺任務(wù)中擴(kuò)展模型容量的優(yōu)勢，較大的模型能夠更有效地捕捉多樣的數(shù)據(jù)特征。

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

推理延遲分析

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

任務(wù)上下文長度的影響

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

結(jié)論

本文提出了LaVin-DiT，一個可擴(kuò)展且統(tǒng)一的計算機(jī)視覺基礎(chǔ)模型，集成了空間-時間變分自編碼器和擴(kuò)散Transformer，能夠高效處理高維視覺數(shù)據(jù)，同時保持空間和視覺一致性。通過上下文學(xué)習(xí)，LaVin-DiT能夠在不進(jìn)行微調(diào)的情況下有效適應(yīng)廣泛的任務(wù)，展現(xiàn)了其顯著的多功能性和適應(yīng)性。廣泛的實(shí)驗(yàn)驗(yàn)證了LaVin-DiT的可擴(kuò)展性和性能，確立了它作為開發(fā)通用視覺模型的有前景框架。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

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