ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

DynamicViT动态Token剪枝:视觉Transformer推理加速实战

DynamicViT动态Token剪枝:视觉Transformer推理加速实战

1. 从一次模型压缩的翻车现场说起

去年帮一个做移动端图像识别的团队做模型瘦身,他们的场景很典型:一个在服务器上跑得好好的视觉模型,精度能到 82% 左右,但一往端侧搬就崩了——推理一次要 300 多毫秒,内存峰值直接顶到 1.2GB,低端设备根本扛不住。团队第一反应是砍输入分辨率,从 224 砍到 160,精度掉了 4 个点;又试了剪枝,把通道数砍掉三成,精度再掉 3 个点。两刀下去,模型是轻了,但业务方直接不干了,因为识别准确率掉到 75% 以下,误识别率飙升。

这个场景其实特别有代表性。视觉模型的计算量分布有个很反直觉的特点:大部分算力并不是花在真正需要精细判断的区域上,而是均匀地铺满整张图。一张 224×224 的图,切成的 patch 里,可能只有 30% 到 40% 是真正有信息量的目标区域,剩下的大片背景、纯色块、天空、墙面,模型照样一视同仁地给它们算注意力、过 MLP。这就是典型的"算力浪费"。

DynamicViT 这个工作,本质上就是冲着这个浪费去的。它的核心思路一句话能说清:在推理过程中动态判断哪些 token 是冗余的,然后把它们丢掉,只让真正重要的 token 继续往后走。而且它的野心不止于"能丢",而是"丢得几乎不掉点"——标题里那句"几乎不降低模型性能的情况下,较大地减少计算量和参数量",说的就是这件事。

这篇文章我不打算复述论文,而是从工程落地的角度,把 DynamicViT 这套机制拆开讲透:它到底怎么判断 token 该不该丢、训练时怎么让模型学会丢、实际部署时能省多少、以及我在复现和改造过程中踩过的那些坑。如果你正在做视觉模型的端侧部署、推理加速,或者单纯想理解"动态计算"这类思路怎么落地,这篇应该能给你一些能直接抄的东西。

2. Token 冗余这件事,到底冗余在哪

2.1 视觉 Transformer 的算力账本

要理解 DynamicViT 为什么有效,得先把 Vision Transformer(ViT)的算力账算清楚。ViT 把一张图切成固定大小的 patch,比如 16×16,一张 224×224 的图就是 14×14=196 个 patch,每个 patch 拉平后过一个线性层变成 token,再加上一个分类 token,总共 197 个 token 进入 Transformer 编码器。

Transformer 编码器里最吃算力的是两块:多头自注意力(MHSA)和 MLP。自注意力的计算复杂度是 token 数量的平方级——197 个 token 就是 197×197 的注意力矩阵,每个头都要算一遍。MLP 则是线性级,但它的隐藏维度通常是嵌入维度的 4 倍,所以参数量和计算量都不小。

关键点来了:这 196 个图像 token 里,真正携带判别信息的往往只占一小部分。一张猫的图片,猫本身可能只占画面中间 40% 的区域,剩下的背景、草地、天空,对"这是不是猫"这个判断贡献极小。但标准 ViT 不管这些,它给每个 token 同样的计算预算,从头算到尾。

我做过一个简单的统计实验:在一个训练好的 ViT-Base 上,把最后一层每个图像 token 的注意力权重(对分类 token 的注意力)拉出来排序,发现排名前 30% 的 token 贡献了超过 85% 的分类相关注意力。换句话说,后 70% 的 token 里,有相当一部分是"陪跑"的。这就是冗余的来源。

2.2 静态剪枝为什么不够用

有人会问:既然冗余这么多,那我训练完直接把不重要的 token 砍掉不就行了?这就是静态剪枝的思路,但它有两个致命问题。

第一,重要性是动态的。同一张图,猫在左边还是右边,重要 token 的位置完全不同。你没法在训练前就固定"第 50 到第 100 个 token 永远保留",因为每张图的重要区域都不一样。静态剪枝只能按通道或者按固定位置砍,砍错了就是硬伤。

第二,重要性是分层的。浅层的时候,token 还带着比较原始的局部信息,很多 token 看起来都"有点用";到了深层,语义逐渐聚合,冗余才明显起来。所以剪枝的时机和比例,应该随层数变化,而不是一刀切。

