Llama也能做图像生成！港大字节推出开源自回归文生图模型，在线体验已开放

克雷西 2024-07-03 16:09:37 来源：量子位

LlamaGen团队 投稿
量子位 | 公众号 QbitAI

只需Image Tokenizer，Llama也能做图像生成了，而且效果超过了扩散模型。

来自港大和字节的研究人员，提出了基于自回归模型Llama的图像生成方法。

目前该模型已经开源，并在GitHub斩获了近900颗星标。

扩散模型出现后，取代了自回归方法，一度成为图像生成的主流技术路线。

但在ImageNet测试基准上，作者提出的LlamaGen表现超越了LDM、DiT等扩散模型。

作者的这一发现，证明了最原始的自回归模型架构同样可以实现极具竞争力的图像生成性能。

△LlamaGen生图示例，第一行为class调控生成，第二行为文生图

那么，基于自回归模型，或者说基于Llama的图像生成，是如何实现的呢？

用自回归模型做图像生成

作者介绍，开源社区对自回归模型做图像生成的印象大多停留在2020年的VQ-GAN的ImageNet基准上取得的15左右的FID分数。

然而，早在2021年的ViT-VQGAN已经达到了FID 3.0左右的性能，DALL-E 1，Parti等更是在文生图领域展现了巨大的潜力。

不过这些工作都没有开源，于是，研究团队将目标设定成了推出开源版的基于自回归图像生成模型。

针对现有的先进的图像生成模型，作者总结出其成功的三点关键设计：

• 图像压缩/量化器（Image Compressors/Tokenizers）
• 可scale up的图像生成模型（Scalable Image generation models）
• 高质量的训练数据（High-quality Training Data）

于是，作者采用了与VQ-GAN同样的CNN架构，将连续的图像转化成离散的Token。

相比2020年的VQ-GAN，作者对Image Tokenizer有了更多的认知：

一个优秀的Tokenizer需要更大的Codebook Size，更低的Codebook Vector Dimension，同时，更好的图像重建需要更多的Token数量。

△VQ-GAN架构，非本项目

架构方面，LlamaGen的模型架构主要基于Llama语言模型，包括使用RMSNorm的Pre-Normalization、SwiGLU和RoPE。

尽管图像生成领域一些常用的技术（如AdaLN）可能进一步提高性能，但作者还是尽可能保持与Llama语言模型一模一样的架构。

在Class-Conditional和Text-Conditional（文生图）图像生成模型中，作者采用了使用最简单的实现：

Class或文本嵌入直接作为起始Token，后续的Image Token应用next-Token预测范式产生。

训练的过程则分为两个阶段进行。

在第一阶段，模型在LAION-COCO的50M子集上进行训练，图像分辨率为 256×256。

LAION-COCO原始数据集有6亿图文对，作者通过有效的图像URL、美学分数、水印分数、CLIP图文相似度分数和图像大小来筛选这些图像。

在第二阶段，模型在1千万规模的内部高美学质量图像上进行微调，图像分辨率为512×512。

这些美学图像的文本描述由LLaVa产生。

到了部署阶段，基于原生自回归模型架构的图像生成模型可以无缝采用现有的LLM部署框架，例如vLLM。这也是统一模型架构的一大优势。

同时，基于vLLM的框架部署方式，为LlamaGen带来了326%-414%的加速。

效果不输扩散模型

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

先说作者重新训练的Image Tokenizer，它在ImageNet和COCO上优于以前的Tokenizers，包括VQGAN，ViT-VQGAN和MaskGI等。

重要的是，基于离散表征的Tokenizer与基于连续表征的VAE性能持平（例如在扩散模型中被广泛使用的SD VAE），这表明图像量化的离散表征不再是图像重建的一大瓶颈。

实际生成过程中，在ImageNet测试集上，LlamaGen在FID、IS、Precision和Recall等指标上都表现出了极强的竞争力。

其中，LlamaGen-3B模型优于广为流行的扩散模型 LDM和DiT。这表明最朴素的自回归模型架构有能力作为先进图像生成系统的基础模型。

同时，与之前的自回归模型相比，LlamaGen在各个参数量级上均优于以前的模型。

作者分析，这样的成绩是得益于更好的Image Tokenizer和Llama架构更好的扩展性。

文生图方面，经过第一阶段的训练，模型基本拥有了图文对齐的能力，但其生成图像的视觉质量有待提高。

第二阶段的训练显著提高了生成图像的视觉质量，作者认为这种提高来自两个方面——

• 第二阶段的训练使用了高质量的美学图像；
• 第一阶段的图像分辨率是256×256，第二阶段是512×512，更大的图像分辨率会带来更好的视觉效果。

当输入更长的文本时，LlamaGen也可以生成兼具图文对齐与视觉质量的图像。

不过作者也坦言，如果类比扩散模型的发展路线，目前的LlamaGen只是做到了Stable Diffusion v1阶段，未来的改进方向包括SDXL（更大的分辨率，更多的Aspect Ratio），ControlNet（更高的可控性），Sora（视频生成）。

从多模态大模型的视角看，自回归模型分别实现理解任务和生成任务都被证明了可行性，下一步就是在同一个模型中联合训练。

目前该项目已经开源，而且还支持在线体验，感兴趣的话不妨一试。

在线体验：
https://huggingface.co/spaces/FoundationVision/LlamaGen
论文地址：
https://arxiv.org/abs/2406.06525
项目主页：
https://peizesun.github.io/llamagen/
GitHub：
https://github.com/FoundationVision/LlamaGen
Hugging Face：
https://huggingface.co/FoundationVision/LlamaGen