当前位置: 首页 > news >正文

(论文阅读)An Image is Worth 32 Tokens for Reconstruction and Generation

(论文阅读)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), \]

这就完成了重构。

image-20251118111307766

4. 实验:

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

image-20251118111251729

5. 讨论:

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

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