（2024，ViT，小波变换，图像标记器，稀疏张量）基于小波的 ViT 图像标记器

Wavelet-Based Image Tokenizer for Vision Transformers

公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

0 摘要

1 引言

3 基于小波的图像压缩简介

4 图像标记器

4.1 像素空间标记嵌入

4.2 语义标记嵌入

4.3 Transformer 层的操作计数

5 实验结果

5.1 设置

5.2 标记器比较

5.3 标记嵌入大小的影响

6 结论与未来工作

0 摘要

非重叠逐补丁卷积（Non-overlapping patch-wise convolution）是所有最先进的视觉 Transformer（ViT）模型的默认图像标记器（tokenizer）。尽管已经提出了许多 ViT 变体来提高其效率和准确性，但在改进图像标记器本身的研究文献中却很少见。在本文中，我们提出了一种基于小波变换的新图像标记器。我们展示了采用新标记器的 ViT 模型在 ImageNet 验证集上既实现了更高的训练吞吐量，又获得了更好的 top-1 精度。我们对为什么所提出的标记器在不改变 ViT 模型架构的情况下提高了训练吞吐量进行了理论分析。我们的分析表明，新标记器能够有效处理高分辨率图像，并且天然对抗对抗性攻击。此外，所提出的图像标记器为基于 ViT 模型设计的重要新研究方向提供了新的视角，例如基于非均匀网格的图像标记以进行图像理解。

1 引言

自然语言固有地使用固定字母表中的一组离散字符或字节（bytes）。字节序列可以直接作为 Transformer 模型的输入，如在 ByT5 中所做的那样。然而，这样的字节序列的潜在表示可能不携带足够的信息，且序列长度可能不必要地长。更常见的替代方法是将单词分割成显性或隐性的子词标记（token）。显性子词标记器（tokenizer）已成为许多最先进的大型语言模型（如 T5、GPT-3、ULM、PaLM 和 Chinchilla）的默认组件。

图像像素也是离散的，因为 RGB 通道由 8 位无符号整数表示。这意味着原始像素词汇量为 224，这个词汇量过大以至于不实际。此外，即使是中等分辨率的 256 × 256 图像的序列长度也太长。视觉 Transformer（ViT）常用的图像标记器是非重叠的补丁卷积。尽管 ViT 有许多变体，但基于补丁卷积的标记器仍然是最先进的 ViT 模型的默认选择。由于 ViT 模型是 CLIP、ALIGN、CoCa 和 PaLI 等大型视觉-语言模型的图像编码器，基于补丁卷积的图像标记器也成为这些视觉-语言模型的默认组件。

图像标记器执行两个主要任务：1）将给定图像划分为非重叠的补丁，使最终的标记数量比图像像素数量小得多；2）将每个补丁中像素的 RGB 通道映射到标记嵌入向量中。这将一组图像补丁转换为输入标记序列。图像标记器不可避免地会丢失与每个图像标记相关的像素之间的相对位置信息。因此，一个好的图像标记器必须在位置信息丢失和计算节省之间找到平衡，这两者都是由于标记数量减少而导致的。同时，不同的视觉任务对这两个相互冲突的方面有不同的要求。对于高分辨率图像的数据集，标记数量由于标准 Transformer 注意力的二次复杂性而容易成为瓶颈。对于语义分割和目标检测等视觉任务，相对位置或空间信息起着关键作用。我们认为，一个好的标记器应该有效地压缩补丁（patch）像素中的冗余信息。这使得能够灵活地适应这两个相互冲突的要求，并激发我们从图像压缩领域寻找灵感。

图像压缩领域有许多已建立且理解良好的算法。图像压缩的中心目标是去除像素空间中的冗余信息，使得重构后的图像在视觉上接近原始图像。目前行业标准 JPEG2000 中使用的流行算法是小波变换。

在本文中，我们从 JPEG2000 中的关键见解中汲取灵感：通过小波变换揭示的大多数高频信号可以在视觉上无显著差异地被安全截断。这是大部分图像压缩发生的地方。简而言之，我们建议用基于小波的图像标记器替代标准 ViT 架构中的基于补丁卷积的图像标记器。

