（论文阅读）An Image is Worth 32 Tokens for Reconstruction and Generation

1. 论文

1. Yu et al. An Image is Worth 32 Tokens for Reconstruction and Generation. NeurIPS, 2024.

摘要：

在视觉中，尤其是高分辨率设置下，tokenization 将图像转换为隐层表示以降低计算需求。然后，现在的方法（VQ）通常将隐层表示限制到一个潜在的2D网格上（并不是显式设计的）。这是因为每一个token通常都对应一个特定于它的local region（这是因为 这个token表示了这个区域内的信息，两者之间存在位置上的对应关系），此时tokens就被隐式的限制在了一个潜在2D网格上。

然而，在图像中，相邻的区域通常会有相似的信息，那么早先的工作无法克服这种冗余，因为tokenization只会专注于每一个local。为此，本文将图像表示为 1D 的tokens（去除早先的2D网格约束，每一个token都表示全局信息），来获得更加高效的表示。

---

2. 背景：

类似于NLP，为了能够提高训练效率且更有利于缩放模型大小，图像生成模型通常需要将输入转换到一个紧的隐空间。这借助了 tokenizer 和 re-tokenizer。比如，连续方法（VAE）和离散方法（VQ）。

然后现在的tokenizers会蕴含一个潜在的2D网格，例如，第一个token一般代表左上角的第一个patch。这可能会阻碍模型更好的压缩信息。

那么，潜在的2D网格真的会影响表现吗？实际上，对于高水平任务，比如图像分别，分割，目标识别，是不需要tokenizer的。它们会将图像压缩到1D序列，用来提取任务相关的高级信息。然而这个结论不能直接迁移，图像生成不仅仅需要高级信息，也需要细节来辅助生成。

为此本文就提出了一个基于transformer的框架，带有基于transformer的1D tokenizer：一个vit encoder，一个VQ，一个vit decoder。

---

3. 核心思想：

如何提取1D 的tokens呢？本文的方法其实很简单。

1. 在tokenization期间，首先将输入图像划分为patch \(P \in \mathbb{R}^{\frac{H}{f} \times \frac{w}{f} \times D}\)，\(f\) 表示下采样因子，\(D\) 表示embedding dimension。

2. 额外的，构建一个latent tokens \(L \in \mathbb{R}^{K\times D}\)。

3. 全部送入vit的encoder中：

   \[Z_{1D} = Enc(P \oplus C), \]

   特别的，输出仅仅保留latent tokens。

在de-tokenization阶段，我们在额外构建一个mask tokens \(M \in \mathbb{R}^{\frac{H}{f} \times \frac{w}{f} \times D}\)(与P相同大小），与 VQ 后的latent tokens拼接送入vit的解码器：

\[\hat{I} = Dec(Quant(Z_{1D}) \oplus M), \]

这就完成了重构。

---

4. 实验：

能够用更少的token获得更高的结果（重构or生成）。

---

5. 讨论：

remark 1：简单把2Dtokens展平并不等价1D，因为每个tokens中隐含的位置信息是固定的。而上述方法中，每一个token并不被限制到单个local，而是能够访问全局信息（通过SA)。

remark 2：我感觉有点像分类中的 cls token。