英偉達開源福利：影片生成、機器人都能用的SOTA tokenizer

机器之心發表於2024-11-23
機器人

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。英偉達開源福利：影片生成、機器人都能用的SOTA 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 的效能也優於以前的方法，顯示了出色的壓縮質量權衡。
圖片
參考連結：
https://developer.nvidia.com/blog/state-of-the-art-multimodal-generative-ai-model-development-with-nvidia-nemo/
https://research.nvidia.com/labs/dir/cosmos-tokenizer/
https://mp.weixin.qq.com/s/Hamz5XMT1tSZHKdPaCBTKg

相關文章