Chasing Sparsity in Vision Transformers: An End-to-End Exploration 解读

一、论文基本信息论文题目Chasing Sparsity in Vision Transformers: An End-to-End Exploration作者Tianlong Chen、Yu Cheng、Zhe Gan、Lu Yuan、Lei Zhang、Zhangyang Wang发表会议NeurIPS 2021方法名称主要包括三部分SViTESparse Vision Transformer ExplorationS²ViTEStructured Sparse Vision Transformer ExplorationSViTESparse Vision Transformer Co-Exploration官方代码VITA-Group/SViTE。GitHub 页面说明这是 NeurIPS 2021 论文的代码并给出了 SViTE-Small、SViTE-Base 等训练命令。(GitHub)这篇论文的核心目标是不要先训练一个完整 dense ViT再做后剪枝和重训练而是在训练过程中就动态探索稀疏连接最终直接得到一个稀疏 ViT。二、这篇论文要解决什么问题传统剪枝流程通常是先训练 dense model。再剪掉不重要权重、head 或 channel。再 fine-tune / retrain 恢复精度。这个流程在 CNN 里已经比较常见但对 ViT 来说成本很高。因为 ViT 本身训练成本就大如果再做多轮剪枝和重训练整体训练预算会更重。论文开头明确指出传统 post-training pruning 虽然能降低推理复杂度但往往需要额外重训练而 ViT 的一次完整训练已经很慢、很不稳定因此后剪枝成本更难接受。所以这篇论文想解决的问题是能不能从训练一开始就只训练一个稀疏子网络并且在训练过程中动态调整稀疏连接最终直接得到一个高质量 sparse ViT这就是标题里End-to-End Exploration的含义稀疏性不是训练结束后的补丁而是贯穿整个训练过程。三、核心思想这篇论文的核心思想可以概括为在训练过程中同时优化模型参数和稀疏结构先随机初始化一个稀疏 ViT训练若干步后剪掉当前不重要连接再根据梯度长出新的连接如此循环直到训练结束。它借鉴的是动态稀疏训练思想也就是 sparse training / dynamic sparse training。论文把这个框架迁移到 Vision Transformer 上并进一步扩展出结构化稀疏和 token 稀疏。具体有三个层次SViTE非结构化权重稀疏。训练过程中动态 prune-and-grow 单个权重连接。S²ViTE结构化稀疏。动态 prune-and-grow attention heads 和 MLP neurons更硬件友好。SViTE模型稀疏 数据稀疏联合。在稀疏 ViT 的基础上再加入可学习 token selector选择最有信息量的 patch tokens进一步减少中间特征计算。论文摘要也明确说SViTE 会引入 learnable token selector 来选择当前最重要的 patches。所以这篇论文不是单一剪枝方法而是一次比较系统的探索从 dense 到 sparse。从 unstructured 到 structured。从 model sparsity 到 model data sparsity。四、它“剪”的是什么这篇论文里有三种不同粒度的稀疏。4.1 SViTE单个权重连接SViTE 做的是unstructured weight sparsity。它剪的是 Transformer 里的单个权重元素例如 Q/K/V projection、MLP 里的线性层权重等。论文说明ViT 的 first linear projection layer 和 classifier 不参与 sparsification主要 sparsify Transformer 层内部参数。这种方法可以减少参数量和理论 FLOPs但因为是非结构化稀疏实际硬件加速不一定明显。普通 GPU dense matmul 通常不能自动利用零散 0 权重。4.2 S²ViTEattention head 和 MLP neuronS²ViTE 做的是structured sparsity。它不再剪单个权重而是剪更大的结构单元self-attention heads。MLP neurons。这更接近结构化剪枝能带来更实际的运行时间减少。论文也强调结构化稀疏更硬件友好并用一阶重要性近似来指导 attention heads 的 prune-and-grow。4.3 SViTEpatch tokensSViTE 在模型稀疏之外还做data sparsity。这里的数据稀疏不是减少训练样本数量而是减少输入图像中的patch token embeddings。它用一个 learnable token selector 给 patch tokens 打分然后用 top-k 选择最重要 tokens。论文算法 1 中写得很清楚每次训练迭代先对 input token embeddings 打分并选择 top-k informative tokens然后再训练稀疏 ViT 或更新稀疏连接。所以 SViTE 和你前面看的 DynamicViT、ATS、EViT 有联系都是减少视觉 tokens但 SViTE 的重点是把 token sparsity 和 architecture sparsity 放在同一个 sparse training 框架里。五、SViTE动态稀疏训练怎么做SViTE 不是把 dense ViT 训练完再剪而是从一开始就训练一个稀疏网络。整体流程是第一初始化一个随机稀疏 ViT。第二训练 ΔT 个 iteration。第三根据权重幅值剪掉一部分连接。第四根据梯度幅值长出同样数量的新连接。第五重复这个 train–prune–grow 过程直到稀疏结构停止更新。论文图 1 对这个流程做了总结上半部分展示了训练若干步后执行 prune-and-grow再继续训练底部左边展示 Transformer 层内部的结构化或非结构化稀疏底部右边展示 SViTE 的 token selection。这里的关键是稀疏率固定但连接位置会变。例如目标是 50% sparsity那么整个训练过程中始终只有 50% 权重是活跃的但哪些权重活跃会在训练过程中不断更新。六、SViTE 的 prune-and-grow 规则SViTE 的稀疏训练有四个关键因素sparsity distribution。update schedule。pruning criterion。grow criterion。论文使用 Erdős–Rényi 类型的稀疏分布让不同层按照规模分配不同稀疏率更新策略中包含 update interval ΔT、停止更新的 Tend、初始可变连接比例 α 和 cosine decay schedule。具体到 prune 和 growPrune按权重幅值剪。每次 connectivity update 时删除 layer-wise 权重幅值最小的一部分参数。Grow按梯度幅值长。在被剪掉的位置之外选择梯度幅值最大的连接重新激活。这个设计的直觉是小幅值权重当前贡献小可以删。大梯度位置说明如果激活可能最有利于降低 loss应该长出来。新长出的连接初始化为 0数量和被剪连接数量相同因此总稀疏率保持不变。论文明确写到剪掉的 mask 和 grow 的 mask 非零数相等整体非零元素数量固定。七、S²ViTE为什么要做结构化稀疏SViTE 的非结构化稀疏可以显著减少参数和理论 FLOPs但实际 running time 不一定明显下降。因为普通硬件对不规则稀疏支持有限。所以论文进一步提出S²ViTE把动态稀疏训练扩展到结构化单元attention head。MLP neuron。S²ViTE 的 prune-and-grow 仍然是动态的只是对象从“单个权重”变成了“结构单元”。八、S²ViTE 怎么判断 head 重要性S²ViTE 用一阶 Taylor 近似估计 attention head 的重要性。直观上它关心的是如果移除某个 head会让 loss 增加多少这个影响可以用 head 输出和 loss 对 head 输出的梯度之间的乘积来近似。论文在公式 1 中定义了 head importance proxy并说明每次 topology update 时删除重要性最小的 attention heads。对于 MLP neurons论文使用该 neuron 相关权重向量的 L1 norm 作为重要性分数grow 时则选择梯度幅值最大的 head 或 neuron 重新激活。所以 S²ViTE 不是静态 head pruning而是训练过程中反复剪掉低重要 head / neuron再长出高梯度 head / neuron。九、SViTE数据稀疏和架构稀疏联合SViTE 进一步考虑一个问题ViT 的计算不仅来自模型权重也来自输入 token 数。因为 ViT 每层都处理 patch tokenstoken 数越多中间特征越大attention 和 MLP 计算也越多。论文指出由于 skip connectionsinput tokens 数量会决定 intermediate features 的维度因此减少输入 token embeddings 可以直接压缩中间特征并带来效率收益。SViTE 的做法是先用 token selector 给 patch tokens 打分。通过 learnable top-k selection 选择重要 tokens。只把这些 tokens 送入稀疏 ViT。训练时 top-k selector 使用类似 Gumbel trick straight-through 的方式让选择过程可以训练。论文算法 2 给了 PyTorch-like 的 top-k selector 代码对 logits 加 Gumbel noise、除以温度 τ、softmax再用 top-k 和 straight-through 构造硬选择。这和 DynamicViT 的 prediction module 很像但 SViTE 的目的不是单独做 token pruning而是把 token sparsity 和 sparse architecture exploration 结合起来。十、它和 Lottery Ticket / Sparse Training 的关系这篇论文和 Lottery Ticket Hypothesis 有直接关系。Lottery Ticket 的核心发现是一个 dense network 中存在可以单独训练到高精度的 sparse subnetwork。但传统 lottery ticket 需要先训练 dense model再迭代剪枝寻找 winning ticket仍然很贵。SViTE 的思路是不要先训练 dense ViT 再找 ticket。而是在训练过程中动态探索 sparse topology直接训练 sparse ViT。所以它更接近 RigL、SET 这类 dynamic sparse training而不是传统 pruning。这也是它标题里End-to-End Exploration的含义稀疏结构不是训练后挖出来而是在训练过程中边学边找。十一、实验设置论文主要在ImageNet-1K上实验backbone 包括DeiT-TinyDeiT-SmallDeiT-Base训练基本沿用 DeiT 设置包括 AdamW、cosine decay、warmup、weight decay、label smoothing、data augmentations 等但为了更好探索稀疏连接所有实验都训练600 epochs。不同 backbone 的 update schedule 不同例如 DeiT-Tiny 的 ΔT 是 20000DeiT-Small 是 15000DeiT-Base 是 7000。对比方法包括OMPone-shot magnitude pruning。GMPgradual magnitude pruning。TPTaylor pruning。Small-Dense参数量相近的小 dense 模型。SSPstructured sparsity baseline。十二、SViTE 非结构化稀疏结果12.1 DeiT-TinyDense DeiT-Tiny5.72M 参数72.20% Top-1。SViTE-Tiny 30% sparsity4.02M 参数25.56% FLOPs saving71.78% Top-1。SViTE-Tiny 40% sparsity3.46M 参数34.16% FLOPs saving71.75% Top-1。同样稀疏率下OMP、GMP、TP 都明显低于 SViTE。例如 40% sparsity 时OMP 66.52、GMP 68.36、TP 65.45而 SViTE 仍有 71.75。这说明动态稀疏训练比训练后剪枝更适合保持 ViT 精度。12.2 DeiT-SmallDense DeiT-Small22.1M 参数79.90% Top-1。SViTE-Small 50% sparsity11.1M 参数46.26% FLOPs saving79.72% Top-1。SViTE-Small 60% sparsity8.9M 参数55.44% FLOPs saving79.41% Top-1。同样 50% sparsity 下OMP 76.32、GMP 76.88、TP 76.30远低于 SViTE 的 79.72。这说明DeiT-Small 可以剪掉一半参数和接近一半 FLOPs精度几乎不掉。12.3 DeiT-BaseDense DeiT-Base86.6M 参数81.80% Top-1。SViTE-Base 50% sparsity43.4M 参数47.95% FLOPs saving81.51% Top-1。SViTE-Base 60% sparsity34.8M 参数57.50% FLOPs saving81.28% Top-1。这说明更大的 ViT 也存在大量参数冗余而且动态稀疏训练可以在较高稀疏率下保持性能。十三、S²ViTE 结构化稀疏结果结构化稀疏结果更关注真实 running time。在DeiT-Small上Dense DeiT-Small79.90% Top-1。S²ViTE-Small 40% structured sparsity14.6M 参数31.63% FLOPs saving22.65% running time reduction79.22% Top-1。在DeiT-Base上Dense DeiT-Base81.80% Top-1。S²ViTE-Base 40% structured sparsity56.8M 参数33.13% FLOPs saving24.70% running time reduction82.22% Top-1。这个结果很关键S²ViTE-Base 不仅更省而且 Top-1 还超过 dense DeiT-Base。论文总结说S²ViTE 在 30%–40% structured sparsity 下能带来 23.79%–33.63% FLOPs saving并有 10.57%、22.65%、24.70% 的 running time reduction。这说明结构化动态稀疏训练可以带来真实速度收益而不只是理论 FLOPs 下降。十四、SViTE模型稀疏 token 稀疏结果SViTE 的代表实验是 DeiT-Small。在50% unstructured architecture sparsity下保留 100% tokens46.26% FLOPs saving79.72% Top-1。保留 95% tokens49.32% FLOPs saving4.40% time reduced80.18% Top-1。保留 90% tokens52.38% FLOPs saving7.63% time reduced79.91% Top-1。保留 70% tokens63.95% FLOPs saving19.77% time reduced77.90% Top-1。最有意思的是保留 95% tokens 时Top-1 从 79.72 提升到 80.18甚至超过原始 dense DeiT-Small 的 79.90。论文摘要也强调SViTE 在 DeiT-Small 上使用 5% data sparsity 50% architecture sparsity 时Top-1 提升 0.28%同时节省 49.32% FLOPs 和 4.40% running time10% data sparsity 时也没有精度下降并节省 52.38% FLOPs 和 7.63% running time。这就是论文说的sparsity as a free lunch适度稀疏不仅不伤精度反而像正则化一样提高泛化。十五、可视化发现论文有两个重要可视化结论。第一attention heads 存在结构冗余。Figure 4 可视化了 dense DeiT-Base、S²ViTE-Base 和 SViTE-Base 的 attention maps。论文观察到多个 attention heads 行为相似说明存在 head-level 结构冗余S²ViTE 能在一定程度上删掉不必要 heads。第二SViTE 学到的 token selector 会保留物体区域。Figure 5 展示了 SViTE-Small 的 patch selection patterns。论文观察到被删除的无用 patches 通常分布在主体周围或背景中而目标物体内部 patches 大多被保留。这说明 token selector 不是随机删 patch而是在学视觉相关性。十六、和 DynamicViT 的区别DynamicViT 是token sparsification方法。它的重点是通过 prediction module 动态删除冗余 tokens减少后续层计算。这篇论文更全面SViTE动态非结构化权重稀疏。S²ViTE动态结构化 head / MLP neuron 稀疏。SViTE再加入 token selector 做 data sparsity。所以 DynamicViT 主要解决token 冗余SViTE 系列解决的是模型连接冗余 token 冗余。另外DynamicViT 通常从已有 dense ViT 出发进行 token pruning 训练SViTE 的核心是从训练开始就训练 sparse subnetwork。十七、和 ViT-Slim 的区别ViT-Slim 是通过可学习 soft masks 和 L1 稀疏约束做连续空间多维搜索然后把搜索出的结构重新训练。SViTE 系列是动态稀疏训练ViT-Slimmask-based continuous search retrain。SViTEtrain–prune–grow 动态稀疏训练不先训练 dense model。ViT-Slim 最终得到比较规整的维度变窄结构SViTE 的非结构化版本得到 sparse weight topologyS²ViTE 得到结构化 head/neuron 稀疏。所以两者都探索 ViT 稀疏但机制不同ViT-Slim 更像结构化搜索。SViTE 更像 sparse training / dynamic sparse topology exploration。十八、和 ToMe / ATS / TokenLearner 的区别ToMe合并相似 tokens通常 training-free。ATS用 class attention × value norm 做 adaptive token sampling无额外参数。TokenLearner学习生成少量新 tokens。SViTE学习 top-k token selector但它不是 standalone token reduction 插件而是和 sparse model training 联合。所以ToMe / ATS / TokenLearner 主要关注 data/token 维度。SViTE 系列同时关注 architecture/model sparsity 和 data/token sparsity。十九、它是不是剪枝严格说SViTE 不是传统后剪枝而是动态稀疏训练。但它确实属于广义模型稀疏 / 模型压缩SViTE非结构化动态稀疏。S²ViTE结构化动态稀疏。SViTE稀疏训练 token selection。如果写综述建议把它放在Sparse training for ViTs / End-to-end sparse exploration。而不是简单放在 post-training pruning 或 token pruning 里。二十、方法优点第一不需要先训练 dense ViT 再剪。从一开始就训练 sparse subnetwork可以降低训练期间的参数和内存负担。第二动态探索稀疏拓扑。不是固定随机稀疏结构而是训练中根据 weight magnitude 和 gradient magnitude 不断 prune-and-grow。第三同时探索非结构化、结构化、token 稀疏。论文从 SViTE 到 S²ViTE 再到 SViTE覆盖了权重、head/neuron、patch token 三个层面。第四实验结果强。例如 SViTE-Small 50% sparsity 以 11.1M 参数和 46.26% FLOPs saving 达到 79.72% Top-1S²ViTE-Base 40% structured sparsity 达到 82.22% Top-1同时节省 33.13% FLOPs 和 24.70% running time。第五适度稀疏可提升泛化。SViTE 在 50% architecture sparsity 5% token sparsity 下Top-1 达到 80.18比 dense DeiT-Small 的 79.90 更高。二十一、方法局限第一非结构化 SViTE 的真实速度收益有限。单个权重稀疏虽然减少 FLOPs 和参数但普通 GPU 不一定能利用不规则稀疏。论文表中 SViTE unstructured 的 running time reduction 通常记为 0%结构化 S²ViTE 才有明显 running time reduction。第二训练周期长。论文实验训练 600 epochs比常规微调或 training-free token reduction 方法重很多。第三动态稀疏训练超参数较多。包括 sparsity distribution、update interval ΔT、Tend、α、decay schedule、prune/grow criterion、batch size 等。论文也专门做了 update interval 和 batch size 的消融说明这些因素会明显影响结果。第四SViTE 的 token selector 只保留少量 token 稀疏时效果最好。在 50% architecture sparsity 下保留 95% 或 90% tokens 很稳但只保留 70% tokens 时精度降到 77.90说明 aggressive data sparsity 会损伤性能。第五主要验证 ImageNet 分类。论文主要在 DeiT-Tiny/Small/Base 和 TNT-S 上验证。检测、分割、视频、现代大规模预训练 ViT 上是否保持同样结论需要额外实验。二十二、整体评价这篇论文的核心贡献是把 ViT 压缩从“训练后剪枝”推进到“训练全过程稀疏探索”。它不是问训练好的 ViT 里哪些东西可以删而是问ViT 能不能从一开始就以稀疏结构训练并在训练中动态寻找更好的稀疏拓扑这个视角很重要。因为 ViT 的训练成本很高如果必须先训练 dense model 再 prune/retrain那么压缩本身会变得很昂贵。SViTE 的意义就在于它试图把训练和压缩合并成一个过程。从方法谱系看它处在一个很特殊的位置不是单纯 token pruning。不是单纯 head pruning。不是单纯 NAS。而是 dynamic sparse training for ViTs。如果和你前面看的方法串起来AutoFormer搜索一个静态 ViT 架构。ViT-Slim用 mask 搜索多维 slim 子结构。DynamicViT / EViT / ATS / ToMe减少 token 序列计算。SViTE从训练开始动态探索稀疏权重、结构和 token。所以它是 ViT 稀疏训练方向非常代表性的早期论文。二十三、一句话总结《Chasing Sparsity in Vision Transformers: An End-to-End Exploration》提出 SViTE 系列方法把稀疏性贯穿 ViT 的整个训练过程SViTE 从随机稀疏 ViT 出发通过“训练—按权重幅值剪枝—按梯度幅值生长”动态探索非结构化稀疏连接S²ViTE 将该思想扩展到 attention heads 和 MLP neurons 的结构化稀疏SViTE 再加入 learnable top-k token selector同时探索模型稀疏和数据/token 稀疏。实验表明适度稀疏不仅能显著减少 FLOPs 和参数有时还能像正则化一样提升 ViT 泛化是 ViT 动态稀疏训练方向的重要代表。