[Paper Reading] Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model

link
时间：24.08
机构：Waymo & University of Southern California

TL;DR

提出一种使用混合模态token来训练transformer，名为transfusion，是一种生成式AI模型。主要工作使用了2T的tokens结合语言模型的next token prediction以及diffusion训练了一个7B的模型，实验证明在图像与文本领域都有很好的效果。

比如，图像生成方面超过DALLE2与SDXL，文本生成达到Llama1的水平。

token准备：text模态被tokenize为token，image经过VAE Encoder编码为feature再通过Linear或者UNet提取image token(推测实际用的时候还是保留spatial信息的feature map，再输入diffusion model)。
Attention Mask：text使用mask将feature token屏蔽，而image token前后不需要屏蔽。推理过程text token部分是通过Auto-regressive模式逐个预测，而image token使用diffusion模式迭代T次输出后，再进行后续的text token regressive。
多模态Transformer：文本方面不解释了，image feature token经过Transformer是在不断生成噪声并对图像去噪，猜测使用DIT这种Diffusion模型，所以可以和text复用Transformer。参考这篇知乎文章也有可能是利用transformer直接seq2seq预测的(见下图)。