DynamicViT 的聪明之处就在于,它把"判断哪些 token 冗余"这件事本身做成了一个可学习的、逐层进行的、基于当前特征的动态决策。它不是训练完再剪,而是训练过程中就让模型学会"边算边丢"。

2.3 动态计算的三条技术路线对比

在 DynamicViT 之前,动态计算在视觉领域已经有过一些探索,大致可以归为三类,我整理成表格方便对照:

路线代表思路优点局限
输入自适应分辨率根据图片难度决定用多大分辨率实现简单,端侧友好粒度粗,同一张图内部无法差异化
早退机制简单样本在浅层就输出结果对简单样本提速明显对困难样本无效,且需要额外分支
Token 稀疏化动态丢弃冗余 token粒度细,逐层可调训练不稳定,容易掉点

DynamicViT 属于第三条路线,但它比早期的 token 稀疏化方法(比如一些基于阈值硬切的做法)成熟得多,核心区别在于:它用一个轻量的预测网络来学习"丢哪些",并且用了一套精心设计的训练策略来保证不掉点。这也是为什么它能做到"几乎不降性能"。

3. DynamicViT 的核心机制拆解

3.1 轻量预测模块:谁来决定 token 的生死

DynamicViT 在每个需要做决策的层前面,插入一个轻量的 token 评分与选择模块。这个模块的结构非常克制,通常就是两层 MLP 加一个 softmax,参数量相比主干网络可以忽略不计。

它的工作流程是这样的:假设当前有 N 个 token,每个 token 的特征向量过一个 MLP,输出一个标量分数,表示这个 token 的"保留价值"。然后对所有分数做 softmax 归一化,得到一个概率分布。接着按概率从高到低排序,保留前 K 个(K 是预设的保留比例,比如 0.7),丢掉剩下的。

这里有个工程上很关键的细节:丢掉 token 不是简单地删掉,而是要把被丢 token 的信息"合并"到保留 token 上。因为如果直接删,信息就永久丢失了,深层可能就缺了那块区域的语义。DynamicViT 的做法是,对被丢弃的 token,按它们的分数做加权,把信息聚合到保留的 token 上。这样既减少了 token 数量,又尽量不丢信息。

我一开始觉得这个"合并"步骤有点多余,后来实测发现它很关键。直接硬删的版本,在 ImageNet 上掉了差不多 1.5 个点;加上合并之后,掉点缩小到 0.3 个点以内。这个差距在端侧业务里就是能不能上线的区别。

3.2 逐层递进的丢弃节奏

DynamicViT 不是一次性把 token 砍到位,而是分多个阶段逐步丢。以 ViT-Base 的 12 层为例,它通常在第 3、6、9 层后面各插一个决策模块,每次丢掉一定比例的 token。

为什么要分阶段?因为浅层的 token 特征还不够"语义化",这时候做判断容易误伤。到了中层,语义开始聚合,冗余逐渐显现,这时候丢比较准。而且分阶段丢,模型有缓冲的余地,不会因为一次丢太多而崩掉。

具体比例上,论文里常用的配置是:第一次保留 0.7,第二次保留 0.7,第三次保留 0.7,最终 token 数量变成原来的 0.7³≈0.34。也就是说,最终只有约三分之一的 token 进入深层计算。而深层的计算量是大头,所以整体 FLOPs 能降 30% 到 40%。

这里有个经验:丢弃比例不要设得太激进。我试过第一次就保留 0.5,结果精度掉了 2 个多点。后来改成 0.8、0.7、0.6 这种递减节奏,精度稳住了,FLOPs 也降了差不多 35%。递减的好处是浅层少丢、深层多丢,符合冗余随层数增加的规律。

3.3 训练时的"软硬兼施"策略

DynamicViT 最容易被低估的部分,其实是它的训练策略。如果只是加个预测模块然后端到端训练,模型大概率学不好——因为"保留/丢弃"这个决策是离散的,没法直接反向传播。

论文用了一个很巧妙的松弛 + 退火方案。训练初期,决策是"软"的:每个 token 不是被硬性保留或丢弃,而是按概率加权,所有 token 都参与计算,只是权重不同。这样梯度可以顺畅地传回预测模块。随着训练推进,温度参数逐渐降低,决策越来越"硬",最终逼近真正的二值选择。

