字節豆包全新圖像Tokenizer:生成圖像最低只需32個token,最高提速410倍 精華
在生成式模型的迅速發展中,Image Tokenization 扮演著一個很重要的角色,例如Diffusion依賴的VAE或者是Transformer依賴的VQGAN。這些Tokenizers會將圖像編碼至一個更為緊湊的隱空間(latent space),使得生成高分辨率圖像更有效率。
然而,現有的Tokenizer通常會將輸入圖像映射為隱空間的一個降采樣后的2D矩陣,這一設計隱式的限制了token與圖像之間的映射關系,導致其很難有效的利用圖像中的冗余信息(比如相鄰的區域經常會有類似的特征)來獲得一個更加有效的圖像編碼。
為了解決這一問題,字節跳動豆包大模型團隊和慕尼黑工業大學提出了全新的1D圖像Tokenizer:TiTok,這一Tokenizer打破了2D Tokenizer的設計局限,可以將整個圖片壓縮至更為緊湊的Token序列。
- 論文鏈接:??https://arxiv.org/abs/2406.07550 ??
- 項目鏈接:??https://yucornetto.github.io/projects/titok.html??
- 代碼鏈接:https://github.com/bytedance/1d-tokenizer
對于256 x 256分辨率的圖片,TiTok最少僅需32個Token就可以表達,比通常2D Tokenizer的256或1024個Token顯著減少。對于512 x 512分辨率的圖片,TiTok最少僅需64個Token,64倍小于Stable Diffusion的VAE Tokenizer。此外,在ImageNet圖像生成這一任務上,使用TiTok作為Tokenizer的生成器在生成質量和生成速度上都有顯著提高。
在256分辨率,TiTok獲得了1.97的FID,顯著超過使用同樣生成器的MaskGIT 4.21。在512分辨率TiTok可以獲得2.74的FID,不僅超過了DiT(3.04),并且相比DiT在圖像生成上加速了驚人的410倍!TiTok的最好變種取得了2.13的FID,顯著超過DiT的同時仍舊有著74倍的加速。
TiTok僅用32個Token就可以完成高質量的圖像重建與生成
圖像所需Token的顯著減少帶來了明顯更快的生成速度,但是同時維持了高質量的圖像生成。
模型結構
TiTok的結構非常簡單,編碼器和解碼器部分各自是一個ViT,在編碼過程中,一組latent tokens會拼接在image patches后,在過完編碼器后,僅保留latent tokens并進行quantization的過程。獲得的quantized latent tokens將會與一組mask tokens拼接在一起,一并送入解碼器,從mask token序列中重建出圖像。
1D Tokenization 性質研究
研究者進行了一系列實驗研究關于不同數量的用于表示圖像的token,不同的tokenizer大小,重建表現,生成表現,linear probing準確率,以及訓練和推理速度的比較。在這一過程中,研究者發現(1)僅需32個Token便能取得很好的重建與生成效果(2)通過增大Tokenizer的模型大小,研究者可以使用更少的Token來表示圖片(3)當圖片使用較少的Token來表示時,Tokenizer會學到更強的語義信息(4)使用更少的Token來表示圖片時,訓練和推理速度都有了顯著的提升。
此外,視頻中展示了使用不同的Tokenizer大小以及Token數目時所重建出的圖片,可以看到更大的Tokenizer可以在有限的Token下重建出質量更好的圖像。此外,當僅有有限Token時,模型更傾向于保留顯著區域有更好的重建效果。
實驗驗證
研究者主要在ImageNet-1k的256 x 256分辨率以及512 x 512分辨率上進行了與其他方法的比較。可以看到,盡管TiTok使用有限的Token數目,但是可以和其他使用更多Token的方法取得相當的重建效果(rFID),使用較少的Token數目讓TiTok在維持較高的生成圖片質量(gFID)的同時有著顯著快于其他方法的生成速度。
例如TiTok-L-32獲得了2.77的gFID score,同時可以以每秒101.6張圖片的速度進行圖片生成,這一速度顯著快于其他Diffusion Models (169倍快于DiT)或者Transformer Models (339倍快于ViT-VQGAN).
TiTok使用較少Token的優勢在更高分辨率的圖像生成上更加明顯,其中TiTok-L-64僅使用64個Token就能完成高質量的512分辨率圖像的重建與生成,生成圖像的質量不僅高于DiT (2.74 v.s. 3.04),同時生成速度提高了近410倍。
結論
在本文中,研究者專注于一個全新的1D Image Tokenizer,并提出了一種全新的Tokenizer來打破現有2D Tokenizer中的局限,進而更好的利用圖像中的冗余信息。TiTok僅需少量的Token(比如32個)來表示圖像,同時仍舊能進行高質量的圖像重建與生成。在ImageNet的256分辨率和512分辨率生成實驗中,TiTok不僅取得了超過Diffusion Models的生成質量,同時有著百倍更快的生成速度。
本文轉自 機器之心 ,作者:機器之心
