1. 动态视觉Transformer的核心思路拆解

1.1 从静态计算到动态剪枝的范式转变

视觉Transformer（ViT）这两年在图像分类、目标检测等任务上表现抢眼，但有个绕不开的痛点：计算量太大。一张224×224的图片切成16×16的patch，就是196个token，再加上一个cls token，总共197个token要过每一层Transformer。12层下来，中间的特征图计算量和显存占用都相当可观。更关键的是，图片里大量区域其实是冗余的——背景、纯色块、重复纹理，这些patch对最终分类结果的贡献微乎其微，但模型依然一视同仁地给它们分配了完整的计算资源。

DynamicViT的核心洞察就一句话：不是所有token都值得算到底。它引入了一个轻量级的预测模块，在Transformer的若干层之后，动态判断哪些token可以丢弃，只保留信息量高的token继续往后传。这个思路和NLP里的token pruning、计算机视觉里的动态网络一脉相承，但DynamicViT把这件事做得足够优雅——通过端到端的训练让模型自己学会“该丢谁、该留谁”。

我最初看到这个标题时的第一反应是：剪枝不难，难的是剪完之后性能不掉。很多剪枝方案在ImageNet上掉个1%到2%的top-1精度是常态，但DynamicViT宣称“几乎不降低性能”，这背后必然有精巧的设计。拆开来看，它的技术路线可以概括为三个关键词：可微分掩码、层级式剪枝、注意力保留。

1.2 为什么选择层级式渐进剪枝而非一次性裁剪

一个自然的疑问是：为什么不直接在某一层把不重要的token一次性全砍掉？答案在于浅层特征和深层特征的语义差异。浅层Transformer主要捕捉局部纹理和边缘信息，这时候很多patch看起来都“有用”，因为语义还没聚合；到了深层，cls token已经和关键区域建立了强关联，冗余token的辨识度才真正显现出来。

DynamicViT的做法是在第3、6、9层之后分别插入剪枝模块，每次丢掉一部分token，逐步降低序列长度。这种渐进式策略给了模型“缓冲”的空间——每次剪枝后，保留的token还能通过后续的注意力层重新整合信息，而不是一刀切导致信息断裂。实测数据也支持这个判断：一次性剪枝50%的token，精度掉得厉害；分三次每次剪20%左右，最终保留40%到50%的token，精度几乎无损。

这里有个经验性的比例分配：假设目标是把token保留率降到50%，那么三次剪枝的保留率可以设为0.7、0.7、0.7，连乘约0.34，或者0.8、0.8、0.8，连乘约0.51。具体怎么设，要看你的任务对精度的敏感程度。分类任务可以激进一点，检测和分割任务建议保守，因为空间定位信息更依赖完整的token分布。

1.3 可微分掩码：让“丢弃”这件事变得可训练

剪枝最大的技术障碍是离散决策不可导。你没法直接对“保留还是丢弃”这个二值操作求梯度。DynamicViT的解决方案是引入一个轻量级的MLP预测头，对每个token输出一个保留概率，然后用Gumbel-Softmax技巧在训练时做可微分的采样。推理时则直接按概率阈值硬截断，只保留概率高的token。

这个设计的好处是：训练和推理的行为一致。很多剪枝方法训练时用软掩码、推理时用硬掩码，中间存在gap，导致实际部署时精度下降。DynamicViT通过Gumbel-Softmax的退火策略，让训练后期的采样结果越来越接近硬决策，从而缩小了这个gap。

另一个细节是注意力掩码的同步更新。token被丢弃后，后续层的注意力矩阵维度要相应缩小，否则计算会出错。DynamicViT在实现上维护了一个token索引映射表，每次剪枝后更新这个表，确保注意力计算只在保留的token之间进行。这个工程细节看似简单，但实际写代码时很容易在位置编码和残差连接上翻车。

2. 核心模块的技术细节与实操要点

2.1 预测模块的网络结构设计

DynamicViT的预测模块本身非常轻量，通常由两层MLP加一个LayerNorm组成，参数量只占主干网络的百分之几。输入是当前层的token特征，输出是每个token的保留概率。具体来说，对于维度为D的token特征，第一层MLP把维度降到D/4，第二层再映射到1维，最后过Sigmoid得到概率值。