这个策略我在复现时深有体会。如果一上来就用硬决策,预测模块根本学不到东西,因为它不知道自己的选择是对是错。软决策相当于给了一个"试错"的缓冲期,让模型慢慢学会哪些 token 值得留。

另外还有一个细节:预测模块的损失函数里加了一个"预算约束"。意思是,模型不能为了保精度而把所有 token 都留下,它必须满足预设的保留比例。这个约束通过一个额外的损失项实现,逼着模型在"保精度"和"省算力"之间找平衡。这个设计很关键,否则模型会倾向于"全保留"来偷懒。

4. 实测数据:到底省了多少,掉了多少

4.1 精度与算力的权衡曲线

我在自己的数据集上复现了 DynamicViT,主干用的是 ViT-Base,输入 224×224,分三个阶段丢弃,保留比例分别是 0.8、0.7、0.6。对比基线是原始 ViT-Base。结果如下:

模型配置Top-1 精度FLOPs(相对值)参数量(相对值)推理延迟(相对值)
ViT-Base 基线81.8%100%100%100%
DynamicViT(0.8/0.7/0.6)81.4%66%100%72%
DynamicViT(0.7/0.7/0.7)80.9%58%100%65%
DynamicViT(0.6/0.6/0.6)79.8%51%100%58%

几个观察值得说。第一,精度掉点确实很小,最温和的配置只掉了 0.4 个点,FLOPs 降了三分之一。第二,参数量基本没变,因为预测模块本身很小,主干参数也没动。标题里说"减少参数量",其实主要指的是如果配合剪枝一起做,或者把预测模块算进去后的整体账,单靠 token 丢弃本身省的是计算量而非参数量。这一点要拎清楚,别被标题误导。

第三,推理延迟的下降幅度小于 FLOPs 的下降幅度。FLOPs 降了 34%,但延迟只降了 28%。原因是 token 丢弃带来的稀疏性,在通用硬件上不一定能完全转化成实际加速,尤其是 batch size 小的时候,GPU 的并行度吃不满。这个问题在端侧反而好一些,因为端侧更吃内存带宽和实际计算量。

4.2 不同主干上的适配表现

DynamicViT 不是只能用在 ViT-Base 上。我还在 DeiT-Small 和 Swin-Tiny 上试过,结论是:模型越小,能省的冗余越少,掉点越敏感。

DeiT-Small 本身参数就少,冗余空间没那么大,用 0.8/0.7/0.6 的配置,精度掉了 0.9 个点,FLOPs 只降了 25%。Swin 因为本身有层级结构,token 数量在浅层就比较大,DynamicViT 插进去反而效果不错,FLOPs 降了 30%,精度掉 0.5 个点。

所以选型建议是:大模型、冗余多的模型,收益更明显;小模型要谨慎,别为了省算力把精度搭进去。如果你的端侧场景对精度极其敏感,宁可少丢一点,保留比例设高一些。

4.3 和知识蒸馏、量化能不能叠加

这是工程上最关心的问题:DynamicViT 能不能和其他压缩手段叠加?答案是能,但要注意顺序和相互影响。

我试过 DynamicViT + 量化(INT8)。单独量化掉 0.6 个点,单独 DynamicViT 掉 0.4 个点,两个一起上掉了 1.1 个点,比简单相加多一点。原因是 token 丢弃后,保留下来的 token 特征分布发生了变化,量化的校准如果还是按原始分布做,就会有偏差。解决办法是在 DynamicViT 训练完之后,再用量化感知训练(QAT)微调几轮,让量化校准适配新的特征分布,这样能把叠加损失压到 0.7 个点以内。

和知识蒸馏叠加则比较友好。用一个大的 ViT 当教师,DynamicViT 当学生,蒸馏损失加上去之后,精度甚至能反超单独训练的 DynamicViT 0.2 个点。因为教师模型提供的软标签,帮学生更好地保留了关键 token 的语义。

5. 复现与改造中的踩坑记录

5.1 预测模块的位置放错了

我第一次复现的时候,把预测模块放在了 MHSA 之后、MLP 之前。结果训练很不稳定,精度波动很大。后来对照论文才发现,预测模块应该放在整个 Transformer block 之后,也就是 MHSA 和 MLP 都算完之后,再决定丢哪些 token。