3 基于小波的图像压缩简介

我们首先以卷积核的形式描述离散小波变换。我们建议读者参考优秀的书籍，如 [15, 48]，以全面了解小波变换理论。

一个分辨率为 N × N 且具有 RGB 通道的彩色图像可以表示为一个张量 P ∈ R^(N×N×3)。在图像压缩中，通常会通过一个固定且可逆的线性投影将 RGB 格式转换为亮度-色度格式 [23]。这可以表示为另一个张量 [Y,Cb,Cr] ∈ R^(N×N×3)，其中 Y ∈ R^(N×N) 是亮度矩阵，Cb 和 Cr ∈ R^(N×N) 分别是色度-蓝和色度-红矩阵。众所周知，人类视觉感知对亮度或亮度（brightness or luminance）更敏感，而对颜色或色度的敏感度较低 [23]。因此，通常在应用小波变换之前，将 Cb 和 Cr 的原始图像分辨率按 2 × 2 的比例减少。这将总体像素数减少了一半。

小波有许多家族。这里我们关注两种在图像标记器中有效的小波家族：Daubechies (db) 和 Coiflets (coif)。每个家族中可以选择不同的阶数。例如，db 和 coif 家族中的最低阶一维 (1D) 核分别是 db1（也称为 Haar）核：

coif1 内核：

二维 (2D) 核由一维核衍生而来，如下所示

使用这四个二维卷积核对 Y、Cb 和 Cr 分别进行卷积。这也称为组卷积（group convolution）。例如，对矩阵 Y 进行卷积得到

其中，Y^(0)=Y, conv_2d (M) 是使用核 M 和固定步长 (2, 2) 的二维卷积。方程（4）-（5）中的相同卷积可以递归地在矩阵 Y^(k) 上执行，k=1,2,…。步长 (2, 2) 决定了每次递归步骤将图像分辨率按因子 2 × 2 减少。矩阵

可以合并成一个矩阵。例如，对矩阵 Y 进行简单的二级小波变换得到

方程（4）-（6）中的过程同样适用于矩阵 Cb 和 Cr，并且我们可以得到与方程（6）中相同的矩阵结构。

图 1 显示了使用方程 (1) 中的 db1 卷积核对具有 YCbCr 通道的图像进行的 10 级小波变换。请注意，大多数小波系数几乎为零，这在通道 Cb 和通道 Cr 中特别明显。这种分布是自然图像的典型特征，它为我们提供了足够的空间，可以安全地稀疏化或将这些接近于零的小波系数置零，而对重建图像的质量影响最小。所有非零系数都经过量化，然后使用熵编码算法进一步压缩。有关此步骤的详细信息，请参阅 [1]。

我们可以轻松逆转方程 (4)-(6) 中的步骤并重建原始图像。主要区别在于，我们需要在方程 (4)-(5) 中使用转置卷积。

4 图像标记器

在本节中，我们将离开标准 JPEG2000 算法背后的图像压缩概念，专注于设计一个有效的图像标记器（tokenizer）。共同的目标是去除图像中的冗余信息，因此小波变换是一个共同的组成部分。然而，最终目标是不同的。在图像压缩算法中，每一步都必须是可逆的，以便能够从 JPEG 字符串中重建原始图像，且视觉感知质量损失最小。另一方面，图像标记器的设计旨在捕捉语义上有意义的信息。生成的标记嵌入是 ViT 模型后续处理的起点，即注意力层中的 “回音室”（echo chambers）效应 [25] 和前馈层中的键值对记忆 [22]。

4.1 像素空间标记嵌入

与 JPEG2000 相同，我们采用 R-G-B 到 Y-Cb-Cr 的线性变换，并将 Cb 和 Cr 通道按 2 × 2 的比例进行下采样。因此，对于分辨率为 N × N 的图像，小波变换的输入是形状为 Y ∈ R^(N×N) 和 Cb, Cr ∈ R^(N/2 × N/2) 的 YCbCr 通道。

