在討論圖像、視頻生成模型時(shí)，人們的焦點(diǎn)更多地集中在模型所采用的架構(gòu)，比如大名鼎鼎的 DiT。但其實(shí)，tokenizer 也是非常重要的組件。

谷歌等機(jī)構(gòu)的研究者曾在一篇題為「Language model Beats diffusion - tokenizer is key to visual generation」的論文中證明，一個(gè)好的 tokenizer 接入到語(yǔ)言模型后，能夠立即獲得比當(dāng)時(shí)最好的 diffusion 模型還要好的效果。論文作者蔣路在后來(lái)接受采訪時(shí)表示，「我們的研究可能會(huì)讓社區(qū)意識(shí)到 tokenizer 是被嚴(yán)重忽視的一個(gè)領(lǐng)域，值得發(fā)力去做」。

在圖像、視頻生成模型中，tokenizer 的核心作用是將連續(xù)的、高維的視覺數(shù)據(jù)（如圖像和視頻幀）轉(zhuǎn)換成模型可以處理的形式，即緊湊的語(yǔ)義 token，它的視覺表示能力對(duì)于模型的訓(xùn)練和生成過(guò)程至關(guān)重要。就像上述論文作者所說(shuō)，「tokenizer 的存在就是通過(guò)建立 token 之間的互聯(lián)，讓模型明確『我現(xiàn)在要做什么』，互聯(lián)建立得越好、LLM 模型越有機(jī)會(huì)發(fā)揮它的全部潛力?！?/span>

tokenizer 是生成式 AI 的關(guān)鍵組件，它通過(guò)無(wú)監(jiān)督學(xué)習(xí)發(fā)現(xiàn)潛在空間，從而將原始數(shù)據(jù)轉(zhuǎn)換為高效的壓縮表示。視覺 tokenizer 專門將圖像和視頻等高維視覺數(shù)據(jù)轉(zhuǎn)化為緊湊的語(yǔ)義 token，從而實(shí)現(xiàn)高效的大型模型訓(xùn)練，并降低推理的計(jì)算需求。圖中展示了一個(gè)視頻 token 化過(guò)程。

當(dāng)前，業(yè)界有很多可用的開源視頻、圖像 tokenizer，但這些 tokenizer 經(jīng)常生成質(zhì)量不佳的數(shù)據(jù)表示，這會(huì)造成采用該 tokenizer 的模型生成失真的圖像、不穩(wěn)定的視頻。此外，低效的 token 化過(guò)程還會(huì)導(dǎo)致編解碼速度變慢、訓(xùn)練和推理時(shí)間變長(zhǎng)，從而對(duì)開發(fā)人員的工作效率和用戶體驗(yàn)產(chǎn)生負(fù)面影響。

為了解決這些問(wèn)題，來(lái)自英偉達(dá)的研究者開源了一套名為 Cosmos 的全新 tokenizer。

• 研究地址：https://research.nvidia.com/labs/dir/cosmos-tokenizer/
• HuggingFace 地址：https://huggingface.co/collections/nvidia/cosmos-tokenizer-672b93023add81b66a8ff8e6

一般來(lái)說(shuō)，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)不同的計(jì)算限制。

Cosmos tokenizer 基于輕量級(jí)時(shí)間因果架構(gòu)，使用因果時(shí)間卷積和注意力層來(lái)保持視頻幀的順序。這種統(tǒng)一的設(shè)計(jì)允許對(duì)圖像和視頻進(jìn)行無(wú)縫 token 化。

英偉達(dá)的研究者在高分辨率圖像和長(zhǎng)視頻上訓(xùn)練 Cosmos tokenizer，涵蓋不同類別數(shù)據(jù)的寬高比（包括 1:1、3:4、4:3、9:16 和 16:9）。在推理過(guò)程中，它不受時(shí)間長(zhǎng)度的影響，可以處理比訓(xùn)練時(shí)間更長(zhǎng)的數(shù)據(jù)。

• GitHub 地址：https://github.com/NVIDIA/Cosmos-Tokenizer

研究者在標(biāo)準(zhǔn)數(shù)據(jù)集（包括 MS-COCO 2017、ImageNet-1K、FFHQ、CelebA-HQ 和 DAVIS）上對(duì) Cosmos tokenizer 進(jìn)行了評(píng)估。為了使視頻 tokenizer 評(píng)估標(biāo)準(zhǔn)化，他們還策劃了一個(gè)名為 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 上編碼長(zhǎng)達(dá) 8 秒的 1080p 和 10 秒的 720p 視頻?？臻g壓縮率為 8 倍和 16 倍、時(shí)間壓縮率為 4 倍和 8 倍的預(yù)訓(xùn)練模型可在 GitHub 上獲取。

試用過(guò) Cosmos 的 1x 機(jī)器人公司 AI 副總裁 Eric Jang 表示，Cosmos 是一個(gè)非常好的 tokenizer，比根據(jù)他們自己的數(shù)據(jù)進(jìn)行微調(diào)的 Magvit2 好得多?？磥?lái)，這個(gè)新工具值得一試。

以下是 Cosmos 的一些技術(shù)細(xì)節(jié)。

Cosmos tokenizer 架構(gòu)

Cosmos tokenizer 采用復(fù)雜的編碼器 - 解碼器結(jié)構(gòu)，旨在實(shí)現(xiàn)高效率和高效學(xué)習(xí)。其核心是采用 3D 因果卷積塊，這是聯(lián)合處理時(shí)空信息的專門層，并利用因果時(shí)間注意力捕捉數(shù)據(jù)中的長(zhǎng)程依賴關(guān)系。

因果結(jié)構(gòu)確保模型在進(jìn)行 token 化時(shí)只使用過(guò)去和現(xiàn)在的幀，而避免使用未來(lái)幀。這對(duì)于與許多真實(shí)世界系統(tǒng)的因果性質(zhì)保持一致至關(guān)重要，例如物理 AI 或多模態(tài) LLM 中的系統(tǒng)。

Cosmos tokenizer 架構(gòu)圖。

使用 3D wavelet 對(duì)輸入進(jìn)行降采樣，這種信號(hào)處理技術(shù)能更有效地表示像素信息。數(shù)據(jù)處理完成后，通過(guò)反向 wavelet 變換重建原始輸入。

這種方法提高了學(xué)習(xí)效率，使 tokenizer 編碼器 - 解碼器可學(xué)習(xí)模塊專注于有意義的特征，而不是多余的像素細(xì)節(jié)。這些技術(shù)與其獨(dú)特的訓(xùn)練方法相結(jié)合，使 Cosmos tokenizer 成為了一個(gè)高效、強(qiáng)大的架構(gòu)。

實(shí)驗(yàn)結(jié)果

定性結(jié)果

圖 6 顯示了使用連續(xù)視頻 tokenizer 重建的視頻幀。

圖 9 顯示了使用不同離散圖像 tokenizer 重建的圖像。

圖 8 則顯示了連續(xù)圖像 tokenizer 的誤差圖，以突出重建差異。與之前的方法相比，Cosmos tokenizer 能更有效地保留結(jié)構(gòu)和高頻細(xì)節(jié)（如草地、樹枝、文本），同時(shí)將視覺失真（如人臉、文本）和偽影降到最低。

這些定性結(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)衡。