詳解Latte:去年底上線的全球首個開源文生影片DiT

机器之心發表於2024-03-28

隨著 Sora 的成功釋出,影片 DiT 模型得到了大量的關注和討論。設計穩定的超大規模神經網路一直是視覺生成領域的研究重點。DiT [1] 的成功為影像生成的規模化提供了可能性。

然而,由於影片資料的高度結構化與複雜性,如何將 DiT 擴充套件到影片生成領域卻是一個挑戰,來自上海人工智慧實驗室的研究團隊聯合其他機構透過大規模的實驗回答了這個問題。

早在去年 11 月,該團隊就已經開源了一款與 Sora 技術相似的自研模型:Latte。作為全球首個開源文生影片 DiT,Latte 受到了廣泛關注,並且模型設計被眾多開源框架所使用與參考,如 Open-Sora Plan (PKU) 和 Open-Sora (ColossalAI)。

圖片

  • 開源連結:https://github.com/Vchitect/Latte

  • 專案主頁:https://maxin-cn.github.io/latte_project/

  • 論文連結:https://arxiv.org/pdf/2401.03048v1.pdf

先來看下Latte的影片生成效果。

圖片

方法介紹

總體上,Latte 包含兩個主要模組:預訓練 VAE 和影片 DiT。預訓練 VAE 編碼器將影片逐幀從畫素空間壓縮到隱空間,影片 DiT 對隱式表徵提取 token 並進行時空建模,最後 VAE 解碼器將特徵對映回畫素空間生成影片。為了得到最優的影片質量,作者著重探究了 Latte 設計中兩個重要內容,(1) 影片 DiT 模型整體結構設計以及 (2) 模型與訓練細節的最優設計(The best practices)。

(1)Latte 整體模型結構設計探究

圖片

圖 1. Latte 模型結構及其變體

作者提出了 4 種不同的 Latte 變體 (圖 1),從時空注意力機制的角度設計了兩種 Transformer 模組,同時在每種模組中分別研究了兩種變體(Variant):

1. 注意力機制模組,每個模組中只包含時間或者空間注意力

  • 時空交錯式建模 (Variant 1): 時間模組插入到各個空間模組之後。

  • 時空順序式建模 (Variant 2): 時間模組整體置於空間模組之後。

2. 注意力機制模組,每個模組中同時包含時間與空間注意力機制 (Open-sora所參考變體)

  • 串聯式時空注意力機制 (Variant 3): 時空注意力機制序列建模。

  • 並聯式時空注意力機制 (Variant 4): 時空注意力機制並行建模並特徵融合。

實驗表明 (圖 2),透過對 4 種模型變體設定相同的引數量,變體 4 相較於其他三種變體在 FLOPS 上有著明顯的差異,因此 FVD 上也相對最高,其他 3 種變體總體效能類似,變體 1 取得了最優異的效能,作者計劃未來在大規模的資料上做更加細緻的討論。

圖片

圖 2. 模型結構 FVD

(2)Latte 模型與訓練細節的最優設計探究(The best practices)

除了模型總體結構設計,作者還探究了其他模型與訓練中影響生成效果的因素。

1.Token 提取:探究了單幀 token(a)和時空 token(b)兩種方式,前者只在空間層面壓縮 token,後者同時壓縮時空資訊。實驗顯示單幀 token 要優於時空 token(圖 4)。與 Sora 進行比較,作者猜測 Sora 提出的時空 token 是透過影片 VAE 進行了時間維度的預壓縮,而在隱空間上與 Latte 的設計類似都只進行了單幀 token 的處理。

圖片

圖 3. Token 提取方式,(a) 單幀 token 和 (b) 時空 token

圖片

圖 4. Token 提取 FVD

