图像 Token 化：视觉数据转换的关键技术

type

status

date

slug

summary

VQ-VAE

Understanding VQ-VAE (DALL-E Explained Pt. 1)

Code: https://github.com/apple/ml-4m/blob/main/fourm/vq/vqvae.py

VAE使用连续向量表示隐状态，而 VQ-VAE (VQ: vector quantization) 使用离散方式表示隐状态。

思路其实也挺简单的，encoder的输出还是连续向量，但会维护一个大小为的 codebook （连续向量集合），然后把这个连续向量对应到codebook中离它距离最近的向量，对应的序号就变成了离散的隐状态了。传给 decoder 的输入就是向量。

除了模型参数外，codebook中的向量集合也是通过训练获得的。

如果一个输入的图片只用codebook中的一个向量表示的话，表达能力一般是不够的（decoder 只能看到个向量）。但可以让一个输入输出一个向量序列，表达能力就够了。比如把一张图片按横竖线划分为 32 x 32 个网格，然后每个网格对应codebook中的一个向量，这个表达能力就足够了 ()。

训练 loss：

其中表示 stop gradient ，即不进行求导。

上面第二项是期望调整codebook以尽量与encoder输出接近，第三项（codebook commitment loss）是期望调整encoder输出以尽量与codebook接近。如果输出一个序列的向量，最后两项会对序列所有向量的结果进行平均。

先验分布

训练的时候是假设先验分布是均匀分布而忽略的。但实际推断使用时，是可以基于应用数据获得先验分布的。比如有一堆应用数据，可以通过VQ-VAE计算获得它们对应的离散表示，其中每个是一个离散序列，类似文本中的一句话。有了这些离散序列，就可以类似语言模型一样，训练一个语言模型，然后用它来预测一个离散序列的概率值。

训练得到先验分布后，就可以先从先验分布中抽取得到，然后再通过 VQ-VAE 的 decoder 生成图片。

<ins/>

discrete VAE (dVAE)

How is it so good ? (DALL-E Explained Pt. 2)

Code: https://github.com/lucidrains/DALLE-pytorch/blob/main/dalle_pytorch/dalle_pytorch.py#L101

VQ-VAE 每次只对应到 Codebook 中的一个元素，而 dVAE 则每次计算与 Codebook 各元素的对应分布），然后从此分布中抽取一个 Codebook 元素作为 Decoder 的输入。

下面具体介绍每个部分：

Encoder

常见的 Encoder 就是 CNN，它把维度为的原始图片，转换成维度为的输出，其中的为 Codebook 中向量数量。其中为高度维度上 patch 化后的 patch 数量。比如原始图片高度，每个 patch 高度为，那。类似，是宽度维度上 patch 化后的 patch 数量。

当然，现在较大模型 Encoder 经常使用 CNN + Transformer 的结构，比如 ViT 。

Gumbel-Softmax 离散化

Encoder 输出结果的维度为，也就是每个 patch 对应一个长度为的向量。中的每个元素对应着 Codebook 中的各个向量，例如第个元素的值可以理解成是选择 Codebook 中第个向量的 logit，做下 softmax 变换后就可以认为是选择 Codebook 中第个向量的概率值。

那怎么从这个包含个类别的多项式分布中抽取出类别（Codebook 中哪个向量）呢。在推断时当然好做，但是在训练时要保证可微性就不能直接抽样。这时候就用到了一种叫 Gumbel-Softmax Trick 的方法。这里不细说 Gumbel-Softmax Trick，只要知道用了它之后我们在训练时也可以做到从多项式分布中抽样（或近似抽样）且不破坏训练过程的可微性。抽样过程包括两种模式：

hard 模式：这种模式得到的分布是取值为 0 或 1，抽中的那个类值为 1，其他类为 0。比如选中第 2 个类，那通过 Gumbel Softmax 处理后得到的向量是。

soft 模式：这种模式类似带 temperature 的 softmax 操作，它让之前概率高的类概率值更高，之前概率低的类概率值更低，可以认为是 hard 模式的一种 soft 近似。比如之前第 2 个类的概率值就很高，那通过 Gumbel Softmax 处理后得到的向量可能是。

⚠️ 注：2 中上下两个分布对应不同 patch 的输出分布 — **⚠️ 注：**2 中上下两个分布对应不同 patch 的输出分布

所以维度为的 Encoder 输出，通过 Gumbel-Softmax 离散化后，其维度仍是，只是在最后一个维度上的含义从 logits 变成了概率值。这个结果再与 Codebook 矩阵相乘就得到了最终的抽样向量，维度为：

代码示例：

Decoder

Decoder 通常就是转置 CNN （transposed CNN），它把输入维度为的张量恢复成与原始图片同样大小，即大小的张量。