这里有个关键设计：预测模块的输入不只是单个token的特征，还融合了全局信息。具体做法是把cls token的特征广播到每个token上，和局部特征拼接后再送入MLP。这样做的理由是，cls token聚合了全局语义，能帮助判断某个局部patch对最终任务是否重要。比如在分类“狗”的图片时，cls token已经编码了“狗”的语义，背景patch的保留概率自然会被压低。

实操中，预测模块的学习率通常设得比主干网络高一些，比如主干用1e-4，预测模块用1e-3。原因是预测模块是随机初始化的，需要更快地学到有效的剪枝策略。如果学习率设成一样，训练前期预测模块输出接近均匀分布，剪枝效果出不来，后期又可能因为梯度太小而收敛缓慢。

2.2 剪枝比例与层数的配置策略

剪枝配置没有万能公式，但有一些经过验证的经验规律。对于DeiT-Small（12层，384维），一个常见的配置是在第3、6、9层后剪枝，保留率分别为0.7、0.7、0.7，最终保留约34%的token。对于DeiT-Base（12层，768维），可以稍微激进一点，保留率0.65、0.65、0.65，最终约27%。

但要注意，剪枝太早会伤精度。第1、2层的特征还很底层，token之间的区分度不高，这时候剪枝容易误伤。一般建议从第3层之后开始剪。如果模型层数更多，比如24层的ViT-Large，可以在第6、12、18层后剪，给浅层更多的计算资源来提取基础特征。

还有一个容易被忽略的点：cls token永远不参与剪枝。它是最终分类的依据，必须全程保留。实现上，在计算保留概率时把cls token的概率强制设为1，或者在索引映射时单独处理。我见过有人忘了这一点，结果cls token被剪掉，模型直接输出随机结果。

2.3 训练技巧：退火温度与正则化

Gumbel-Softmax的温度参数τ控制采样的随机性。τ越大，采样越接近均匀分布；τ越小，采样越接近argmax的硬决策。DynamicViT的训练策略是让τ从1.0逐渐退火到0.1，前期鼓励探索不同的剪枝模式，后期固定下来。

退火的速度需要调。退火太快，模型还没学好哪些token重要就固化了策略，精度会掉；退火太慢，训练后期还在随机采样，收敛不了。一个实用的做法是前50%的epoch保持τ=1.0，后50%线性降到0.1。如果训练资源充足，可以用余弦退火，效果更平滑。

另外，剪枝模块需要额外的正则化。如果不加约束，模型倾向于保留所有token，因为这样精度最高，但计算量没降下来。DynamicViT引入了一个目标保留率的损失项，惩罚实际保留率与目标保留率的偏差。这个损失项的权重通常设在0.01到0.1之间，太大影响分类精度，太小起不到约束作用。

3. 完整实操流程与关键环节实现

3.1 环境准备与代码框架搭建

动手复现DynamicViT之前，先把环境理清楚。PyTorch 1.9以上、torchvision 0.10以上是基础，timm库用来加载预训练的ViT权重。硬件方面，单卡24G显存可以跑DeiT-Small的剪枝实验，batch size设256左右；DeiT-Base建议用32G以上的卡，或者用梯度累积。

代码结构上，核心是三个类：DynamicViT继承原始的ViT，在指定层后插入PruningModule；PruningModule包含MLP预测头和Gumbel-Softmax采样逻辑；TokenIndexManager负责维护每次剪枝后的token索引映射。下面是一个简化的预测模块实现：

