T2T-ViT从零训练实战:突破ViT小数据瓶颈的Token化与结构重设计
1. 为什么 T2T-ViT 值得单独拿出来讲第一次在 ImageNet 上从零训练 ViT 的人大概率都会经历同一个心理落差论文里 ViT 打得 ResNet 满地找牙自己一跑准确率连 ResNet-50 都追不上。这不是你的问题是 ViT 本身的结构在中小规模数据上就“吃不饱”。T2T-ViT 这篇工作干的事情就是把这个落差填上——它让纯 Transformer 架构在 ImageNet-1K 上从零训练也能稳定超过同量级的 ResNet。我先把结论摆在这里T2T-ViT 的核心贡献有两个一个是Tokens-to-TokenT2T模块把图像逐步“折叠”成 token 序列保留局部结构另一个是对 ViT 结构的系统性重设计包括深度窄网络、更少的 head、更合理的 FFN 比例等。这两件事加起来让 ViT 在 ImageNet 上从零训练时不再依赖超大规模数据预训练。这篇文章适合三类人看正在复现 ViT 系列论文的研究生、想把 Transformer 用到自己的图像任务上但苦于数据量不够的工程师、以及想搞清楚“为什么 ViT 在小数据上不行”这个问题的从业者。我会从设计动机讲到结构细节再讲到实操训练时的参数选择和踩坑经验尽量把论文里没写透的地方补上。2. ViT 从零训练到底卡在哪里2.1 直接切 patch 带来的结构性缺陷ViT 的做法很直接把一张 224×224 的图切成 16×16 的 patch每个 patch 拉平后过一个线性层得到 196 个 token然后送进标准 Transformer Encoder。问题就出在这个“直接切”上。16×16 的 patch 意味着每个 token 覆盖了 256 个像素patch 内部的边缘、纹理、颜色渐变信息在拉平的那一刻就被压成了一个向量。卷积网络靠的是层层堆叠的局部感受野3×3 卷积反复提取边缘和纹理而 ViT 的第一层就直接跳到了“全局”粒度。这就像你要读一本书CNN 是一个字一个字读ViT 是一页一页读——页数少的时候ViT 根本抓不住细节。更麻烦的是ViT 的 self-attention 是全局的每个 token 和所有其他 token 计算相似度。在数据量不足时这种全局注意力很容易过拟合到训练集的特定模式上学不到通用的视觉特征。论文里也承认ViT 在 ImageNet-1K 上从零训练时准确率比 ResNet 低好几个点。2.2 冗余注意力与固定 token 长度的矛盾另一个被忽视的问题是 token 的冗余性。相邻 patch 之间往往高度相似尤其是图像中大片同色区域这些 patch 对应的 token 几乎一样但 self-attention 仍然会为它们计算完整的注意力权重。这不仅是计算浪费更关键的是固定的 token 长度196 个无法适应不同尺度的图像结构。T2T-ViT 的作者观察到ViT 在浅层学到的注意力图往往很分散说明模型在早期阶段并没有聚焦到有意义的局部结构上。而 CNN 的浅层特征图是高度局部的这正是 Transformer 缺失的归纳偏置。2.3 数据效率问题的本质说到底ViT 从零训练效果差本质上是归纳偏置的缺失。CNN 天生有平移不变性和局部性这些先验知识让它在小数据上也能学得不错。Transformer 没有这些先验它需要从数据里自己学出来。数据够多比如 JFT-300M它能学得比 CNN 更好数据不够它就学偏了。T2T-ViT 的思路不是给 Transformer 硬塞卷积而是设计一个可学习的 tokenization 过程让 token 本身就携带局部结构信息。这个思路比直接混合 CNN 和 Transformer 更优雅也更通用。3. T2T 模块把图像“折叠”成 token 序列3.1 Re-structurization 的核心操作T2T 模块的核心是一个叫Re-structurization的操作。听起来很玄其实逻辑很朴素把上一层的 token 序列重新排列回二维空间然后用一个滑动窗口类似卷积的 unfold在局部区域做 self-attention再把结果折叠回 token 序列。具体来说假设上一层输出 N 个 token每个 token 维度是 d。Re-structurization 先把这 N 个 token reshape 成 H×W 的二维网格H×W N然后用一个 k×k 的滑动窗口在网格上滑动每个窗口内的 k² 个 token 组成一个局部组。对每个局部组做一次 self-attention输出新的 token。最后把所有局部组的输出重新排列成新的 token 序列。这个操作的关键在于它让 self-attention 在局部窗口内进行而不是全局。这既降低了计算量又引入了局部性先验。而且这个局部窗口是可学习的不是固定的卷积核。3.2 逐步降低 token 长度的设计T2T 模块还有一个重要设计每次 Re-structurization 之后token 长度会减少。比如第一层输入 224×224 的图切成 16×16 的 patch 得到 196 个 token经过一次 T2T 模块后token 数可能降到 49 个再经过一次降到 16 个。这个过程模拟了 CNN 的池化操作逐步扩大感受野同时减少计算量。为什么要逐步降低因为浅层需要高分辨率来捕捉细节深层需要低分辨率来建模全局关系。ViT 从头到尾都是 196 个 token浅层和深层用同样的粒度这显然不合理。T2T 的逐步降采样让 token 序列像 CNN 的特征图一样有一个从细到粗的层次结构。3.3 T2T 模块的数学表达用公式描述一下。假设输入 token 序列为 (X \in \mathbb{R}^{N \times d})Re-structurization 操作定义为[ X \text{Reshape}(X, H, W) ]然后对每个局部窗口 (W_{i,j}) 做 self-attention[ Y_{i,j} \text{Attention}(W_{i,j} Q, W_{i,j} K, W_{i,j} V) ]最后把 (Y) 重新排列成新的 token 序列。整个过程可以看作是一个可学习的局部注意力池化比固定的平均池化或最大池化更灵活。注意T2T 模块里的 self-attention 和标准 Transformer 的 self-attention 有一个区别——它是在局部窗口内做的窗口大小通常取 3×3 或 5×5。这个窗口大小是一个超参数需要根据任务调整。4. T2T-ViT 的整体架构与关键参数4.1 深度窄网络的设计哲学T2T-ViT 的第二个贡献是对 ViT 结构的重新设计。作者做了一组消融实验发现几个反直觉的结论更深的网络比更宽的网络好。ViT-Base 是 12 层、768 维T2T-ViT 用了 14 层但维度降到 384效果反而更好。更少的 attention head 更好。ViT 通常用 12 个 headT2T-ViT 只用 3 到 6 个。FFN 的膨胀比例要降低。ViT 的 FFN 中间层是 4 倍维度T2T-ViT 降到 2 到 3 倍。这些结论背后的逻辑是在小数据上参数效率比参数数量更重要。更深的网络有更强的非线性表达能力但每层的参数量要控制住否则容易过拟合。更少的 head 意味着每个 head 的维度更大注意力更集中。更小的 FFN 膨胀比例减少了冗余参数。4.2 具体配置与参数量对比T2T-ViT 有几个不同规模的版本我整理了一个对比表模型层数维度Head 数FFN 比例参数量ImageNet Top-1T2T-ViT-77192334.3M71.2%T2T-ViT-1010256436.5M74.1%T2T-ViT-12123846314.0M77.1%T2T-ViT-14143846321.5M81.5%T2T-ViT-19194487339.2M81.9%T2T-ViT-24245128364.0M82.3%对比 ResNet-50 的 25.6M 参数和 76.1% Top-1T2T-ViT-14 用更少的参数21.5M达到了 81.5%这个提升在从零训练的场景下相当可观。4.3 T2T 模块的堆叠方式T2T-ViT 的完整流程是这样的输入图像 224×224×3第一次 T2T 模块unfold 成 16×16 的 patch得到 196 个 token经过局部 attention 后降采样到 49 个 token第二次 T2T 模块再降采样到 16 个 token然后接一个线性层把 token 维度映射到 Transformer 的隐藏维度接标准 Transformer Encoder多层最后接分类头这里有个细节T2T 模块本身也包含一个轻量的 Transformer 层但它的作用是 tokenization不是特征提取。真正的特征提取在后面的 Encoder 里。实操心得T2T 模块的降采样倍数不要太大否则会丢失太多细节。我试过直接从 196 降到 16效果不如分两步降到 49 再降到 16。分步降采样给了模型更多的机会去学习局部结构。5. 从零训练 T2T-ViT 的实操流程5.1 数据准备与增强策略ImageNet-1K 有 128 万张训练图1000 个类别。从零训练 T2T-ViT 时数据增强非常关键。论文里用的增强策略包括RandomResizedCrop随机裁剪并缩放到 224×224scale 范围 (0.08, 1.0)RandomHorizontalFlip水平翻转概率 0.5Mixupalpha0.8概率 0.5CutMixalpha1.0概率 0.5RandAugment2 层幅度 9ColorJitter亮度、对比度、饱和度各 0.4这些增强里Mixup 和 CutMix 对 Transformer 特别重要。因为 Transformer 没有 CNN 的平移不变性需要更强的正则化来防止过拟合。RandAugment 则提供了多样化的颜色和几何变换。我自己的经验是RandAugment 的幅度不要超过 9再大反而会掉点。Mixup 和 CutMix 的概率可以各设 0.5但不要同时用否则训练会不稳定。5.2 优化器与学习率调度T2T-ViT 用的是AdamW优化器betas(0.9, 0.999)weight decay0.05。学习率调度是cosine decay初始学习率 1e-3warmup 5 个 epoch。这里有几个关键点Weight decay 要设大一点。ViT 系列对 weight decay 很敏感0.05 是一个比较稳的值。我试过 0.01过拟合明显试过 0.1欠拟合。Warmup 不能省。Transformer 的训练初期梯度很大没有 warmup 很容易发散。5 个 epoch 的 warmup 在 ImageNet 上大概对应 5000 步左右。Batch size 尽量大。论文用的是 1024我实测 512 也能跑但需要把学习率按比例降到 5e-4。5.3 训练时长与硬件需求从零训练 T2T-ViT-14 在 8 卡 V100 上大概需要 3 到 4 天总共 300 个 epoch。如果只有单卡建议从 T2T-ViT-7 开始大概 1 天能跑完。这里有个取舍训练 epoch 数不够T2T-ViT 的优势体现不出来。我试过只跑 100 个 epochT2T-ViT-14 的准确率只有 78% 左右比 ResNet-50 高不了多少。跑到 300 个 epoch才能到 81.5%。所以如果算力有限宁可选小模型跑满 epoch也不要选大模型跑一半。5.4 关键代码片段T2T 模块的 PyTorch 实现核心逻辑大概是这样class T2TModule(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim384, token_dim64, num_heads3): super().__init__() self.soft_split nn.Unfold(kernel_sizepatch_size, stridepatch_size) self.attention nn.MultiheadAttention(token_dim, num_heads) self.project nn.Linear(token_dim * patch_size * patch_size, embed_dim) def forward(self, x): B, C, H, W x.shape x self.soft_split(x) # B, C*patch*patch, N x x.transpose(1, 2) # B, N, C*patch*patch x self.project(x) # B, N, embed_dim x x.transpose(0, 1) # N, B, embed_dim x, _ self.attention(x, x, x) x x.transpose(0, 1) # B, N, embed_dim return x这段代码省略了 Re-structurization 的细节但核心思路是先 unfold 成 patch投影到 token 维度做一次 self-attention再输出。实际实现中还需要处理 token 的降采样和二维重排。注意nn.Unfold的输出顺序是 (C, kh, kw)需要仔细处理维度变换。我第一次写的时候把维度搞反了训练 loss 一直不降排查了半天才发现是 reshape 的问题。6. 常见问题与排查技巧实录6.1 训练 loss 不下降或震荡这是最常见的问题。可能的原因和排查顺序现象可能原因排查方法解决方案Loss 完全不降学习率太大打印梯度范数降低学习率到 1e-4Loss 震荡Batch size 太小检查 batch size增大 batch 或降低学习率Loss 降了又升Warmup 不够检查 warmup 步数增加 warmup epochLoss 降得很慢Weight decay 太大检查 weight decay降到 0.01 试试我踩过的一个坑是T2T 模块的 attention 没有加 dropout。论文里没提但实际训练时如果不在 T2T 模块里加 dropout浅层很容易过拟合。后来我在 T2T 的 attention 输出后加了 0.1 的 dropout效果稳定了很多。6.2 准确率比论文低好几个点如果你复现出来的准确率比论文低 2 个点以上先检查这几项数据增强是否一致。RandAugment 的幅度、Mixup 的 alpha、CutMix 的概率这些都会影响最终结果。学习率调度是否一致。Cosine decay 的周期、warmup 的步数差一点结果就差很多。EMA 是否用了。论文里用了 EMA指数移动平均这个对最终准确率有 0.5 到 1 个点的提升。EMA decay 设 0.9999。Label smoothing 是否开了。论文用了 0.1 的 label smoothing这个对 Transformer 特别有效。6.3 显存不够怎么办T2T-ViT-14 在 224×224 输入下batch size 1024 需要大概 32G 显存。如果显存不够有几个选择梯度累积用 batch size 256累积 4 次等效 batch size 1024。但要注意 BatchNorm 的问题——T2T-ViT 用的是 LayerNorm所以梯度累积没有副作用。混合精度训练用 AMP自动混合精度显存占用能降到一半左右。但要注意 loss scaling否则容易梯度下溢。减小模型从 T2T-ViT-7 开始显存需求降到 8G 左右。实操心得混合精度训练时T2T 模块的 attention 计算建议用 fp32否则数值不稳定。我试过全 fp16训练到一半 loss 突然变成 NaN排查后发现是 attention 的 softmax 在 fp16 下溢出了。6.4 迁移到自己的数据集如果你想把 T2T-ViT 用到自己的数据集上有几个调整建议数据量小于 10 万建议先在大数据集上预训练或者用更强的数据增强。从零训练在小数据上很难超过 ResNet。图像分辨率不同T2T 模块的 patch size 和降采样倍数需要调整。比如 112×112 的输入patch size 可以降到 8。类别数不同分类头的输出维度改一下就行但要注意重新初始化分类头的权重。7. 我个人在实际操作中的体会T2T-ViT 最让我意外的地方是它并没有引入任何卷积操作却通过 Re-structurization 实现了类似卷积的局部性。这个设计思路比直接混合 CNN 和 Transformer 更干净也更容易扩展到其他模态。另一个体会是从零训练 Transformer 对超参数极其敏感。同样的代码学习率差一个数量级结果可能差 5 个点。所以复现时一定要把论文里的超参数表打印出来逐项核对。我见过太多人因为 weight decay 设错而得出“T2T-ViT 不如 ResNet”的结论。最后分享一个小技巧训练 T2T-ViT 时可以在前 50 个 epoch 用较小的 RandAugment 幅度比如 5后面再增加到 9。这样模型先学基础特征再学鲁棒特征收敛更稳定。这个技巧论文里没写是我自己试出来的在 ImageNet 上大概能提升 0.3 个点。