与傅里叶或离散余弦变换不同，小波变换使用局部化的卷积核，见方程 (1) 和 (2)。因此，每个小波系数仅与一个局部像素区域相关。我们可以通过一个具体的例子来说明这一点。给定一个分辨率为 128 × 128 的彩色图像，我们按照方程 (4)-(5) 中的步骤对 YCbCr 通道进行二级小波变换。我们可以将方程 (6) 中的结果矩阵及其对应的 Cb 和 Cr 重新写为

注意，我们在分辨率为 128 × 128 时跳过了对 Cb 和 Cr 的卷积，因为它们已经被下采样到 64 × 64。这就是为什么在方程 (9) 中没有 Cb 或 Cr 成分。还要注意，P^(2) 只是一个分辨率为 32 × 32 的 YCbCr 通道的粗化图像。我们可以重新整形张量 D^(1)，使 D^(1) ∈ R^(32×32×12)，并将P^(2)、D^(2) 和 D^(1) 沿其最后一个维度连接成一个张量：

更一般地，对于分辨率为 N×N 的图像进行 L 级小波变换，结果是一个沿最后一个维度连接的张量集合：

方程 (11) 中向量 W(i, j, :) 的条目数量为：

我们将长度为 C 的向量 W(i,j,:) 视为分辨率为 ~N × ~N 的粗略图像中像素 (i,j) 的像素空间嵌入。对于方程 (10) 中的简单二级示例，我们只需要包含 24 个条目的向量 W(i,j,:) 就可以重建对应于粗略图像 P^(2) 中像素 (i,j) 的 4×4 原始像素区域。这可以通过按照相反顺序执行方程 (4)-(10) 中的步骤来完成。同样的程序也适用于方程 (11) 中的通用情况。如果没有截断，即没有将小的小波系数归零，则此重建是精确的。如图 1 所示，可以轻松通过阈值截断来实现截断。JPEG2000 已经证明，即使将 90% 以上的小波系数归零，重建误差也很小。因此，从这一点开始，像素空间标记嵌入总是指向量 W(i,j,:) 的截断后版本。

几点备注如下：

像素空间嵌入张量 W 中的标记（token）数量为 ~N^2 = N^2 / 4^L。因此，级别 L 作为超参数用于实现标记数量减少和位置信息损失之间的权衡，如第 1 节所讨论的。
小波卷积核是固定的，如方程 (1) 和 (2) 所示。因此，任何基于梯度的对抗样本 [24] 都会失败，因为相对于这些卷积核没有梯度。一般来说，任何对输入图像的小白噪声扰动都会变成高频小波系数 [3, 52]，在稀疏化步骤中将这些系数归零。因此，基于小波的标记器自然对图像分类器的对抗攻击具有抵抗力（例如，[44]）。
小波变换暴露出具有平滑 RGB 像素的图像区域。每个像素空间嵌入向量中的稀疏性与对应像素区域的信息熵强相关。这可以作为设计和实现图像标记器的有用信号。例如，它可以用于引导非均匀图像分区，并在非均匀网格上生成图像标记（非均匀网格是偏微分方程数值方法中的标准技术【45, 35】。它是处理变量（如电势或气流压力）作为空间坐标函数时的大变化的不可或缺的工具。自然图像中的 RGB 或 YCbCr 强度也随着像素坐标的变化而显著变化）。这与 [61, 41] 中的目标类似，但小波变换提供了一个更有原则的框架。这涉及对 ViT 模型本身的修改，因此不在本文讨论范围之内。

4.2 语义标记嵌入

这一步的目标是将像素空间嵌入映射到语义上有意义的标记嵌入。这个映射可以是一个简单的可训练线性投影：

其中，E ∈ R^(~N × ~N × H) 是语义标记嵌入张量，H 是目标嵌入大小，Q ∈ R^(C×H) 是可训练的投影矩阵，W ∈ R^(~N × ~N × C) 定义在方程 (11) 中。