转置 CNN 原理如下图，具体可见：conv_arithmetic/README.md at master · vdumoulin/conv_arithmetic 。

当然，Decoder 也可以使用 Transformer 架构。

训练 Loss

dVAE 的训练 loss 和 VQ-VAE 类似，只是使用了KL距离来让分布尽量分散：

其中利用 Gumbel-Softmax 从中抽样得到，是个等概率的多项式分布。

上面第一项 loss 为 L1，L2 或者 Smooth L1（靠近 0 为 L2，否则为 L1）。

第二项 loss 是 Encoder 输出的 logits 和 uniform 分布的 KL 距离，示例代码如下。

<ins/>

VQGAN

Taming Transformers for High-Resolution Image Synthesis

Code: https://github.com/CompVis/taming-transformers/tree/master#training-on-custom-data

图中 Transformer 与 VQGAN 训练过程无关，是指 VQGAN 训练完后可以把离散 tokens 作为 Transformer 模型输入训练图片生成模型。

VQ-VAE 中的损失函数是：

上面第一项为重构 Loss。

而 VQGAN 则把上式中的替换为感知（Perceptual）Loss，同时加入了生成 patch 的判别器 loss：

📌

看官方实现，重构 Loss 其实同时包含了 L1 Loss 和 Perceptual Loss 😅。效果怎么好就怎么来吧。

所以 VQGAN 的训练 Loss 为：

通过以下方式自适应调整，其中表示输入的梯度值：

感知损失（Perceptual Loss）

Perceptual Loss（感知损失）论文笔记 - stardsd

VQGAN 中实现，参考：
lpips.py
CompVis

感知损失（Perceptual Loss）是一种用于衡量生成图像质量的损失函数。它与传统的像素级损失（如均方误差，MSE）不同，更关注图像的高层次特征和感知质量。感知损失通常基于预训练的卷积神经网络（CNN），如VGG网络，用来捕捉图像的内容和样式特征。

感知损失的工作原理：

预训练模型的使用：感知损失利用了在大规模数据集上预训练的卷积神经网络（例如VGG-16）。这些网络在分类任务上已经学到了丰富的图像特征表示。

特征提取：将原始图像和生成图像分别通过预训练模型，提取出不同层次的特征图。这些特征图包含了图像的高层次语义信息。

计算损失：在特征空间中计算原始图像和生成图像之间的差异。感知损失通常由内容损失和样式损失组成：

，其中和是权重系数，用于平衡内容损失和样式损失。

内容损失：主要关注图像的高层语义信息和全局结构，主要作用于较高层。在特定层次上比较原始图像和生成图像的特征图，衡量它们的相似度。

样式损失：关注图像的纹理、颜色等风格特征，通常作用于从低到高的多个特征层。通过计算特征图的 Gram 矩阵，比较原始图像和生成图像的样式。

其中，是第层的权重，用来平衡不同层次的样式损失。

具体的计算过程：

假设我们有一个经过预训练卷积神经网络的特征图（activation map），其维度为，其中是通道数（channels）。Gram 矩阵的计算步骤如下：

展平特征图：将特征图展平成一个行，列的二维矩阵。

其中表示通道，表示展平后的空间位置。

计算Gram矩阵：Gram矩阵是通过特征图矩阵的自身转置相乘得到的。

Gram 矩阵的每个元素表示特征图中第和第个通道之间的相关性。

样式损失是通过比较原始图像和生成图像的 Gram 矩阵来计算的。假设原始图像的特征图的 Gram 矩阵为，生成图像的特征图的 Gram 矩阵为，样式损失可以定义为：

其中，是第层的权重，用来平衡不同层次的样式损失。

感知损失在许多计算机视觉任务中都有广泛应用，包括但不限于：

图像超分辨率：通过感知损失，生成的高分辨率图像可以更好地保持原始图像的细节和感知质量。

图像风格迁移：在图像风格迁移（style transfer）中，通常会同时使用内容损失和样式损失。内容损失确保生成图像与原始内容图像在内容上相似，而样式损失则确保生成图像与目标样式图像在样式上相似。通过在不同层次上衡量内容和样式的相似度，可以生成具有特定艺术风格的图像。

图像生成和重建：在生成对抗网络（GANs）中使用感知损失，可以提升生成图像的真实感和视觉质量。

感知损失通过关注图像的高层次特征和感知质量，弥补了传统像素级损失在图像生成任务中的不足，使生成图像更加逼真和自然。

<ins/>

TiTok 1D Tokenizer

An Image is Worth 32 Tokens for Reconstruction and Generation

Code:
1d-tokenizer
bytedance • Updated Jul 14, 2024

