Introduction

Transformer 已经革新了语言处理领域，但将其成功迁移到计算机视觉任务并非易事。Vision Transformer (ViT) 证明了可以将图像视为令牌序列并应用纯 Transformer 层来进行分类。然而，ViT 通常需要大规模的预训练数据集 (如 JFT-300M) ，才能在像 ImageNet 这样的中等规模数据集上达到与精调良好的卷积神经网络 (CNN) 相当的精度。

“Tokens-to-Token ViT” 的作者指出，从零训练的原始 ViT 存在两个主要问题:

  1. 缺乏局部结构建模: ViT 的初始分词 (tokenization) ，即对图像进行简单的基于网格的非重叠块切分，无法捕捉到诸如边缘、线条和纹理等丰富的局部结构。相比之下，CNN 通过卷积操作本质上能建模这些局部特征。
  2. 主干网络设计低效: ViT 主干通常表现为“浅而宽”的架构，这会导致通道冗余和容量未被充分利用，尤其在训练数据有限时更为明显。

他们提出的解决方案, T2T-ViT , 是一种优雅且务实的方法，重新思考了图像如何被转换为令牌以及 Transformer 主干的结构。其结果是一个能够在 ImageNet 上从零训练并有效收敛的模型，在参数和计算量更少的情况下，准确率优于可比的 ResNet 和现有的 ViT 变体。

如图 1 所示，T2T-ViT 在相似的计算成本 (MACs) 或模型规模下，相比 ViT、ResNets 和 MobileNets 一直取得更高的精度。

图 1: 在 ImageNet 上从零训练时，T2T-ViT 与 ViT、ResNets 和 MobileNets 的比较。左图: MACs 与 top-1 精度的性能曲线。右图: 模型大小与 top-1 精度的性能曲线。

图 1。将 T2T-ViT 与 ViT、ResNets 和 MobileNets 在 ImageNet 上从零训练时的比较。左图: MACs 与 top-1 精度的性能曲线。右图: 模型大小与 top-1 精度的性能曲线。

本文将剖析 T2T-ViT 的核心思想: 创新性的 Tokens-to-Token (T2T) 逐步分词模块、“深窄”主干设计，以及验证这些架构选择的实验证据。

背景: 为什么 ViT 在局部结构上表现不佳

Transformer 层通过自注意力机制擅长建模全局交互。然而，当图像被简单地切分为例如 16x16 的非重叠补丁时，每个补丁就成为了一个令牌，但并没有明确地包含其细粒度的局部结构或与邻近补丁的关系。相反，CNN 通过多个阶段使用小卷积核逐步构建局部特征的层次表示 (例如先检测边缘、再是纹理、然后是物体部件) 。这种固有的归纳偏置使得 CNN 在较少训练样本时也能有效学习。

为研究这种差异，作者可视化了 ResNet-50、ViT-L/16 和他们提出的 T2T-ViT-24 学到的特征。图 2 展示了一个显著对比: ResNet-50 在早期和中间层逐步发展出可识别的边缘和纹理检测器。相比之下，ViT 的令牌往往过早地编码高层或全局信号，且许多通道似乎是“失效”的 (全为零) ，这表明可能存在效率低下或冗余。T2T-ViT-24 则展示了在早期能够学习到有意义的局部特征，同时仍然利用 Transformer 的全局建模能力。

图 2: 在 ImageNet 上训练的 ResNet50、ViT-L/16 和 T2T-ViT-24 的特征可视化。绿色框突出显示学习到的低级结构特征 (如边缘和线条) ；红色框突出显示无效的特征图。

图 2。对在 ImageNet 上训练的 ResNet50、ViT-L/16 [12] 和我们提出的 T2T-ViT-24 的特征可视化。绿色框突出显示学习到的低级结构特征，例如边缘和线条；红色框突出显示值为零或过大的无效特征图。请注意，这里为 ViT 和 T2T-ViT 可视化的特征图不是注意力图，而是从令牌重塑得到的图像特征。为了更好的可视化，我们将输入图像缩放到大小 $1024 \times 1024$ 或 $2048 \times 2048$。

核心思想一: Tokens-to-Token (T2T) 逐步分词

Tokens-to-Token (T2T) 模块通过反复进行“重结构化” (Re-structurization) 和“软切分” (Soft Split) 来逐步将像素转换为令牌，从而解决了 ViT 简单分词的局限。每次迭代都将局部上下文注入令牌的表示中，同时减少令牌序列长度。这确保后续的 Transformer 层在更紧凑但具有局部信息的令牌上进行操作。