方程 (13) 表明像素空间嵌入大小随着小波级别 L 的增加呈指数增长。然而，方程 (11) 中的每个向量 W(i,j,:) 是极其稀疏的，并且具有可预测的稀疏图样（pattern）。对于图 1 中显示的 10 级小波，几乎所有 W^(1) 中的条目都是零，而 W^(L) 几乎没有零条目。一般来说，随着级别索引 k 的减少，W^(k) 中零的比例迅速增长，可能是指数级的。我们的实验（表 1-2）显示，将 W 中的 80% 条目稀疏化实际上提高了图像分类的 top-1 精度。我们认为原因是这些小的小波系数包含的噪声多于信号，过滤掉它们可以作为模型训练的一种正则化形式。

我们可以在矩阵 Q 中使用块稀疏结构来利用这种稀疏模式。新的投影是：

其中 C_k 在方程 (12) 中定义，H_k 是分配给级别 k 的标记嵌入大小，满足约束条件：

根据方程 (11)，我们有：

关键问题是：给定方程 (12) 中的 C_k，H_k 的合适值是什么？这显然取决于向量W^(k) (i, j, :) 的稀疏性。对于较高的小波级别，如 k = L 或 k = L-1，张量 W^(k) 是密集的，因此我们可以简单地设置 H_k = C_k。然而，像 Q^(1) 这样较大的投影矩阵对参数数量和运行时间有显著影响。幸运的是，张量 W^(1) 是非常稀疏的，这为找到远小于 C_k 的 H_k 提供了良好的机会。

为了便于说明，我们专注于计算向量 E^(k) (i, j, :) ：

或等效地：

注意，向量 w_ij 中非零条目的位置是随机的，取决于正在处理的图像。我们通过两个步骤表明，我们可以将 H_k 设置为矩阵 Q^(k) 的预期数值秩，这对于较小的 k 来说远小于 C_k。第一步可以形式化为以下命题：

命题 1：对于线性投影：

且 x 中只有一个随机子集的条目是非零的。一个条目在 x 中为非零的概率是 p。假设我们对这样一个伯努利分布进行采样以生成一个随机 x 并将生成的 y 累积到一个矩阵

那么矩阵 B 的期望列秩是 pN。

证明：矩阵 A 的列秩为 N，但每个 y_i 的计算仅使用其中的 pN 个。因此，矩阵 B 的列只能跨越概率上覆盖矩阵 A 中 pN 列的空间。这是矩阵 B 的预期列秩。

假设我们将方程 (20) 中的所有嵌入向量 e_ij 收集到一个单独的矩阵 S 中。让总的标记（token）数为 S_t，因此 S ∈ R^(Hk × S_t)。注意 St 可能是一个很大的数字，因为矩阵 S 包括从训练数据集中提取的所有图像的张量 W^(k) 中的每个级别-k 的像素空间嵌入向量。显然，Hk ≪ St。第二步可以陈述为以下命题。

命题 2：矩阵 S 的期望数值秩是 min(H_k，p·C_k)。

证明：根据方程 (20) 和 (21)，我们可以看到 e_ij 对应于 y，w_ij 对应于 x。同样，矩阵 S 对应于方程 (22) 中的 B。由于命题 1 的直接结果，矩阵 S 的预期列秩为 p·C_k。因此，rank(S) = min(H_k，p·C_k)。

由于我们可以灵活地设置 Hk，以便在方程 (20) 中的投影后保留 w_ij 中的信息在 eij 中，我们可以简单地设置 Hk = p·C_k。即使方程 (12) 中的 Ck 对于小的 k（如 k = 1）可能很大，但根据经验，我们发现概率 p 对于小的 k 非常小。这意味着 Hk = p·C_k 在小的 k 上几乎是恒定的。对于大的k，例如 k = L，我们设置 Hk = Ck，因为像素空间嵌入是完全密集的，或者等效地 p = 1。根据方程 (16) 的观察，语义 token 嵌入大小 H 接近于一个仿射函数，而不是小波级别 L 中的指数函数。