2. 條件注入模式:探究了(a)S-AdaLN 和(b)all tokens 兩種方式 (圖 5)。S-AdaLN 透過 MLP 將條件資訊轉換為歸一化中的變數注入到模型中。All token 形式將所有條件轉化為統一的 token 作為模型的輸入。實驗證明,S-AdaLN 的方式相較於 all token 對於獲得高質量的結果更加有效 (圖 6)。原因是,S-AdaLN 可以使資訊被直接注入到每一個模組。而 all token 需要將條件資訊從輸入逐層傳遞到最後,存在著資訊流動過程中的損失。

圖片

圖 5. (a) S-AdaLN 和 (b) all tokens。

圖片

圖 6. 條件注入方式 FVD

3. 時空位置編碼:探究了絕對位置編碼與相對位置編碼。不同的位置編碼對最後影片質量影響很小 (圖 7)。由於生成時長較短,位置編碼的不同不足以影響影片質量,對於長影片生成,這一因素需要被重新考慮。

圖片

圖 7. 位置編碼方式 FVD

4. 模型初始化:探究使用 ImageNet 預訓練引數初始化對模型效能的影響。實驗表明,使用 ImageNet 初始化的模型具有較快的收斂速度,然而,隨著訓練的進行,隨機初始化的模型卻取得了較好的結果 (圖 8)。可能的原因在於 ImageNet 與訓練集 FaceForensics 存在著比較大的分佈差異,因此未能對模型的最終結果起到促進作用。而對於文生影片任務而言,該結論需要被重新考慮。在通用資料集的分佈上,影像與影片的內容空間分佈相似,使用預訓練 T2I 模型對於 T2V 可以起到極大的促進作用。

圖片

圖 8. 初始化引數 FVD

5. 影像影片聯合訓練:將影片與影像壓縮為統一 token 進行聯合訓練,影片 token 負責最佳化全部引數,影像 token 只負責最佳化空間引數聯合訓練對於最終的結果有著顯著的提升 (表 2 和表 3),無論是圖片 FID,還是影片 FVD,透過聯合訓練都得到了降低,該結果與基於 UNet 的框架 [2][3] 是一致的。

6. 模型尺寸:探究了 4 種不同的模型尺寸,S,B,L 和 XL (表 1)。擴大影片 DiT 規模對於提高生成樣本質量有著顯著的幫助 (圖 9)。該結論也證明了在影片擴散模型中使用 Transformer 結構對於後續 scaling up 的正確性。

圖片

表 1. Latte 不同尺寸模型規模

圖片

圖 9. 模型尺寸 FVD

定性與定量分析

作者分別在 4 個學術資料集(FaceForensics,TaichiHD,SkyTimelapse 以及 UCF101)進行了訓練。定性與定量(表 2 和表 3)結果顯示 Latte 均取得了最好的效能,由此可以證明模型整體設計是具有優異性的。

圖片

表 2. UCF101 圖片質量評估

圖片

表 3. Latte 與 SoTA 影片質量評估

文生影片擴充套件

為了進一步證明 Latte 的通用效能,作者將 Latte 擴充套件到了文生影片任務,利用預訓練 PixArt-alpha [4] 模型作為空間引數初始化,按照最優設計的原則,在經過一段時間的訓練之後,Latte 已經初步具備了文生影片的能力。後續計劃透過擴大規模驗證 Latte 生成能力的上限。

討論與總結

Latte 作為全世界首個開源文生影片 DiT,已經取得了很有前景的結果,但由於計算資源的巨大差異,在生成清晰度,流暢度上以及時長上與 Sora 相比還具有不小的差距。團隊歡迎並在積極尋求各種合作,希望透過開源的力量,打造出效能卓越的自主研發大規模通用影片生成模型。

參考文獻

[1] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[2] Ho, Jonathan, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303 (2022)

[3] Wang, Yaohui, et al. "Lavie: High-quality video generation with cascaded latent diffusion models." arXiv preprint arXiv:2309.15103 (2023).

[4] Chen, Junsong, et al. "PixArt-$\alpha $: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis." arXiv preprint arXiv:2310.00426 (2023).

相關文章