【NeurlPS 2024】MAR：无矢量量化的自回归图像生成

1 前言

​ 本期内容，我们将MAR（掩码自回归模型）。与传统的视觉自回归需要把图像离散化成一个个token的方式不同。MAR构建的是一个连续型的自回归模型。
视频：MAR：无矢量量化的自回归图像生成

2 引入

​ 对于传统的自回归模型而言，比如VQGAN+GPT2，一般来说是需要先用VQGAN把图像量化成1D的token。然后再使用仅解码器的Transformer进行图像生成，比如GPT2。而本篇论文提出了一个问题——“图像自回归的离散化是否是必须的？”

​ 论文认为这并不是必须的，因此，他们便想去构建一个连续型的视觉自回归模型。如何去构建呢？在我们的印象中，什么样的东西是直接对连续型数据直接建模的呢？用脚想一下，扩散模型不就是吗？让我们先来回顾一下传统的离散视觉自回归模型

3 离散视觉自回归模型

​ 对于离散视觉自回归模型，往往我们会把图像离散化成一个个的token，比如我们离散化成n个token—— { x 1 , x 2 , ⋯   , x n } \{x^1,x^2,\cdots,x^n\} {x1,x2,⋯,xn}。里面的每一个token可以表示一个个的整数。一般来说，视觉自回归使用的是的仅解码器架构，所以注意力使用的是因果注意力，我们可以把这些token的关系分解为
$$p(x^1,x^2,\cdots ,x^n)=\prod _{i=1}^{n}p(x^i|x^1,\cdots ,x^{i-1})$$
​ 通俗来讲，就是当前时刻的状态，仅依赖于过去的所有时刻，而与未来无关。

​ 为了与论文保持一致的结果，接下来我们定义$x$为我们下一个要预测的token。记词汇表的大小为K（表示$x$的取值只有K种，即$0\le x\le K$）。自回归模型进行生成的时候，会先生成一个连续型的D维向量，即$z\in {\mathbb{R}}^D$。接着，会使用一个投影矩阵$W\in {\mathbb{R}}^{K\times D}$将$z$投影到K个类别大小，然后使用使用softmax将其转换为概率，我们可以把该过程用概率去表示$p(x|z)=\text{softmax}(Wz)$。

​ 所以目前来说，我们认为下一个要预测的token $x$ ，可以通过建模概率分布$p(x|z)$来获得，离散自回归是将其表示为$\text{softmax}(Wz)$。那么问题就来了，我们是否可以用其他的方法去获得这个概率分布呢？

​ 当然可以！我们之前所学习过的生成模型，不就是能够学习概率分布吗。假设我们能够学习出$p(x|z)$，那么当我们提供$z$之后，便可获得$x$，也就是生成$x$。如果我们用其他生成模型去建模$p(x|z)$，比如扩散模型，那么也就没有离散化的硬性要求了，因为扩散模型本身就可以处理连续型的数据。

4 Diffusion Loss

​ 我们用扩散模型去拟合$p(x|z)$这个分布，拟合生成之后，给定一个$z$，我们就可以获得$x$了。

​ 训练：

​

​ 即图中所示，先获取自回归模型的$z$，把他作为条件，然后利用$z$和$x$去训练扩散模型。即
$$\mathcal{L}(z,x)={\mathbb{E}}_{ϵ,t}\left[\Vert ϵ-{ϵ}_{\theta }(x_t|t,z){\Vert }^2\right]\text{\hspace{1em}(1)}$$
​ 其中，$ϵ$是一个高斯白噪声，$x_t=\sqrt{{\bar{\alpha }}_t}x+\sqrt{1-{\bar{\alpha }}_t}ϵ$，${\bar{\alpha }}_t$是噪声调度，t是噪声调度的时间步。如图中所示，扩散模型的网络结构只是一个小型的MLP网络，${ϵ}_{\theta }(x_t|t,z)$表示将$x_t$作为网络的主输入，而$t$和$z$作为旁系条件输入。里面的$z$是由自回归网络$f$生成的，即$z=f(\cdot )$

​ 注意到期望不仅对$ϵ$求，还对时刻t求了。为什么呢？论文提到，他们设计的去噪网络的结构很小，因此他们对于给定一个编码向量$z$，会采样多个t。这有助于提高损失函数的利用率，而无需重新计算z。在训练得时候，每张图像抽样4次。

​ 采样：
$$x_{t-1}=\frac{1}{\sqrt{{\bar{\alpha }}_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t|t,z)\right)+{\sigma }_t\delta$$
​ 其中，$\delta$是高斯白噪声，${\sigma }_t$是时间步t的噪声水平。初始化$x_T\sim N(0,I)$，然后条件去噪得到$x_0$，即得到$x_0\sim p(x|z)$。去噪的时候，会引入一个温度系数$\mathcal{T}$，用于缩放${\sigma }_t\delta$。直观上，就是通过调整方差来控制样本多样性。

5 Diffusion Loss for Autoregressive Models

​ 给定序列图像划分成一堆token { x 1 , x 2 , ⋯   , x n } \{x^1,x^2,\cdots,x^n\} {x1,x2,⋯,xn}，我们的目标为
$$p(x^1,\cdots ,x^n)=\prod _{i=1}^{n}p(x^i|x^1,\cdots ,x^{i-1})$$
​ 请注意，这里的$x^i$是连续的，我们记自回归网络为$f$。在预测$x^i$的时候，它依赖于$x^1$到$x^{i-1}$时刻的状态，所以，我们把他们送给自回归网络$f$以获得$z^i:z^i=f(x^1,\dots ,x^{i-1})$。然后，利用扩散模型计算$p(x^i|z^i)$来生成$x^i$。

6 双向注意力可以执行自回归

​ 传统观念认为，只有因果注意力才能够执行自回归生成图像。而本文设计了一个双向注意力的自回归模型。此模型是基于MAE（掩码自编码器），运行流程难以描述出来，见视频当中。

7 MAR(掩码自回归模型)

​ 掩码自回归模型能够不仅可以预测下一个token，还可以直接预测下一个token集合。
$$p(x^1,\cdots ,x^n)=p(X^1,\cdots ,X^K)=\prod _k^{K}p(X^k|X^1,\cdots ,X^{k-1})$$
​ 其中， X k = { x i , x i + 1 , ⋯   , x j } X^k=\{x^i,x^{i+1},\cdots,x^j\} Xk={xi,xi+1,⋯,xj}，表示第k步要预测的token集合。见图

​

8 结束

​ 好了，本期内容到此为止，如有问题，还望指出，阿里嘎多！