考虑一个分辨率为 N ×N 的图像，基于补丁（patch）卷积的标记器（tokenizer）以补丁大小 p × p 生成 T = N^2 / p^2个 token。当 p 太大时，补丁卷积在提取视觉信息时变得无效，这就是为什么 p = 8、16、32 是常见选择的原因。因此，标准 ViT 与补丁卷积标记器的操作次数增长为 O(T^2) = O(N^4 / p^4)。当 N > 2048 时，这变得非常昂贵。然而，这种高分辨率图像可以通过基于小波的标记器降低到用户指定的分辨率 ˜N = N / 2^L，其中 L = log2(N / ˜N)。由于方程 (16) 中的语义 token 嵌入大小 H 是 L 的仿射函数，正如上面讨论的，我们有 H = O(log2(N / ˜N))，随着 N 的增长缓慢增长。因此，基于小波的标记器提供了一种更有效的处理高分辨率图像的方法，因为方程 (15) 中的块稀疏投影。我们可以在方程 (15) 中的单个线性层之上添加更多的密集层，以潜在地提高 token 嵌入质量。由于添加层导致的参数数量和训练成本的增加为 O(H2)。这比方程 (15) 中的 O(CH) 要小得多，因为嵌入维度已经从方程 (13) 中的大值 C 减少到方程 (16) 中的小值 H。

4.3 Transformer 层的操作计数

标准 Transformer 由许多相同的层组成，具有以下超参数：标记嵌入大小 H，模型隐藏状态大小（通常与 H 相同），注意力查询、键和值特征大小 d_model（通常也与 H 相同），以及前馈中间层激活大小 d_ff = m×H（其中整数 m 通常在范围 [2, 10] 中）。与之前一样，令 T 为标记数。每个 Transformer 层的操作计数的分解如下：

注意力：2·T^2·H + 4·T·H^2
- 查询/键/值多头投影和输出投影：4·T·H^2
- 查询键内积和注意力值乘法：2·T^2·H
前馈：2·m·T·H^2

我们将注意力放在两个二次项上：2·T^2·H 和 (2m + 4)·T·H^2。如果我们使用典型的 m=4，那么它们之间的比率为

许多高效 Transformer 研究的重点是改善 Transformer 注意力中的 O(T^2) 二次复杂度 [49, 68, 42, 9, 58, 29, 64, 2, 4]。这对于具有长上下文的任务（如长文档摘要和长视频理解和生成）非常重要，因为 T≫H，因此在方程 (23) 中 r≫1。然而，对于许多标准的语言或视觉模型和数据集来说，另一个二次项 O(H^2) 占主导地位。例如，在 T5-large 模型 [39] 中使用 T=512,H=1024，在 SOTA 视觉语言模型 PaLI 中使用 T=256,H=4096。在这种情况下，在方程 (23) 中 r≪1，因此通过使用第 4.2 节中的基于小波的图像标记器来降低 H 更具成本效益。

值得注意的是，最近提出的 Flash Attention 中的矩阵切片和操作融合思想实质上消除了 Transformer 模型的二次内存复杂度。因此，我们在分析中不包括内存使用情况。这里的讨论仅关注模型训练中的前向路径。反向梯度传播路径的分析遵循相同的推理，相同的结论适用。

5 实验结果

5.1 设置

数据集。我们使用 ILSVRC-2012 ImageNet 数据集，该数据集包含 1k 类和 130 万张图像【43】。图像分辨率有两种设置：256 × 256 和 512 × 512。

ViT模型。我们的基准模型是【17】中的 ViT-base 模型，使用相同的模型参数：H = d_model = 768, d_ff = 3072, 12 层 Transformer 层, 12 个注意力头。我们将 d_model 和 H 视为两个独立的超参数，因此可以分别检查它们对训练吞吐量和 top-1 精度的影响。所有模型在 TPUs v4【19, 28】上进行训练，使用【17】中报告的相同训练超参数，例如批量大小、训练轮数、学习率、权重衰减等。我们遵循【47】中的数据增强方法，结合了 Mixup【66】和 RandAugment【12】。下面的结果清楚地表明，基于小波的标记器不会干扰这些数据增强技术。我们要强调的是，在我们用基于小波的标记器（tokenizer）替换补丁（patch）卷积图像标记器时，ViT 模型架构和训练方法保持不变。这样，改进可以直接归因于小波分词器。

