只需Image Tokenizer，Llama也能做圖像生成了，而且效果超過了擴(kuò)散模型。

來自港大和字節(jié)的研究人員，提出了基于自回歸模型Llama的圖像生成方法。

目前該模型已經(jīng)開源，并在GitHub斬獲了近900顆星標(biāo)。

圖片

擴(kuò)散模型出現(xiàn)后，取代了自回歸方法，一度成為圖像生成的主流技術(shù)路線。

但在ImageNet測試基準(zhǔn)上，作者提出的LlamaGen表現(xiàn)超越了LDM、DiT等擴(kuò)散模型。

作者的這一發(fā)現(xiàn)，證明了最原始的自回歸模型架構(gòu)同樣可以實(shí)現(xiàn)極具競爭力的圖像生成性能。

LlamaGen生圖示例，第一行為class調(diào)控生成，第二行為文生圖

那么，基于自回歸模型，或者說基于Llama的圖像生成，是如何實(shí)現(xiàn)的呢？

用自回歸模型做圖像生成

作者介紹，開源社區(qū)對(duì)自回歸模型做圖像生成的印象大多停留在2020年的VQ-GAN的ImageNet基準(zhǔn)上取得的15左右的FID分?jǐn)?shù)。

然而，早在2021年的ViT-VQGAN已經(jīng)達(dá)到了FID 3.0左右的性能，DALL-E 1，Parti等更是在文生圖領(lǐng)域展現(xiàn)了巨大的潛力。

不過這些工作都沒有開源，于是，研究團(tuán)隊(duì)將目標(biāo)設(shè)定成了推出開源版的基于自回歸圖像生成模型。

針對(duì)現(xiàn)有的先進(jìn)的圖像生成模型，作者總結(jié)出其成功的三點(diǎn)關(guān)鍵設(shè)計(jì)：

• 圖像壓縮/量化器（Image Compressors/Tokenizers）
• 可scale up的圖像生成模型（Scalable Image generation models）
• 高質(zhì)量的訓(xùn)練數(shù)據(jù)（High-quality Training Data）

于是，作者采用了與VQ-GAN同樣的CNN架構(gòu)，將連續(xù)的圖像轉(zhuǎn)化成離散的Token。

相比2020年的VQ-GAN，作者對(duì)Image Tokenizer有了更多的認(rèn)知：

一個(gè)優(yōu)秀的Tokenizer需要更大的Codebook Size，更低的Codebook Vector Dimension，同時(shí)，更好的圖像重建需要更多的Token數(shù)量。

圖片

△VQ-GAN架構(gòu)，非本項(xiàng)目

架構(gòu)方面，LlamaGen的模型架構(gòu)主要基于Llama語言模型，包括使用RMSNorm的Pre-Normalization、SwiGLU和RoPE。

盡管圖像生成領(lǐng)域一些常用的技術(shù)（如AdaLN）可能進(jìn)一步提高性能，但作者還是盡可能保持與Llama語言模型一模一樣的架構(gòu)。

在Class-Conditional和Text-Conditional（文生圖）圖像生成模型中，作者采用了使用最簡單的實(shí)現(xiàn)：

Class或文本嵌入直接作為起始Token，后續(xù)的Image Token應(yīng)用next-Token預(yù)測范式產(chǎn)生。

訓(xùn)練的過程則分為兩個(gè)階段進(jìn)行。

在第一階段，模型在LAION-COCO的50M子集上進(jìn)行訓(xùn)練，圖像分辨率為 256×256。

LAION-COCO原始數(shù)據(jù)集有6億圖文對(duì)，作者通過有效的圖像URL、美學(xué)分?jǐn)?shù)、水印分?jǐn)?shù)、CLIP圖文相似度分?jǐn)?shù)和圖像大小來篩選這些圖像。

圖片

在第二階段，模型在1千萬規(guī)模的內(nèi)部高美學(xué)質(zhì)量圖像上進(jìn)行微調(diào)，圖像分辨率為512×512。

這些美學(xué)圖像的文本描述由LLaVa產(chǎn)生。

圖片

到了部署階段，基于原生自回歸模型架構(gòu)的圖像生成模型可以無縫采用現(xiàn)有的LLM部署框架，例如vLLM。這也是統(tǒng)一模型架構(gòu)的一大優(yōu)勢。

同時(shí)，基于vLLM的框架部署方式，為LlamaGen帶來了326%-414%的加速。

圖片

效果不輸擴(kuò)散模型

那么，作者研究出的這款模型效果究竟怎樣呢？

先說作者重新訓(xùn)練的Image Tokenizer，它在ImageNet和COCO上優(yōu)于以前的Tokenizers，包括VQGAN，ViT-VQGAN和MaskGI等。

重要的是，基于離散表征的Tokenizer與基于連續(xù)表征的VAE性能持平（例如在擴(kuò)散模型中被廣泛使用的SD VAE），這表明圖像量化的離散表征不再是圖像重建的一大瓶頸。

圖片

實(shí)際生成過程中，在ImageNet測試集上，LlamaGen在FID、IS、Precision和Recall等指標(biāo)上都表現(xiàn)出了極強(qiáng)的競爭力。

其中，LlamaGen-3B模型優(yōu)于廣為流行的擴(kuò)散模型 LDM和DiT。這表明最樸素的自回歸模型架構(gòu)有能力作為先進(jìn)圖像生成系統(tǒng)的基礎(chǔ)模型。

同時(shí)，與之前的自回歸模型相比，LlamaGen在各個(gè)參數(shù)量級(jí)上均優(yōu)于以前的模型。

作者分析，這樣的成績是得益于更好的Image Tokenizer和Llama架構(gòu)更好的擴(kuò)展性。

圖片

文生圖方面，經(jīng)過第一階段的訓(xùn)練，模型基本擁有了圖文對(duì)齊的能力，但其生成圖像的視覺質(zhì)量有待提高。

第二階段的訓(xùn)練顯著提高了生成圖像的視覺質(zhì)量，作者認(rèn)為這種提高來自兩個(gè)方面——

• 第二階段的訓(xùn)練使用了高質(zhì)量的美學(xué)圖像；
• 第一階段的圖像分辨率是256x256，第二階段是512x512，更大的圖像分辨率會(huì)帶來更好的視覺效果。

圖片

當(dāng)輸入更長的文本時(shí)，LlamaGen也可以生成兼具圖文對(duì)齊與視覺質(zhì)量的圖像。

圖片

不過作者也坦言，如果類比擴(kuò)散模型的發(fā)展路線，目前的LlamaGen只是做到了Stable Diffusion v1階段，未來的改進(jìn)方向包括SDXL（更大的分辨率，更多的Aspect Ratio），ControlNet（更高的可控性），Sora（視頻生成）。

從多模態(tài)大模型的視角看，自回歸模型分別實(shí)現(xiàn)理解任務(wù)和生成任務(wù)都被證明了可行性，下一步就是在同一個(gè)模型中聯(lián)合訓(xùn)練。

目前該項(xiàng)目已經(jīng)開源，而且還支持在線體驗(yàn)，感興趣的話不妨一試。

在線體驗(yàn)：https://huggingface.co/spaces/FoundationVision/LlamaGen

論文地址：https://arxiv.org/abs/2406.06525

項(xiàng)目主頁：https://peizesun.github.io/llamagen/

GitHub：https://github.com/FoundationVision/LlamaGenHugging Face：https://huggingface.co/FoundationVision/LlamaGen