高层次上，一个 T2T 迭代执行如下步骤 (如图 3 所示) :

  1. 应用小型 Transformer: 来自前一层的当前令牌序列 \(T\) 被一个自注意力块处理 (图 3 中称为 T2T Transformer) :

    \[ T' = \text{MLP}(\text{MSA}(T)) \]

    这里 MSA 表示带有层归一化的多头自注意力操作，MLP 是多层感知机，同样带有层归一化，遵循标准的 Transformer 架构。

  2. 重结构化 (Restructurization) : 变换后的令牌序列 \(T'\) 被重塑为二维的空间网格，有效地恢复其空间邻域信息:

    \[ I = \text{Reshape}(T') \]

    这个 Reshape 操作将令牌 \(T' \in \mathbb{R}^{l \times c}\) 重新组织为一张“图像” \(I \in \mathbb{R}^{h \times w \times c}\)，其中 \(l\) 是 \(T'\) 的长度，\(h, w, c\) 分别是高、宽和通道维，且 \(l = h \times w\)。

  3. 软切分 (Soft Split) : 在得到重结构化的图像 \(I\) 后，应用“软切分”。这一步从 \(I\) 中提取重叠补丁以建模局部结构并减少令牌序列长度。每个提取的补丁随后被展平为新的令牌。使用重叠补丁至关重要，因为它允许每个新令牌编码来自其局部邻域的信息，从而增强相邻令牌间的相关性并为局部结构 (例如边缘和线条) 建立先验。

在数学上，对于重构后的图像 \(I \in \mathbb{R}^{h \times w \times c}\)，以补丁大小 \(k \times k\)、重叠量 \(s\) (和填充 \(p\)) 进行切分，其输出令牌长度 \(l_o\) 为:

\[ l_o = \left\lfloor \frac{h + 2p - k}{k - s} + 1 \right\rfloor \times \left\lfloor \frac{w + 2p - k}{k - s} + 1 \right\rfloor. \]

每个大小为 \(k \times k \times c\) 的切分补丁被展平成一个令牌，得到输出令牌 \(T_o \in \mathbb{R}^{l_o \times ck^2}\)。这些新令牌随后被送入下一次 T2T 处理。

图 3: T2T 过程的示意图。令牌被重结构化为图像，带重叠地切分，然后串联形成新令牌。

图 3。T2T 过程的示意图。令牌 $T_i$ 在变换与重塑后被重结构化为图像 $I_i$；然后对 $I_i$ 进行带重叠的切分以再次生成令牌 $T_{i+1}$。具体地，如粉色面板所示，输入 $I_i$ 的四个令牌 (1,2,4,5) 被拼接以形成 $T_{i+1}$ 中的一个令牌。T2T transformer 可以是一个普通的 Transformer 层 [37]，也可以在 GPU 内存有限时使用诸如 Performer 层 [34] 等更高效的 Transformer 变体。

通过堆叠这些 T2T 步骤迭代进行 (作者通常使用两次迭代) ，该模块逐步将细粒度的局部结构信息聚合到越来越紧凑的令牌中。重要的是，T2T 模块内部使用的 T2T transformer 可以是标准的 Transformer 层，也可以是如 Performer 这样更高效的变体以管理 GPU 内存使用。

核心思想二: 深-窄 (Deep-Narrow) 主干

在 T2T 模块生成固定长度的令牌序列 \(T_f\) 之后，标准的 Transformer 主干 (包括一个分类标记、位置嵌入和堆叠的 Transformer 块) 会处理这些令牌以进行分类。这里的关键创新在于主干的拓扑结构: 作者主张采用“深-窄”设计——更多层但更小的通道宽度，而不是 ViT 的“浅而宽” (层数少、隐藏维度大) 结构。

具体地，主干按如下方式处理令牌:

\[ \begin{aligned} T_{f_0} &= [t_{\text{cls}}; T_f] + E, \quad E \in \mathbb{R}^{(l+1)\times d},\\ T_{f_i} &= \text{MLP}(\text{MSA}(T_{f_{i-1}})), \quad i=1\ldots b,\\ y &= \text{fc}(\text{LN}(T_{f_b})), \end{aligned} \]

其中 \(b\) 表示主干的深度 (层数) ，\(d\) 是隐藏维度。核心思想是让 \(d\) 更小而 \(b\) 更大 (例如对于 T2T-ViT-14，\(d=384\) 且 \(b=14\)) 。实验证明，这样可以在有效控制参数量和计算成本的同时提升特征的丰富性。

图 4 展示了完整的 T2T-ViT 架构: 输入图像首先经过初始的软切分，随后进行多个 T2T 迭代。得到的固定令牌连同分类标记和位置嵌入一起被送入深-窄的 Transformer 主干，最终通过一个 MLP 头进行分类。

图 4: T2T-ViT 的整体网络架构。T2T 模块逐步减少令牌数量；主干在紧凑的令牌上应用全局注意力。

图 4。T2T-ViT 的整体网络架构。在 T2T 模块中，输入图像首先被软切分为补丁，然后展开为令牌序列 $T_0$。在 T2T 模块中令牌长度被逐步减少 (此处使用两次迭代并输出 $T_f$) 。然后 T2T-ViT 主干接收固定长度的令牌作为输入并输出预测。两个 T2T 区块与图 3 相同，PE 表示位置嵌入 (Position Embedding) 。

设计选择与模型变体

论文引入了若干 T2T-ViT 模型尺寸，旨在与现有的 CNN 可比:

  • T2T-ViT-14: 主干具有 14 层，隐藏维度为 384，约 2150 万参数和 4.8G MACs。该模型的设计目的是与 ResNet50 大小相当。
  • T2T-ViT-19 与 T2T-ViT-24: 这些是更大的变体，通过增加主干深度来提供更高的精度。
  • 轻量变体 (T2T-ViT-7、T2T-ViT-12) : 这些较小的模型专门设计用于在效率和性能方面与 MobileNets 竞争。

关键实现细节包括:

  • T2T 模块使用相对较小的通道维度 (例如 32 或 64) 。这对管理计算负担至关重要，特别是因为 T2T 模块早期阶段的令牌长度通常比 ViT 传统的 16x16 切分要大。
  • 在 GPU 内存受限的情况下，T2T 模块内部的 Transformer 层可以选择性地被更高效的变体 (如 Performer 层) 替换。

实验与结果

所有实验均在 ImageNet 数据集 (约 130 万张图像) 上从零训练模型完成。所有模型采用一致的训练方案，包括标准的数据增强 (mixup、cutmix) 、AdamW 优化器、余弦学习率衰减，并训练约 310 个 epoch，以确保公平比较。

T2T-ViT 与 ViT 的比较

论文中的表 2 明确表明在从零训练时 T2T-ViT 相较于原始 ViT 的优越性能。例如，ViT-S/16 模型拥有 4860 万参数和 10.1G MACs，得到 78.1% 的 top-1 精度。相比之下，T2T-ViT-14 仅有约 2150 万参数和 4.8G MACs (大致是参数和 MACs 的一半) ，却达到了 81.5% 的 top-1 精度——绝对提升超过 3%。

表 2: 在 ImageNet 上从零训练时，T2T-ViT 与 ViT 的比较。

表 2。在 ImageNet 上从零训练时 T2T-ViT 与 ViT 的比较。

T2T-ViT 与 ResNets 及 MobileNets 的比较

作者精心设计了 T2T-ViT 模型，使其在尺寸和计算成本上与流行的 CNN 架构可比。结果显示 T2T-ViT 明显占优:

  • 对比 ResNets: T2T-ViT-14 (21.5M 参数，4.8G MACs) 实现了 81.5% 的精度，优于 ResNet-50 (25.5M 参数，4.3G MACs) ，后者在作者的重现设置中约为 79.1%。这代表了对类似规模的 ResNet 的 1.4% 至 2.7% 的性能提升。
  • 对比 MobileNets: 轻量的 T2T-ViT 变体也显示出强劲的性能。例如，T2T-ViT-12 (6.9M 参数) 达到了 76.5% 的 top-1 精度，超越了他们实验设置中的 MobileNetV2 1.4x。尽管 T2T-ViT 模型由于 Transformer 中的稠密操作在 MACs 上仍高于 MobileNets，但这些结果突出了 T2T-ViT 作为一种高效轻量模型的潜力，尤其是在它并未使用诸如高效卷积之类的专门操作的情况下。

消融研究要点

消融研究为 T2T 模块和深-窄架构设计的有效性提供了有力证据。

  1. T2T 模块至关重要: 从 T2T-ViT-14 中移除 T2T 模块，但保留相同的主干，会导致大约 2% 的明显精度下降 (见表 7) 。用三层卷积 (以模拟补丁聚合) 替换 T2T 模块比没有 T2T 表现要好，但仍比 T2T Transformer 变体低 0.5% 至 1.0%。这表明 T2T 模块将局部聚合与逐令牌的全局建模结合的独特方式，比简单的卷积更有效地捕捉结构化信息。
  2. 深-窄结构优于浅-宽: 当将 T2T-ViT 的主干从深-窄结构重新配置为浅-宽 (层数更少、隐藏维度更大) 时，top-1 精度下降约 2.7% (表 7) 。这验证了这样一个假设: 更多层但更窄的层堆栈可以更高效地构建更丰富的层次特征，即使在通道宽度减少的情况下亦如此。

表 7 (如下复现) 总结了这些关键的消融发现。

表 7: 关于 T2T 模块与深-窄 (DN) 结构的消融研究结果。

表 7。关于 T2T 模块与深-窄 (DN) 结构的消融研究结果。

将 CNN 思想迁移到 ViT: 哪些有用 (哪些有害)

作者进一步探索了将各种 CNN 架构策略借鉴到 ViT 与 T2T-ViT 上的影响。

  • 深-窄结构: 这是最有效的迁移，显著减少了模型大小和 MACs (近 2 倍) ，同时在基线 ViT-S/16 上提升了 0.9% 的性能。
  • Squeeze-and-Excitation (SE) 模块: 实现通道注意力 (如 SE 网络) 对 ViT 与 T2T-ViT 都带来了适度的提升。
  • 密集连接 (DenseNet 风格) : 出人意料的是，密集连接在这些实验中对 ViT 和 T2T-ViT 都带来了负面影响。
  • ResNeXt 结构 (更多的注意力头) : 采用更多的注意力头 (例如 32 个) 对性能影响甚微，但显著增加 GPU 内存使用，使其成为一个低效的选择。
  • Ghost 操作 (GhostNet 风格) : 虽然 Ghost 操作可以进一步压缩模型并减少 MACs，但也导致性能下降 (例如对 T2T-ViT 下降约 2%) 。然而，这种降幅对 T2T-ViT 来说比对原始 ViT 或 ResNet 来说不那么严重。

总体而言，深-窄的重新配置被证明是最显著的架构改进。在所有测试的受 CNN 启发的结构中，T2T-ViT 一直优于 ViT，这进一步强调了所提出 T2T 模块的优势。

实用启示

  • 分词 (Tokenization) 至关重要: 图像的分词方式对视觉 Transformer 的性能有深远影响。一次性、硬性切分会丢失重要的局部结构。像 T2T 模块那样的迭代重叠聚合，在为 Transformer 做全局关系建模准备令牌的同时，有效保留了局部性。
  • 架构归纳偏置很重要: 在没有大规模预训练数据集的情况下训练视觉 Transformer 时，引入架构上的归纳偏置 (如局部聚合和更深更窄的层堆栈) 非常有益。
  • 从零训练也能有竞争力的表现: 通过对输入处理管道和主干架构的合理重设计，可以使 Transformer 从零训练便在 ImageNet 上达到竞争甚至更优的性能。

最后的思考与影响

T2T-ViT 是将两种不同范式的优势——CNN 的归纳偏置与 Transformer 的全局建模能力——通过精心的架构设计和谐结合的优秀示例。T2T 模块显式地将局部结构信息重新引入令牌，直接解决了天真的 ViT 初始分词的核心缺陷。同时，深-窄主干表明，在为视觉任务扩展 Transformer 时，并不仅仅是增加层的宽度；增加深度可能是提升特征丰富性更高效且更有效的杠杆。

该工作为若干有前景的后续方向铺平了道路:

  • 进一步混合化 (Hybridization) : 探索更复杂的可学习卷积与 T2T 步骤的组合，或设计能学习最优切分和重叠参数的 T2T 模块。
  • 高效变体: 为轻量级 T2T-ViT 模型开发更计算高效的注意力机制，可能借助局部注意力、Linformer 或 Performer 等变体来减少稠密计算开销。
  • 下游任务: 将 T2T-ViT 扩展到其他视觉任务，如目标检测、语义分割和图像增强，在这些任务中保留局部性的分词可能带来显著优势。