补丁卷积标记器的变体仅由补丁大小控制。我们在表 1 和表 2 中使用 patch/8, patch/16 和 patch/32 的表示法来表示这些变体。例如，patch/8 表示 8 × 8 像素区域对应一个 ViT 输入 token。

基于小波的标记器变体由两组参数控制：

小波变换参数包括小波类型和级别 L，以及通过阈值零化的百分比小波系数。我们使用 db1-4 和 db1-5 的表示法分别表示级别 4 和级别 5 的小波变换，db1 核定义在方程 (1) 中。同样， coif1-3 和 coif1-4 分别表示级别 3 和级别 4 的小波变换，coif1 核定义在方程 (2) 中。默认情况下，方程 (11) 中 W 的 80% 的小波系数被零化。这个百分比适用于分辨率为 256 × 256 和 512 × 512 的图像（对于更高分辨率的图像，这个百分比应该更高）。
密集投影参数包括方程 (16) 中的 Hk 和额外密集层的数量。如第 4.2 节所述，对于 k = L, L−1，我们设置 Hk = Ck 。我们对 H1 进行超参数搜索，以在方程 (16) 中 H 是 128 的整数倍的约束下优化 top-1 精度。两个表现最好的选择（H = 256, 384）及其影响见表 3。我们的实验表明，增加一个额外的密集层不会带来 top-1 精度的提升。这可能是因为对于 512 × 512 的图像分辨率，L 最多为 5。我们认为，对于分辨率高于 1024 × 1024 的情况，需要更多的密集层。

输入序列长度。需要注意的是，每个基于小波的 token 对应于原始图像中大小为 2L × 2L 的像素区域。因此，使用 coif1-3 和 patch/8 从分辨率为 256 × 256 的图像生成的输入序列长度相同，为(256 ÷ 8)² = 32² = 1024。在所有结果表中，这个序列长度被表示为 tokens。

5.2 标记器比较

我们进行了实验，比较基于补丁卷积的标记器（表示为 patch/N，其中 N 是补丁大小）和基于小波的标记器（db 和 coif）。我们在分辨率为 256×256（表 1）和 512×512（表 2）的图像上报告结果。训练使用 TPUv4 加速器进行，吞吐量以每个 TPU 核心每秒处理的图像数量来衡量。结果显示，基于小波的标记器在 top-1 精度和训练吞吐量上都表现更好。值得注意的是，由于方程 (15) 中的块稀疏投影，基于小波的标记器也显著减少了模型参数数量，因为 H 减小了。

5.3 标记嵌入大小的影响

我们在表 3 中报告了使用不同 H、d_ff 和 d_model 的结果。较小的 H 允许我们进行一些有用的权衡。例如，配置 H = 384, d_model = 768 意味着多头投影大小减少了一半，前馈层大小减少了 3/4。这种节省使我们能够将 d_ff 从 4 × H 扩大到最多 16 × H。这种灵活性允许在精度和训练吞吐量之间进行权衡，如表 3 所示。

6 结论与未来工作

我们描述了一种新的基于小波的图像标记器，作为标准 ViT 模型中使用的补丁卷积标记器的替代。我们证明了装备新标记器的 ViT 模型在 ImageNet 验证集上实现了更高的训练吞吐量和更好的 top-1 精度。我们还提出了一个理论分析，解释了为什么所提出的标记器在不改变 ViT 模型架构的情况下提高了训练吞吐量。我们的分析还表明，新标记器可以有效处理高分辨率图像，并且自然抵抗对抗攻击。总体而言，我们的图像标记器为基于 ViT 模型设计的重要新研究方向提供了新的视角，例如非均匀网格上的图像标记和处理图像分辨率提高的新方法。