魔法降臨!港中文MagicDriveDiT:暴力提升生成質量40%!
寫在前面 & 筆者的個人理解
擴散模型的快速發展極大地改善了視頻合成,特別是在可控視頻生成方面,這對自動駕駛等應用至關重要。然而,現有的方法受到可擴展性和控制條件集成方式的限制,無法滿足自動駕駛應用對高分辨率和長視頻的需求。本文介紹了一種基于DiT架構的新方法MagicDriveDiT,并解決了這些挑戰。我們的方法通過流匹配增強了可擴展性,并采用漸進式訓練策略來管理復雜的場景。通過結合時空條件編碼,MagicDriveDiT實現了對時空延遲的精確控制。綜合實驗表明,它在生成具有更高分辨率和更多幀的逼真街道場景視頻方面具有優越的性能。Magic-DriveDiT顯著提高了視頻生成質量和時空控制,擴大了其在自動駕駛各種任務中的潛在應用。
- 項目主頁:https://gaoruiyuan.com/magicdrivedit/
總結來說,本文的主要貢獻如下:
- 本文設計了一個高效的框架MagicDriveDiT,利用漸進式引導來實現高質量的高分辨率長視頻生成。
- 為目標位置、道路語義和相機軌跡開發了新的時空控制,同時保持了多幀、多視圖的一致性。
- MagicDriveDiT通過混合分辨率和持續時間訓練,從圖像到視頻生成都得到了很好的推廣,具有外推能力,大大超過了之前工作中的分辨率和幀數。
相關工作回顧
自動駕駛中的視頻生成。視頻生成對于自動駕駛至關重要,可應用于訓練感知模型、測試和場景重建。它需要廣闊的視野和動態目標運動處理,需要生成模型可控性和具有更多幀和多相機視角的高分辨率視頻制作。擴散模型改進了可控多視圖視頻生成,但現有模型缺乏足夠的分辨率和幀數,無法用于數據引擎應用和策略測試。如表1所示,我們的工作MagicDriveDiT在自動駕駛視頻生成方面實現了前所未有的分辨率和幀數。
擴散模型和DiT架構。擴散模型通過學習從高斯噪聲到樣本的去噪步驟來生成數據,廣泛應用于圖像和視頻生成。從建模的角度來看,流匹配簡化了擴散模型的建模,提高了訓練和推理的效率。從實現的角度來看,由于更好的可擴展性,擴散模型的架構從UNet轉變為DiT,特別是對于高分辨率任務。我們的MagicDriveDiT還利用流匹配和DiT來擴展到高分辨率和長視頻生成。
有條件生成。條件生成對于使用生成模型的各種應用程序至關重要。LDM和ControlNet的網格形控制信號加性編碼的交叉關注層是可控擴散生成的領先方法。在街景生成中,MagicDrive和MagicDrive3D集成了3D邊界框、BEV地圖、自車軌跡和相機姿態,用于多視圖街景合成。然而,這些方法僅限于空間編碼,并不直接適用于時空VAE延遲,如圖2所示。MagicDriveDiT為時空壓縮的VAE延遲提供了一種新的控制范式,使高分辨率長視頻的生成可控。
MagicDriveDiT方法詳解
對于圖3所示的模型架構,MagicDriveDiT引入了一種基于DiT的新型擴散模型,用于可控街景生成,,該模型利用STDiT-3塊。這些塊具有單獨的空間和時間模塊,以增強信息處理。
架構設計包含兩個重大修改。首先為了促進多視圖生成,多視圖DiT(MVDiT)塊集成了一個跨視圖注意力層,如圖3左側所示。其次,考慮到需要處理多個控制元素,MagicDriveDiT對文本、框、相機視圖和軌跡采用交叉注意力,同時對地圖采用額外的分支來輸入控制信號。
然而如圖2所示,控制信號的空間編碼與時空潛伏期不兼容。因此,每個控制信號的編碼器都被重新制定,如圖3右側的時空框編碼器所示。
此外,MagicDriveDiT采用漸進式引導訓練策略,從圖像過渡到低分辨率和短視頻,最終過渡到高分辨率和長視頻,以增強擴散模型的收斂性。最后的訓練階段包含可變長度和分辨率的視頻,使模型能夠生成不同分辨率的圖像和視頻,并推斷出訓練設置之外的更長幀數。
Design for High-Resolution Long Video
DiT和3D VAE。用于高分辨率和長視頻生成的訓練擴散模型是計算密集型的,需要大量的GPU內存。DiT和3D VAE在這些方面對擴散模型的縮放至關重要。
正如Peebles和Xie所指出的那樣,模型Gflops和FID之間存在很強的負相關關系,這使得DiT在計算效率方面優于UNet等架構。對于GPU內存,3D VAE提供時間降采樣。典型的2D VAE在高度和寬度上都將圖像壓縮8倍,達到64倍的壓縮比。相比之下,3D VAE將時間信息壓縮了4倍,得到了256倍的壓縮比,顯著減少了壓縮后的潛在序列長度和內存消耗,這對轉換器特別有利。此外,架構統一的趨勢允許DiT利用先進的并行化方法,例如序列并行,來克服單個GPU內存的限制,從而促進更高的分辨率和更長的視頻生成。
采用3D VAE的一個主要挑戰是幾何控制。如圖2所示,幾何控制在空間上管理每幀內容。使用2D VAE,T幀被編碼為T延遲。使用T幾何描述符,視頻幾何控制降級為對圖像的空間控制,因為控制信號和延遲在時間上是對齊的。然而,3D VAE會產生T/f延遲(其中f是時間壓縮比),使控制信號與延遲不匹配,并使之前的控制技術無效。
Spatial-Temporal Conditioning Techniques。對于MagicDriveDiT中的幾何控制,我們引入了時空編碼來將控制信號與時空延遲對齊。這涉及重新對齊場景描述符中的貼圖、框和軌跡。以網格數據表示的地圖易于管理。通過擴展ControlNet的設計,我們利用3D VAE中具有新可訓練參數的時間下采樣模塊來對齊控制塊和基塊之間的特征,如圖4(a)所示。
對于3D框,填充應用于不可見的框,以保持視圖和幀之間框序列長度的一致性,如圖3右側所示。我們使用帶有時間變換器和RoPE的下采樣模塊來捕獲時間相關性,創建與視頻延遲對齊的時空嵌入,如圖4(b)所示。通過用相機姿勢的MLP替換盒子的MLP,盒子的時空編碼器也可以適應自車軌跡。所有下采樣比都與采用的3D VAE對齊,即分別以8n或8n+1作為輸入,2n或2n+1作為輸出。
Progressive Bootstrap Training
為了在訓練過程中加快模型收斂,我們根據單次訓練迭代的持續時間來調度訓練數據。具體來說,我們采用三階段訓練方法:最初使用低分辨率圖像進行引導訓練,過渡到低分辨率短視頻,最終采用高分辨率長視頻訓練。
這種訓練策略基于兩個觀察結果。首先,在可控生成中,我們注意到,在學習可控性之前,該模型最初會針對更高的內容質量進行優化,Gao等人也觀察到了這種模式。從頭開始訓練模型需要多次迭代才能收斂。我們的漸進過渡方法使模型能夠更快地獲得可控性。其次,在階段轉換期間,我們發現與長視頻可控性相比,該模型更快地適應高分辨率。因此,我們從第一階段開始訓練可控性,并專注于用(短)視頻而不是圖像優化更多的迭代。
Variable Length and Resolution Adaptation
如前文所述,我們采用三階段訓練流程。在最后兩個階段,我們將不同分辨率和長度的視頻用于訓練。具體來說,在第二階段,我們使用長度不超過65幀、最大分辨率為424×800的視頻混合進行訓練。在第三階段,我們混合了高達241幀(數據集的最大幀數)和高達848×1600(數據集最大分辨率)的視頻。
與單一分辨率和長度的訓練相比,這種混合方法使模型能夠快速理解分辨率和幀數維度的信息。因此,我們的最終模型支持生成各種分辨率(224×400、424×800、848×1600)和幀數(1-241幀)。它還支持跨兩個維度的外推,允許在訓練配置之外進行采樣,例如129幀時的848×1600或241幀時的424×800。
實驗
結果和分析
生成質量。MagicDriveDiT在視頻和圖像生成任務方面都表現出色。在視頻任務中,與MagicDrive相比,它顯著降低了FVD(表2),這是由于DiT架構增強了幀間一致性和時空條件編碼,用于精確控制目標運動和定位。如圖6所示,MagicDriveDiT生成的高分辨率視頻不僅可以提高質量,還可以包含更復雜的細節,與真實相機拍攝的鏡頭非常相似。這一增強是通過我們對可變長度和分辨率的高級訓練實現的,可以實現更真實和詳細的輸出。
MagicDriveDiT也受益于混合訓練方法,能夠生成圖像。如表3所示,MagicDriveDiT在多視圖街景任務中與基線性能相匹配,在車輛分割mIoU和目標檢測mAP中超越了基線。這展示了我們時空條件編碼的強大泛化能力。
可控性。表2和表3中的定量結果表明,MagicDriveDiT生成的圖像和視頻有效地反映了控制條件。此外,圖8提供了可視化結果,表明多個控制條件可以獨立影響生成的內容。例如,天氣可以通過文本輸入(從晴天到雨天)來改變,同時保持道路結構以及其他車輛和自車的軌跡。通過改變條件的組合,MagicDriveDiT能夠制作各種高質量的街景視頻。
消融實驗
街道景觀的VAE比較。在訓練擴散模型之前,我們評估了開源3D VAE(即CogVAE和Open Sora)在街景上的性能,并與2D SD VAE進行了比較。如圖7所示,CogVAE在重建能力方面始終優于同行。此外,如附錄E所示,CogVAE在較長的視頻序列上表現出最小的性能下降,使其特別適合長視頻生成任務。此外,我們發現所有VAE都表現出隨著分辨率的提高而提高的重建能力。這種見解有助于通過專注于更高分辨率的輸出來提高我們的模型生成高質量圖像和視頻的能力。
時空Conditioning。我們通過過擬合實驗中的驗證損失(圖9)和可視化比較(圖10)證明了時空編碼器的有效性。我們比較了兩種基線:全局時間維度縮減(Reduce)和時間維度插值(Interp.)進行對齊。在16個樣本的過擬合訓練中,4×下采樣技術(我們的4×下)加速了收斂,并實現了最低的最終驗證損失,如圖9所示。此外,圖10顯示,與全局縮減基線不同,4×down可以減少偽影并保持精確的運動軌跡。這些結果證實了時空編碼器利用時空延遲增強數據編碼和提高視頻生成性能的能力。
漸進式Bootstrap訓練。與直接的第3階段訓練相比,三階段漸進訓練方法顯著提高了模型訓練效率。表4顯示,例如,在4天內,階段1執行的迭代次數比階段3多約60倍,階段2實現的迭代次數多約7倍。漸進式訓練對于受控生成模型至關重要,可控生成模型需要大量迭代才能有效收斂。漸進式策略能夠快速獲得高質量的視頻生成能力,在早期階段利用更快的迭代來增強收斂性和加快學習速度。
可變長度和分辨率訓練。MagicDriveDiT采用了一種混合不同長度和分辨率視頻的訓練策略。我們的消融研究證明了這一策略的有效性。如表5所示,當僅使用低分辨率視頻(17×224×400)時,VAE的局限性是顯而易見的,與其他設置相比,FVD較差,mAP和mIoU評分較低。合并較長的視頻可以提高模型的可控性(mAP和mIoU都更高),而合并高分辨率視頻可以提高整體質量(所有三個指標都得到了顯著改善)。
盡管混合不同的幀長度會略微降低FVD,但這對于使模型能夠生成各種長度的視頻并推斷出看不見的長度至關重要(見第6節)。因此,我們結合了分辨率和幀長,有效地平衡了視頻質量、可控性和模型功能之間的權衡。
結論和討論
在這篇論文中,我們介紹了MagicDriveDiT,這是一個具有精確控制的高分辨率和長視頻合成的創新框架,專門為自動駕駛等應用量身定制。通過利用DiT架構和流匹配,MagicDriveDiT有效地解決了視頻生成中可擴展性和幾何控制的挑戰。我們的方法引入了新穎的時空條件編碼和具有可變長度和分辨率自適應的漸進自舉策略,增強了模型的泛化能力。大量實驗表明,MagicDriveDiT可以生成保持空間和時間連貫性的逼真視頻,在分辨率和幀數方面大大超過了以前的方法。這項工作推進了視頻合成,并為自動駕駛中的模擬和其他應用開辟了新的可能性。