英偉達開源福利:視頻生成、機器人都能用的SOTA tokenizer
在討論圖像、視頻生成模型時,人們的焦點更多地集中在模型所采用的架構,比如大名鼎鼎的 DiT。但其實,tokenizer 也是非常重要的組件。
谷歌等機構的研究者曾在一篇題為「Language model Beats diffusion - tokenizer is key to visual generation」的論文中證明,一個好的 tokenizer 接入到語言模型后,能夠立即獲得比當時最好的 diffusion 模型還要好的效果。論文作者蔣路在后來接受采訪時表示,「我們的研究可能會讓社區意識到 tokenizer 是被嚴重忽視的一個領域,值得發力去做」。
在圖像、視頻生成模型中,tokenizer 的核心作用是將連續的、高維的視覺數據(如圖像和視頻幀)轉換成模型可以處理的形式,即緊湊的語義 token,它的視覺表示能力對于模型的訓練和生成過程至關重要。就像上述論文作者所說,「tokenizer 的存在就是通過建立 token 之間的互聯,讓模型明確『我現在要做什么』,互聯建立得越好、LLM 模型越有機會發揮它的全部潛力。」
tokenizer 是生成式 AI 的關鍵組件,它通過無監督學習發現潛在空間,從而將原始數據轉換為高效的壓縮表示。視覺 tokenizer 專門將圖像和視頻等高維視覺數據轉化為緊湊的語義 token,從而實現高效的大型模型訓練,并降低推理的計算需求。圖中展示了一個視頻 token 化過程。
當前,業界有很多可用的開源視頻、圖像 tokenizer,但這些 tokenizer 經常生成質量不佳的數據表示,這會造成采用該 tokenizer 的模型生成失真的圖像、不穩定的視頻。此外,低效的 token 化過程還會導致編解碼速度變慢、訓練和推理時間變長,從而對開發人員的工作效率和用戶體驗產生負面影響。
為了解決這些問題,來自英偉達的研究者開源了一套名為 Cosmos 的全新 tokenizer。
- 研究地址:https://research.nvidia.com/labs/dir/cosmos-tokenizer/
- HuggingFace 地址:https://huggingface.co/collections/nvidia/cosmos-tokenizer-672b93023add81b66a8ff8e6
一般來說,tokenizer 有兩種類型:連續型和離散型。連續 tokenizer 將視覺數據映射為連續嵌入,適用于從連續分布中采樣的模型,如 Stable Diffusion。離散 tokenizer 將視覺數據映射為量化指數,適用于 VideoPoet 等依賴交叉熵損失進行訓練的模型,類似于 GPT 模型。下圖比較了這些 token 類型。
tokenizer 必須兼顧高壓縮和高質量,保留潛在空間的視覺細節。Cosmos tokenizer 是一套全面的連續和離散圖像和視頻視覺 tokenizer,可提供出色的壓縮和高質量重建,速度是以前方法的 12 倍。
如表 1 所示,它支持各種圖像和視頻類型,具有靈活的壓縮率,以適應不同的計算限制。
Cosmos tokenizer 基于輕量級時間因果架構,使用因果時間卷積和注意力層來保持視頻幀的順序。這種統一的設計允許對圖像和視頻進行無縫 token 化。
英偉達的研究者在高分辨率圖像和長視頻上訓練 Cosmos tokenizer,涵蓋不同類別數據的寬高比(包括 1:1、3:4、4:3、9:16 和 16:9)。在推理過程中,它不受時間長度的影響,可以處理比訓練時間更長的數據。
- GitHub 地址:https://github.com/NVIDIA/Cosmos-Tokenizer
研究者在標準數據集(包括 MS-COCO 2017、ImageNet-1K、FFHQ、CelebA-HQ 和 DAVIS)上對 Cosmos tokenizer 進行了評估。為了使視頻 tokenizer 評估標準化,他們還策劃了一個名為 TokenBench 的新數據集,涵蓋機器人、駕駛和體育等類別,并在 GitHub 上公開發布。
- TokenBench 地址:https://github.com/NVlabs/TokenBench
結果(圖 1)顯示,Cosmos tokenizer 明顯優于現有方法,在 DAVIS 視頻上的 PSNR 提升了 4 dB。它的 token 化速度是以前方法的 12 倍,并能在配備 80GB 內存的英偉達 A100 GPU 上編碼長達 8 秒的 1080p 和 10 秒的 720p 視頻。空間壓縮率為 8 倍和 16 倍、時間壓縮率為 4 倍和 8 倍的預訓練模型可在 GitHub 上獲取。
試用過 Cosmos 的 1x 機器人公司 AI 副總裁 Eric Jang 表示,Cosmos 是一個非常好的 tokenizer,比根據他們自己的數據進行微調的 Magvit2 好得多。看來,這個新工具值得一試。
以下是 Cosmos 的一些技術細節。
Cosmos tokenizer 架構
Cosmos tokenizer 采用復雜的編碼器 - 解碼器結構,旨在實現高效率和高效學習。其核心是采用 3D 因果卷積塊,這是聯合處理時空信息的專門層,并利用因果時間注意力捕捉數據中的長程依賴關系。
因果結構確保模型在進行 token 化時只使用過去和現在的幀,而避免使用未來幀。這對于與許多真實世界系統的因果性質保持一致至關重要,例如物理 AI 或多模態 LLM 中的系統。
Cosmos tokenizer 架構圖。
使用 3D wavelet 對輸入進行降采樣,這種信號處理技術能更有效地表示像素信息。數據處理完成后,通過反向 wavelet 變換重建原始輸入。
這種方法提高了學習效率,使 tokenizer 編碼器 - 解碼器可學習模塊專注于有意義的特征,而不是多余的像素細節。這些技術與其獨特的訓練方法相結合,使 Cosmos tokenizer 成為了一個高效、強大的架構。
實驗結果
定性結果
圖 6 顯示了使用連續視頻 tokenizer 重建的視頻幀。
圖 9 顯示了使用不同離散圖像 tokenizer 重建的圖像。
圖 8 則顯示了連續圖像 tokenizer 的誤差圖,以突出重建差異。與之前的方法相比,Cosmos tokenizer 能更有效地保留結構和高頻細節(如草地、樹枝、文本),同時將視覺失真(如人臉、文本)和偽影降到最低。
這些定性結果表明,Cosmos tokenizer 能夠編碼和解碼各種視覺內容,并有能力保持圖像和視頻的最高視覺質量。
定量結果
表 2 和表 3 列出了連續和離散視頻 tokenizer 在各種基準上的平均定量指標。Cosmos tokenizer 在 4×8×8 壓縮率的 DAVIS 和 TokenBench 數據集上都達到了 SOTA 性能。即使在更高的壓縮率(8×8×8 和 8×16×16)下,Cosmos tokenizer 的性能也優于以前的方法,顯示了出色的壓縮質量權衡。