import torch import torch.nn as nn import torch.nn.functional as F class PruningModule(nn.Module): def __init__(self, dim, keep_ratio): super().__init__() self.keep_ratio = keep_ratio self.mlp = nn.Sequential( nn.Linear(dim * 2, dim // 4), nn.GELU(), nn.Linear(dim // 4, 1) ) self.norm = nn.LayerNorm(dim) def forward(self, x, cls_token, tau=1.0, hard=False): # x: [B, N, D], cls_token: [B, 1, D] cls_expanded = cls_token.expand(-1, x.size(1), -1) feat = torch.cat([x, cls_expanded], dim=-1) logits = self.mlp(self.norm(feat)).squeeze(-1) # [B, N] if self.training: # Gumbel-Softmax 可微分采样 noise = -torch.log(-torch.log(torch.rand_like(logits) + 1e-8) + 1e-8) logits = (logits + noise) / tau prob = torch.sigmoid(logits) else: prob = torch.sigmoid(logits) # 按保留率确定阈值 k = int(x.size(1) * self.keep_ratio) threshold = torch.topk(prob, k, dim=1).values[:, -1:] mask = (prob >= threshold).float() if hard: mask = (mask > 0.5).float() return mask, prob

这段代码里，cls_token的广播拼接是关键，它让预测模块能看到全局信息。hard参数控制推理时是否用硬掩码。实际部署时，hard=True，直接按阈值截断。

3.2 训练流程与损失函数配置

训练DynamicViT和训练普通ViT最大的区别在于损失函数。总损失由三部分组成：分类交叉熵损失、保留率约束损失、以及可选的蒸馏损失。分类损失是主任务，保留率损失约束实际保留的token比例接近目标值，蒸馏损失用原始未剪枝的ViT作为教师模型，让剪枝后的学生模型模仿教师的输出分布。

保留率损失的计算方式是：对每个剪枝模块，计算实际保留率与目标保留率的L2距离，然后对所有模块求平均。权重系数建议从0.05开始调，如果发现实际保留率远高于目标值，就加大权重；如果分类精度掉得厉害，就减小权重。

蒸馏损失在剪枝比例较高时特别有用。当保留率降到30%以下，单靠分类损失很难维持精度，这时候教师模型的软标签能提供额外的监督信号。蒸馏温度一般设3到5，损失权重设0.5到1.0。不过蒸馏会增加训练开销，因为要同时跑教师和学生两个模型，显存占用翻倍。

训练epoch数建议比原始ViT多20%到30%。因为剪枝模块需要额外的时间来学习有效的策略。以DeiT-Small为例，原始训练300 epoch，DynamicViT可以训400 epoch。优化器用AdamW，weight decay设0.05，主干学习率1e-4，剪枝模块学习率1e-3，余弦退火调度。

3.3 推理部署与加速效果实测

推理阶段，DynamicViT的加速效果来自两个方面：一是序列长度缩短，注意力计算的复杂度从O(N²)降到O(N'²)，N'是剪枝后的token数；二是后续层的MLP计算量按比例减少。以DeiT-Small为例，原始197个token，剪枝到约67个token（保留率34%），注意力矩阵从197×197降到67×67，计算量降到原来的11.6%。MLP部分的计算量降到34%。整体FLOPs大约降到原来的40%到50%。

实测数据方面，在ImageNet-1K上，DeiT-Small原始top-1精度79.8%，DynamicViT剪枝到34% token后精度79.4%，掉了0.4个百分点。DeiT-Base原始81.8%，剪枝后81.3%，掉0.5个百分点。这个精度损失在可接受范围内，尤其是考虑到FLOPs降了将近一半。

部署时要注意，剪枝后的batch内token数可能不一致。因为每张图片的内容不同，保留的token数会有波动。如果直接用PyTorch的tensor操作，需要padding到最大长度，这会浪费一些计算。更好的做法是用自定义的CUDA kernel或者用NVIDIA的FasterTransformer，它们支持变长序列的高效计算。如果只是做实验，padding方案够用；如果要上生产环境，变长支持是必须的。

4. 常见问题排查与避坑经验实录

4.1 精度掉点严重的排查思路

剪枝后精度掉得多，先查三个地方。第一，剪枝起始层是不是太早。如果从第1层后就开始剪，浅层特征还没充分提取，精度必掉。把起始层往后挪，比如从第3层挪到第5层，通常能挽回不少。第二，保留率是不是设得太激进。0.5的保留率对分类任务可能还行，但对检测任务就太狠了。先试0.7，逐步往下压，找到精度和速度的平衡点。第三，预测模块的学习率是不是太低。如果预测模块和主干用同一个学习率，它学不到有效的剪枝策略，相当于随机丢token，精度自然崩。

还有一个隐蔽的坑：位置编码的处理。token被剪枝后，剩余token的位置编码要相应调整。如果直接保留原始位置编码，被剪掉的token位置会出现“空洞”，后续的注意力计算会错位。正确的做法是维护一个索引映射，把保留token的原始位置编码取出来重新排列。这个细节在论文里可能一笔带过，但代码实现时漏掉就会导致精度大幅下降。

4.2 训练不收敛与loss震荡的解决

Gumbel-Softmax的随机性会导致训练前期loss震荡，这是正常现象。但如果震荡幅度过大或者一直不收敛，检查温度参数τ的退火策略。τ退火太快，模型还没探索充分就固化了策略；τ退火太慢，训练后期还在随机采样。建议前30%的epoch保持τ=1.0，中间40%线性降到0.5，最后30%降到0.1。

另外，保留率损失和分类损失的平衡也很关键。如果保留率损失权重太大，模型为了满足保留率目标而牺牲分类精度；权重太小，保留率降不下来。一个实用的调试方法是：先只用分类损失训练几个epoch，观察自然状态下的保留率分布，然后根据这个分布来设目标保留率和损失权重。如果自然保留率是0.9，你想降到0.5，那保留率损失的权重需要设得大一些，比如0.1；如果自然保留率是0.6，想降到0.5，权重设0.01就够了。

4.3 常见问题速查表

4.4 独家避坑技巧与经验总结

第一个技巧：用预训练权重初始化主干，但随机初始化预测模块。如果整个模型都从头训，收敛太慢；如果预测模块也用预训练权重（虽然通常没有），反而不利于学习剪枝策略。主干加载预训练权重后，先冻结几个epoch让预测模块 warm up，然后再联合微调。

第二个技巧：剪枝模块的插入位置不要等间距。等间距（比如第3、6、9层）是常见做法，但更好的策略是前期间隔大、后期间隔小。比如第4、8、10层，因为浅层需要更多计算来提取特征，深层可以更激进地剪。这个策略在DeiT-Base上能多挽回0.2%的精度。

第三个技巧：推理时用top-k而不是阈值。阈值方案在不同图片上的实际保留率波动大，有的图片保留60%，有的只保留20%，导致推理时间不稳定。用top-k固定保留数量，推理时间更可控。具体做法是每张图片按概率排序，取前k个token，k等于目标保留率乘以原始token数。

第四个技巧：蒸馏损失在剪枝比例高时必不可少。当保留率低于40%，单靠分类损失很难维持精度，教师模型的软标签能提供额外的监督。但蒸馏温度要调，太高了软标签太“软”，信息量不足；太低了接近硬标签，和分类损失重复。3到5是比较合适的范围。

5. 扩展方向与个人实操体会

DynamicViT的思路可以扩展到视频理解任务。视频比图片多了时间维度，冗余度更高——相邻帧之间大量内容重复，时间维度上的token剪枝空间更大。一个自然的想法是在空间剪枝的基础上，再加一层时间维度的剪枝模块，把静态背景帧的token大量丢弃。这个方向已经有了一些探索，但如何设计时空联合的剪枝策略还是个开放问题。

另一个扩展方向是与量化、蒸馏的结合。剪枝减少的是token数量，量化减少的是每个token的位宽，两者正交，可以叠加使用。比如先剪枝到50% token，再对保留的token做8bit量化，整体压缩率能到4倍以上。但要注意，剪枝后的token分布会变化，量化校准集需要重新采样，否则量化误差会放大。

我个人在实际操作中的体会是：DynamicViT的工程实现比论文看起来要复杂。论文里几行公式描述的东西，落到代码里涉及大量的张量维度对齐、索引映射维护、训练推理行为一致性处理。第一次复现时，我在位置编码更新上卡了两天，后来发现是索引映射的维度搞错了。建议动手之前先把整个数据流画清楚：输入token序列经过每一层后，哪些token被保留、索引怎么变、位置编码怎么取、注意力矩阵怎么缩。把这几个问题想明白了，代码写起来就顺了。

最后分享一个小技巧：调试剪枝模块时，先把保留率设成1.0（即不剪枝），确认模型能正常训练到基线精度。然后再逐步降低保留率，每次降0.1，观察精度变化。这样能快速定位是剪枝策略的问题还是主干网络的问题。如果保留率1.0时精度就不对，那说明代码有bug，跟剪枝无关。这个“控制变量”的调试方法帮我省了很多时间。