前面介绍的 VQ 类 Tokenizer，一张图片表达后的 token 长度，是跟图片被划分成的 patch 数量一致的。比如的图片，如果 patch 大小为，那 patch 数量也是。如果是的图片呢，那 patch 数量就变成了。所以分辨率高的图片，要表达它所需的 token 数量就越多。另一个问题是，图片中的临近 patch 可能内容很相近，使用单独的 token 分别表达它们会存在冗余，有点浪费。

为了解决上面的问题，字节给出了一种思路：使用固定长度的 tokens 来表达任意分辨率的图片，而且其中的每个 token 都能看到图片的 global 信息，而不是之前的单个 patch 中的 local 信息。模型的名称叫：Transformer-based 1-Dimensional Tokenizer (TiTok)，这名字取得。。

TiTok 依旧是 VQ 类 Tokenizer，它依旧使用 Encoder → VQ → Decoder 的流程。以下是整体结构：

接下来分开细说。

Encoder

Encoder 的主体架构是 BERT。它的输入前一部分是图片 patches。和 ViT 一样，这些 patches 经过 CNN 获得了输入的 embedding 向量。BERT 输入的第二部分是固定长度为的 Latent Tokens。这里的就是表示一张图片的固定 token 数量，如、等。这些 Latent Tokens 对应的 embedding 向量是随机初始化的，然后在训练过程中学习得到。Latent Tokens 和图片 patches 的 embedding 向量一起作为 BERT 的输入。

对于 BERT 的输出，只使用 Latent Tokens 对应的那些输出向量。这个输出向量会作为 VQ 的输入。

Vector Quantizer (VQ)

VQ 部分是为了将 Encoder 输出的个连续向量转换为个离散的 tokens。VQ 的过程是通过比较 Encoder 输出的向量与 Codebook 中的向量，根据距离选择最接近的向量代表原始向量。这个部分和之前介绍的 Tokenizer 一样。

Decoder

Decoder 的主体架构依旧是 BERT。此时 BERT 的输入也包括两部分。

第一部分是从 VQ 中得到的个离散 tokens，它们对应的 embedding 向量就是 Codebook 中对应的向量，只是会利用线性变换转换为 Decoder 需要的输入长度。第二部分是填充的重复 mask token。mask token 只有一个，它对应的 embedding 通过学习得到。这个 mask token 会被重复次，即与 patch 数量相同，追加到第一部分的个 tokens 后面。

BERT 的输出只使用 mask tokens 对应的那些输出，这个输出会被 reshape 为，然后通过一个上采样 CNN 恢复成原始图片。代码中是通过插值获得的上采样，而不是使用转置 CNN。

整个训练过程论文讲的比较简单，相关代码也没开源。训练分成 2 个步骤，作者称之为 Warm-up 和 Decoder fine-tuning。作者实验表明这种训练方法可以改善训练过程的稳定性和最终重构的图片质量。

Warm-up：Codebook 直接从训练好 MaskGIT-VQGAN 模型中拿过来用，这一步训练模型中的 Encoder，Quantizer，Decoder。

Decoder fine-tuning：固定 Encoder 和 Quantizer，只精调 Decoder。具体方法可参考 VQGAN 的训练代码。通过这一步可以改进图片重构质量。

和 VQGAN 可以结合 MaskGIT 做 generative 一样，TiTok 也可以。作者的实验表明 TiTok + MaskGIT 效果更佳。。

重复下 TiTok 的好处：

转换后的 token 长度与图片切分后的 patch 数量无关。

图片转换后的 token 长度与图片分辨率无关，可以是设定好的固定值。

图片中的临近 patch 可能内容很相近，使用单独的 token 分别表达它们存在冗余，有点浪费。

<ins/>

OmniTokenizer

OmniTokenizer

Encoder Code: https://github.com/FoundationVision/OmniTokenizer/blob/main/OmniTokenizer/omnitokenizer.py#L778

OmniTokenizer 是近期复旦的工作，它能同时 token 化图片和视频。OmniTokenizer 依旧是 VQVAE 架构，后者更准确说是 VQGAN 架构。它的 VQ 部分参考 VQGAN 就行。

OmniTokenizer 的 Encoder 和 Decoder 都是使用了魔改过的 Transformer 架构。以下是 Encoder 结构。

Encoder

前面的处理都是类似，图片就是 patch 分块，然后使用 CNN 或者其他 patch embed 方法变成序列向量。而视频则会划分为 3 维（带时间维度）的 patches，但其中的首帧会拿出来单独处理。首帧 patches 和 3 维 patches 会利用各自的 patch embed 模块变成相同长度的序列。把时间维度拉平就能拼接到一块了。图片看作是只有一帧的视频。

Transformer 里面包含两种模块层，一种是让各个 patch 对应序列在空间维度上进行 self-attention，称之为 Spatial Transformer Layers。OmniTokenizer 的 Spatial Layer 使用了 Swin 中的 Window Attention，只考虑空间维度上附近的一些 patches 对应序列，如下图。