詳解Latte:去年底上線的全球首個開源文生視頻DiT
隨著 Sora 的成功發(fā)布,視頻 DiT 模型得到了大量的關(guān)注和討論。設(shè)計穩(wěn)定的超大規(guī)模神經(jīng)網(wǎng)絡(luò)一直是視覺生成領(lǐng)域的研究重點。DiT [1] 的成功為圖像生成的規(guī)模化提供了可能性。
然而,由于視頻數(shù)據(jù)的高度結(jié)構(gòu)化與復(fù)雜性,如何將 DiT 擴展到視頻生成領(lǐng)域卻是一個挑戰(zhàn),來自上海人工智能實驗室的研究團隊聯(lián)合其他機構(gòu)通過大規(guī)模的實驗回答了這個問題。
早在去年 11 月,該團隊就已經(jīng)開源了一款與 Sora 技術(shù)相似的自研模型:Latte。作為全球首個開源文生視頻 DiT,Latte 受到了廣泛關(guān)注,并且模型設(shè)計被眾多開源框架所使用與參考,如 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 包含兩個主要模塊:預(yù)訓(xùn)練 VAE 和視頻 DiT。預(yù)訓(xùn)練 VAE 編碼器將視頻逐幀從像素空間壓縮到隱空間,視頻 DiT 對隱式表征提取 token 并進行時空建模,最后 VAE 解碼器將特征映射回像素空間生成視頻。為了得到最優(yōu)的視頻質(zhì)量,作者著重探究了 Latte 設(shè)計中兩個重要內(nèi)容,(1) 視頻 DiT 模型整體結(jié)構(gòu)設(shè)計以及 (2) 模型與訓(xùn)練細節(jié)的最優(yōu)設(shè)計(The best practices)。
(1)Latte 整體模型結(jié)構(gòu)設(shè)計探究
圖 1. Latte 模型結(jié)構(gòu)及其變體
作者提出了 4 種不同的 Latte 變體 (圖 1),從時空注意力機制的角度設(shè)計了兩種 Transformer 模塊,同時在每種模塊中分別研究了兩種變體(Variant):
1. 單注意力機制模塊,每個模塊中只包含時間或者空間注意力。
- 時空交錯式建模 (Variant 1): 時間模塊插入到各個空間模塊之后。
- 時空順序式建模 (Variant 2): 時間模塊整體置于空間模塊之后。
2. 多注意力機制模塊,每個模塊中同時包含時間與空間注意力機制 (Open-sora所參考變體)。
- 串聯(lián)式時空注意力機制 (Variant 3): 時空注意力機制串行建模。
- 并聯(lián)式時空注意力機制 (Variant 4): 時空注意力機制并行建模并特征融合。
實驗表明 (圖 2),通過對 4 種模型變體設(shè)置相同的參數(shù)量,變體 4 相較于其他三種變體在 FLOPS 上有著明顯的差異,因此 FVD 上也相對最高,其他 3 種變體總體性能類似,變體 1 取得了最優(yōu)異的性能,作者計劃未來在大規(guī)模的數(shù)據(jù)上做更加細致的討論。
圖 2. 模型結(jié)構(gòu) FVD
(2)Latte 模型與訓(xùn)練細節(jié)的最優(yōu)設(shè)計探究(The best practices)
除了模型總體結(jié)構(gòu)設(shè)計,作者還探究了其他模型與訓(xùn)練中影響生成效果的因素。
1.Token 提取:探究了單幀 token(a)和時空 token(b)兩種方式,前者只在空間層面壓縮 token,后者同時壓縮時空信息。實驗顯示單幀 token 要優(yōu)于時空 token(圖 4)。與 Sora 進行比較,作者猜測 Sora 提出的時空 token 是通過視頻 VAE 進行了時間維度的預(yù)壓縮,而在隱空間上與 Latte 的設(shè)計類似都只進行了單幀 token 的處理。
圖 3. Token 提取方式,(a) 單幀 token 和 (b) 時空 token
圖 4. Token 提取 FVD
2. 條件注入模式:探究了(a)S-AdaLN 和(b)all tokens 兩種方式 (圖 5)。S-AdaLN 通過 MLP 將條件信息轉(zhuǎn)換為歸一化中的變量注入到模型中。All token 形式將所有條件轉(zhuǎn)化為統(tǒng)一的 token 作為模型的輸入。實驗證明,S-AdaLN 的方式相較于 all token 對于獲得高質(zhì)量的結(jié)果更加有效 (圖 6)。原因是,S-AdaLN 可以使信息被直接注入到每一個模塊。而 all token 需要將條件信息從輸入逐層傳遞到最后,存在著信息流動過程中的損失。
圖 5. (a) S-AdaLN 和 (b) all tokens。
圖 6. 條件注入方式 FVD
3. 時空位置編碼:探究了絕對位置編碼與相對位置編碼。不同的位置編碼對最后視頻質(zhì)量影響很小 (圖 7)。由于生成時長較短,位置編碼的不同不足以影響視頻質(zhì)量,對于長視頻生成,這一因素需要被重新考慮。
圖 7. 位置編碼方式 FVD
4. 模型初始化:探究使用 ImageNet 預(yù)訓(xùn)練參數(shù)初始化對模型性能的影響。實驗表明,使用 ImageNet 初始化的模型具有較快的收斂速度,然而,隨著訓(xùn)練的進行,隨機初始化的模型卻取得了較好的結(jié)果 (圖 8)。可能的原因在于 ImageNet 與訓(xùn)練集 FaceForensics 存在著比較大的分布差異,因此未能對模型的最終結(jié)果起到促進作用。而對于文生視頻任務(wù)而言,該結(jié)論需要被重新考慮。在通用數(shù)據(jù)集的分布上,圖像與視頻的內(nèi)容空間分布相似,使用預(yù)訓(xùn)練 T2I 模型對于 T2V 可以起到極大的促進作用。
圖 8. 初始化參數(shù) FVD
5. 圖像視頻聯(lián)合訓(xùn)練:將視頻與圖像壓縮為統(tǒng)一 token 進行聯(lián)合訓(xùn)練,視頻 token 負責(zé)優(yōu)化全部參數(shù),圖像 token 只負責(zé)優(yōu)化空間參數(shù)。聯(lián)合訓(xùn)練對于最終的結(jié)果有著顯著的提升 (表 2 和表 3),無論是圖片 FID,還是視頻 FVD,通過聯(lián)合訓(xùn)練都得到了降低,該結(jié)果與基于 UNet 的框架 [2][3] 是一致的。
6. 模型尺寸:探究了 4 種不同的模型尺寸,S,B,L 和 XL (表 1)。擴大視頻 DiT 規(guī)模對于提高生成樣本質(zhì)量有著顯著的幫助 (圖 9)。該結(jié)論也證明了在視頻擴散模型中使用 Transformer 結(jié)構(gòu)對于后續(xù) scaling up 的正確性。
表 1. Latte 不同尺寸模型規(guī)模
圖 9. 模型尺寸 FVD
定性與定量分析
作者分別在 4 個學(xué)術(shù)數(shù)據(jù)集(FaceForensics,TaichiHD,SkyTimelapse 以及 UCF101)進行了訓(xùn)練。定性與定量(表 2 和表 3)結(jié)果顯示 Latte 均取得了最好的性能,由此可以證明模型整體設(shè)計是具有優(yōu)異性的。
表 2. UCF101 圖片質(zhì)量評估
表 3. Latte 與 SoTA 視頻質(zhì)量評估
文生視頻擴展
為了進一步證明 Latte 的通用性能,作者將 Latte 擴展到了文生視頻任務(wù),利用預(yù)訓(xùn)練 PixArt-alpha [4] 模型作為空間參數(shù)初始化,按照最優(yōu)設(shè)計的原則,在經(jīng)過一段時間的訓(xùn)練之后,Latte 已經(jīng)初步具備了文生視頻的能力。后續(xù)計劃通過擴大規(guī)模驗證 Latte 生成能力的上限。
討論與總結(jié)
Latte 作為全世界首個開源文生視頻 DiT,已經(jīng)取得了很有前景的結(jié)果,但由于計算資源的巨大差異,在生成清晰度,流暢度上以及時長上與 Sora 相比還具有不小的差距。團隊歡迎并在積極尋求各種合作,希望通過開源的力量,打造出性能卓越的自主研發(fā)大規(guī)模通用視頻生成模型。