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