2406.11838

paper

Motivation: 现在为了将连续的图像数据匹配给离散的Transformer(Auto Regressive)模型, 需要使用Vector Quantization方法(如, VQ-VAE). 但是VQ Tokenizer训练困难, 对梯度近似敏感, 重建质量往往比连续空间的方法更差(如, VAE). 但是本文提出, 自回归本身是根据已有预测未知, 不依赖于数据是离散的还是连续的. 关键的缺失是一个能够在连续值空间中建模每个token概率分布的loss函数

因此, 提出一种新的方法, 在不使用VQ的情况下建模loss分布:

1. 提出Diffusion Loss. 用一个小型扩散过程来模拟连续值token的概率分布: $p(x|z)$, 使得Auto-Regressive模型能够直接预测连续的latent数据
2. 提出MAR框架, 融合 Diffusion Loss的Auto-Regressive 和 MAE 的方法, 结合Bidirectional Attention机制, 提高生成速度和质量

训练过程如下:

1. 将原始图片经过预训练的VAE, 得到一个latent vector(不需要使用quantize): $$X=\mathbf{Encoder}(I)\in {\mathbb{R}}^{h\times w\times d}$$
2. 对隐空间vector做随机掩码处理: $$X_{vis},X_{mask}=\text{Split}(X,\mathcal{M})$$其中,$\mathcal M$表示一个随机生成的掩码集合,与总token长度$n$的比例$\frac{|\mathcal M|}{n}$通常在$0.7\sim1$左右.分隔的方式为:$X_{vis}={x_i|i\not\in\mathcal M}$,$X_{mask}={x_j|j\in\mathcal M}$
3. 使用Transformer(bidirectional attention)进行观察并生成条件向量$z$, 这个向量总结了周围已知的图像信息, 为后续的denoise提供context, 指导去噪过程: $$z_i=\mathbf{Transformer}(X_{vis},{\text{pos}}_i)\text{for }i\in \mathcal{M}$$
4. 使用一个小的MLP网络作为去噪网络: $${ϵ}_{\theta }=\mathbf{MLP}(x_{i,t}|t,z_i)$$其中:
   • 噪声来源于高斯采样: $ϵ\sim \mathcal{N}(0,I)$
   • $x_{i,t}$表示diffusion加噪的结果: $x_{i,t}=\sqrt{{\bar{\alpha }}_t}+\sqrt{1-{\bar{\alpha }}_t}ϵ$
5. 计算重建损失, 或者叫做Diffusion Loss: $${\mathcal{L}}_{diff}={\sum }_{i\in \mathcal{M}}\mathbb{E}\left[\Vert ϵ-\mathbf{MLP}(x_{i,t}|t,z_i{)}^2\Vert \right]$$
6. Inference的时候, 使用这个公式进行去噪: $$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t|t,z)\right)+{\sigma }_t\delta$$其中:
   • $x_t$是上一步的噪声, 是一个更加混乱的状态. $x_{t-1}$是去噪之后的结果. 下标越小, 表示结果越接近真实值; 下标越大, 表示结果越接近噪声
   • ${ϵ}_{\theta }$是学习到的去噪网络(就是$\text{MLP}$)
   • ${\sigma }_t\delta$是一个随机加噪. ${\sigma }_t$是$t$时间步时的噪声标准差, $\delta$是一个从$\mathcal{N}(0,I)$里面采样得到的新的噪声. 使用这个是为了降低确定性(不要过拟合), 增强泛化能力
7. 同时, 使用MLP(output head, 输出头), 进行多token同时预测