原因其实不难理解:MHSA 和 MLP 是协同工作的,MHSA 做 token 间的信息交互,MLP 做逐 token 的特征变换。如果在中间插一刀,等于把交互和变换割裂了,后面的 MLP 拿到的 token 集合和前面 MHSA 的输出对不上,信息流就断了。放在 block 之后,每个 token 都完整地过了一遍,特征更成熟,判断也更准。

这个坑我踩了差不多一周,训练 loss 一直下不去,还以为是学习率的问题。后来把位置一改,loss 曲线立刻正常了。所以复现论文的时候,模块插入位置这种细节一定要抠准,差一层效果差很多。

5.2 温度退火太快导致决策崩塌

软决策到硬决策的退火过程,温度参数怎么降,是个很讲究的事。我一开始图快,把退火周期设得很短,结果模型还没学会怎么判断,决策就变硬了,导致大量 token 被错误丢弃,精度直接崩到 70% 以下。

后来我把退火周期拉长到总训练轮数的 60% 左右,也就是前 60% 的 epoch 用软决策慢慢学,后 40% 逐渐硬化。这样精度就稳住了。经验值是:退火不能早于训练中段,最好留足一半以上的时间给软决策阶段。另外温度的下限也别设太低,留一点随机性反而有助于泛化。

5.3 保留比例和 batch size 的联动

还有一个容易被忽略的点:token 丢弃后的实际 batch 内 token 数量是不固定的,因为每张图丢的比例虽然一样,但丢的位置不同。这会导致某些实现里,batch 内的张量形状对不齐,需要 padding 或者用掩码处理。

我在小 batch(比如 8)下训练时,这个问题不明显;换到大 batch(64)时,显存反而没省下来多少,因为 padding 的开销上去了。解决办法是用 gather 操作按索引收集保留的 token,而不是用掩码,这样张量是紧凑的,没有浪费。但 gather 本身有开销,需要权衡。实测下来,batch size 在 32 左右时,收益和开销的平衡最好。

5.4 端侧部署时的算子支持问题

训练完了要部署,又是一道坎。DynamicViT 里的 gather、排序、top-k 这些操作,在服务器 GPU 上没问题,但到了某些端侧推理框架里,支持得并不好。我遇到过 top-k 算子缺失,只能自己用排序加切片模拟,性能打了折扣。

后来我的做法是:在导出模型时,把动态的 top-k 换成固定索引的 gather。具体来说,训练完之后,对每个决策层统计一下 token 保留位置的分布,如果发现某些位置几乎总是被保留(比如目标区域),就可以固化下来,变成一个静态的索引表。这样虽然损失了一点动态性,但换来了端侧的兼容性和速度。精度上大概再掉 0.2 个点,但推理速度能提升 15% 以上,这笔账是划算的。

6. 这套思路还能怎么用

DynamicViT 给我的最大启发,其实不是它具体省了多少算力,而是它示范了一种**"让模型自己学会在哪里省"**的范式。这个思路可以迁移到很多地方。

比如视频理解。视频比图像冗余得多,相邻帧之间大量内容重复。可以借鉴 DynamicViT 的思路,让模型动态决定哪些帧、哪些空间区域需要精细计算,哪些可以粗略处理。我试过一个简化版,在视频分类任务上,把冗余帧的 token 合并,FLOPs 降了 40%,精度只掉 0.8 个点。

再比如多模态模型。图像 token 和文本 token 的重要性差异很大,有些图像 token 对文本问答毫无贡献。用类似的动态选择机制,可以在多模态推理时砍掉大量无关的视觉 token,这对降低多模态模型的服务成本很有意义。

还有一个方向是和硬件协同设计。DynamicViT 的稀疏性是动态的,通用硬件吃不满。如果针对这种"动态 token 数量"设计专门的推理芯片或者调度策略,实际加速比还能再上一个台阶。这块目前工业界和学术界都在探索,落地空间很大。

最后说句实在话:DynamicViT 不是银弹,它解决的是"token 层面冗余"这一个特定问题。如果你的瓶颈在别的地方,比如卷积层的计算、内存带宽、或者后处理,那它帮不上忙。但如果你做的正好是视觉 Transformer 的推理加速,而且对精度敏感,那这套机制值得认真试一试。我自己的体会是,先把保留比例设保守一点跑通,再逐步加码,比一上来就追求极致压缩要稳得多。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进