微軟開源視頻Tokenizer新SOTA!顯著優于Cosmos Tokenizer和Open-Sora
Sora、Genie等模型會都用到的Tokenizer,微軟下手了——
開源了一套全能的Video Tokenizer,名為VidTok。
Sora等視頻生成模型工作中,都會利用Tokenizer將原始的高維視頻數據(如圖像和視頻幀)轉換為更為緊湊的視覺Token,再以視覺Token為目標訓練生成模型。
而最新的VidTok,在連續和離散、不同壓縮率等多種設定下,各項指標均顯著優于SOTA模型。
以下是涵蓋PSNR、SSIM、FVD、LPIPS指標的性能比較雷達圖,面積越大表示性能越好。
從圖中可以看出對于離散Tokenizer,VidTok顯著優于英偉達Cosmos Tokenizer;對于連續Tokenizer,VidTok也比Open-Sora、CogVideoX有更高的性能。
這項研究由來自微軟亞研院、上海交通大學、北京大學的研究人員共同完成。
目前,VidTok代碼不僅開源了,還支持用戶在自定義數據集上的微調,為研究者和開發者提供了一個高性能、易用的工具平臺。
性能全面領先,適用各種場景
近年來,視頻生成以及基于此的世界模型已經成為人工智能領域的熱門研究方向,這兩者的核心在于對視頻內容的高效建模。
視頻中蘊含了豐富的視覺信息,不僅能夠提供真實的視覺體驗,更能作為具身場景中模型理解世界的中間媒介。
然而,由于視頻像素級表示信息高度冗余,如何通過Tokenizer對視頻數據進行高效壓縮和表示成為關鍵課題。
當下很多工作如Sora,Genie等都會通過Tokenizer將原始的高維視頻數據(如圖像和視頻幀)轉換為更為緊湊的視覺Token,再以視覺Token為目標訓練生成模型。
可以說,視覺Token的表示能力對于最終的效果至關重要,甚至決定了模型能力的上限。
Tokenizer的主要作用是將高維的原始數據轉換為隱空間中高效的壓縮表示,使得信息的生成和處理可以在該隱空間中進行。上圖展示了一個視頻的Token化過程,通過轉換為Token建模,能夠有效降低模型訓練和推理時的計算需求。
根據不同的使用需求,視頻Tokenizer通常有如下分類:
- 連續型和離散型。根據隱空間的數值分布,Tokenizer可以分為連續型和離散型,分別適用于從連續分布中采樣的模型(如擴散模型等)和從離散分布中采樣的模型(如語言模型等)。
- 因果型和非因果型。因果結構使得模型只依賴歷史幀來對當前幀進行Tokenization,這與真實世界系統的因果性質保持一致。非因果模型則可以同時根據歷史幀和未來幀對當前幀進行Tokenization,通常具有更優的重建質量。
- 不同的壓縮率模型。Sora等眾多工作采用了如4x8x8的視頻壓縮率(時間壓縮4倍、空間壓縮8倍),實現更高的視頻壓縮率而保持高質量的視頻重建是目前的研究趨勢。
目前業界領先的視頻模型多為閉源狀態,而開源的視頻Tokenizer大多受限于單一的模型設定或欠佳的重建質量,導致可用性較差。
由此,來自微軟亞研院、上海交通大學和北京大學的研究人員最近正式發布了開源視頻Tokenizer——VidTok。
在測試中,VidTok性能全面領先,適用各種場景。
如下表所示,VidTok支持多樣化的隱空間且具有靈活的壓縮率,同時支持因果和非因果模型,以適應不同的使用需求。
- 對于連續型Tokenizer,支持不同的視頻壓縮率、不同的隱空間通道數,同時支持因果和非因果模型。
- 對于離散型Tokenizer,支持不同的視頻壓縮率、不同的碼本大小,同時支持因果和非因果模型。
更多模型在持續更新中。
為了全面評估VidTok在各個設定下的重建性能,作者將VidTok與最先進的連續和離散視頻Tokenizer分別進行了對齊設定下的比較。所有模型均為4x8x8倍視頻壓縮率的因果模型,主要包含以下三種設定:
- VidTok-FSQ:離散型,碼本大小各異。基線方法包括MAGVIT-v2,OmniTokenizer,Cosmos-DV等。
- VidTok-KL-4chn:連續型,隱空間通道數為4。基線方法包括CV-VAE,Open-Sora-v1.2,Open-Sora-Plan-v1.2等。
- VidTok-KL-16chn:連續型,隱空間通道數為16。基線方法包括CogVideoX,Cosmos-CV等。
定量實驗結果表明,VidTok在上述三種設定下均達到了SOTA性能,在常見的視頻質量評估指標PSNR、SSIM、FVD、LPIPS上具有全面的優勢。
與現有的離散Tokenizer相比,VidTok即使在使用更小的碼本大小時(例如32,768),也展現出了更優的重建性能。
在連續Tokenizer的設定下,無論隱空間通道數是4還是16,VidTok在所有評估指標上相比基線方法均取得了全面的提升。值得注意的是,這些提升是在沒有模型大小優勢的情況下達成的。
除此之外,團隊還進行了定性分析。
上圖中展示了以上三種設定下的與基線方法的對比結果。
從視頻幀的重建質量可以看出,與現有的方法相比,VidTok在各種模型設定下,均展現出了最優的重建細節保真度和主觀視覺質量。證明了VidTok作為多功能視頻Tokenizer的有效性。
所以VidTok是如何做到的?
VidTok的技術亮點解析
相對于現有的視頻Tokenizer,VidTok在模型架構、量化技術、訓練策略上分別做了創新。
高效的混合模型架構設計
VidTok采用經典的3D編碼器-解碼器結構,同時創新性地結合了3D、2D和1D卷積,有效地解耦空間和時間采樣。
在現有研究中普遍認為,盡管計算成本較高,完全的3D架構提供了更優的重建質量。然而,VidTok發現將部分3D卷積替換為2D和1D卷積的組合,可以有效地解耦空間和時間采樣,在降低計算需求的同時,保持了高水平的重建質量。
詳細的網絡架構如上圖所示。VidTok分別處理空間采樣和時間采樣模塊,并在時間采樣模塊中引入了AlphaBlender操作符。其余組件,包括輸入/輸出層和瓶頸層,則利用3D卷積來促進信息融合。此外,整個架構中引入了層歸一化以增強穩定性和性能。實驗證明該架構在重建質量和計算量之間取得了平衡。
先進的量化技術
VidTok引入了有限標量量化(FSQ)技術,無需顯式學習碼本,顯著提高了模型的訓練穩定性和重建性能。
有限標量量化(FSQ)由「Finite scalar quantization: Vq-vae made simple」提出,其核心原理是,在隱空間表征中,每個標量條目通過四舍五入獨立量化到最近的預定義標量值。
與傳統VQ相比,FSQ無需學習顯式的碼本,從而提高了訓練的穩定性。實驗表明,FSQ在碼本利用率、重建質量和訓練穩定性方面具有顯著優勢,作為一種先進的量化技術,有效提升了離散Tokenizer的性能。
增強的訓練策略
VidTok采用分階段訓練策略,訓練時間減少了50%,而重建質量不受影響。
視頻Tokenizer的訓練通常是計算密集的,要求大量計算資源(例如對于256x256分辨率的視頻需要3,072GPU小時的訓練時長)。這就需要開發有效的策略來降低計算成本,同時保持模型性能。
VidTok采用一種兩階段訓練方法來應對這一挑戰:首先在低分辨率視頻上對完整模型進行預訓練,然后僅在高分辨率視頻上微調解碼器。這種訓練策略顯著降低了計算成本——訓練時間減少了一半(從3,072GPU小時降至1,536GPU 小時),而保持重建視頻質量不變。
該兩階段訓練的另一優勢是,由于第二階段只會微調解碼器,因此模型可以快速適應到新的領域數據中,而不會影響隱空間數據分布。
此外,由于視頻Tokenizer旨在建模輸入視頻的運動動態,因此在模型中有效表示這些動態至關重要。VidTok使用較低幀率的數據進行訓練,顯著增強了模型捕捉和表示運動動態的能力,獲得了更好的重建質量。
VidTok的開源為視頻生成、世界模型領域提供了新的工具,特別是在當前業內許多領先模型仍未開源的背景下。
團隊表示,VidTok支持后續微調也為其他應用提供了更廣闊的使用空間,研究者可輕松將VidTok應用于特定領域數據集,為目標場景優化性能。
更多細節內容感興趣的童鞋可參閱原論文。
論文地址:https://arxiv.org/abs/2412.13061。
項目地址:https://github.com/microsoft/vidtok。