Vision Transformer Slimming: Multi-Dimension Searching in Continuous Optimization Space 解读

一、论文基本信息论文题目Vision Transformer Slimming: Multi-Dimension Searching in Continuous Optimization Space方法简称ViT-Slim作者Arnav Chavan、Zhiqiang Shen、Zhuang Liu、Zechun Liu、Kwang-Ting Cheng、Eric Xing发表会议CVPR 2022官方代码Arnav0400/ViT-SlimGitHub 页面标注为该 CVPR 2022 论文的官方实现。(GitHub)这篇论文的核心目标是不是只剪 token也不是只剪 head 或 MLP channel而是在 ViT 中同时搜索 patch/token、MHSA、MLP 三个维度的高效子结构。论文摘要明确说ViT-Slim 可以端到端搜索输入 tokens、MHSA 和 MLP 模块并通过统一的 L1 稀疏约束在连续搜索空间中完成高效搜索。二、这篇论文要解决什么问题之前很多高效 ViT 方法只压缩一个维度。例如DynamicViT / Patch Slimming主要减少 patch tokens只降低 FLOPs不减少参数量。Head pruning主要减少 attention heads。MLP pruning主要减少 FFN hidden dimensions。NAS 方法可以搜索结构但很多方法需要训练 supernet、评估大量子网搜索成本很高。ViT-Slim 认为ViT 的冗余不只存在于单一维度而是同时存在于输入 patch/token 维度。MHSA 每个 attention head 的内部维度。MLP hidden dimension。所以它想解决的问题是能否用一次高效搜索同时找到 patch、MHSA、MLP 三个维度上的最优压缩结构这就是标题里Multi-Dimension Searching的含义。三、核心思想ViT-Slim 的核心思想是给 ViT 的不同结构维度加上可学习 soft masks通过 L1 稀疏约束让不重要的维度 mask 变小搜索结束后按 mask 大小排序删除低重要维度再重新训练得到最终小模型。它不是手工设定每层剪多少也不是离散枚举候选结构而是把结构搜索转化为一个连续优化问题。论文明确说它在预训练模型上加入 sparsity masks然后联合优化原模型权重和 mask搜索完成后根据 mask 值排序删除低 rank 的维度最后再 retrain 提取出的压缩结构。所以 ViT-Slim 的流程是第一步在预训练 ViT 上插入 soft masks。第二步用交叉熵损失 L1 mask 稀疏正则进行 one-shot search。第三步根据目标 budget 删除低 mask 值结构。第四步把连续 mask 结构转成真实小模型再重新训练。四、它到底“瘦身”哪些维度ViT-Slim 同时考虑三个维度。1. MHSA 维度在 multi-head self-attention 中每个 head 都有自己的 Q、K、V 投影维度。ViT-Slim 给每一层、每个 head 的内部维度加 mask。这意味着它不是简单“删整个 head”而是可以让不同 head 保留不同维度。例如某些 head 可能保留 64 维某些 head 可能只保留 32 维甚至更少。这比整头剪枝更细。论文说明针对 MHSA 模块ViT-Slim 可以探索的结构空间是每层每个 head 的维度组合因此比预定义离散候选结构更灵活。2. MLP 维度ViT block 中的 MLP 通常是两层全连接hidden dimension 通常是 embedding dimension 的 4 倍。这个部分参数和 FLOPs 很大。ViT-Slim 给 MLP hidden dimension 加 mask从而搜索每层 MLP 到底需要保留多少 hidden neurons。这类似结构化 FFN dimension pruning但它不是单独剪 MLP而是可以和 MHSA、patch selection 联合搜索。3. Patch / token 维度ViT-Slim 也搜索每层哪些 patch tokens 可以删除。论文认为MHSA 会逐层聚合 patch 信息使深层 patch 表示越来越相似因此深层可以删除大量 patch浅层也可以删除少量不重要 patch。论文在方法部分提到patch 之间的余弦相似度会随层数快速上升最终层甚至可达到约 0.9这为删除深层 patch 提供了动机。但 patch selection 有一个特殊点它只降低 FLOPs不减少模型参数量。因为删除 patch 是减少 token 数后续 attention 和 MLP 对 token 的计算变少但权重矩阵本身没有变小。五、为什么叫 continuous optimization space很多 NAS 方法会预先定义候选结构例如MLP hidden dimension 只能从 {128, 256, 384} 中选。head 数只能从 {3, 6, 9} 中选。这种是离散搜索空间。ViT-Slim 不这样做。它给每个维度一个连续 mask 值通过优化 mask 大小来判断重要性。论文解释说Transformer 内部主要由 fully-connected layers 组成可以用共享 super weight matrix 的方式让不同候选输出维度从同一个大矩阵中 slice 出来从而构建连续搜索空间。简单理解传统 NAS在几个候选结构里选。ViT-Slim每个维度都有一个连续重要性分数最后按预算切掉一部分。这种方式的优点是搜索空间更细。不用训练大量候选子网。搜索一次后可以根据不同 budget 取不同子模型。论文也强调搜索完成后可以根据真实设备的精度-FLOPs 需求直接确定最终结构不需要像 SPOS 那样额外评估大量子网。六、L1 稀疏 mask 的作用ViT-Slim 的目标可以直观理解为既要分类准确又要让 mask 尽量稀疏。也就是分类损失 mask 的 L1 稀疏损失。L1 正则会推动一部分 mask 变小。搜索完成后mask 小的维度被认为不重要可以删除mask 大的维度被保留。论文中明确写到搜索目标是根据 mask 对最终性能的影响进行排序然后删除 mask 值最低的维度mask 会乘到对应中间 tensor 上优化目标由分类交叉熵和 mask 的 L1 norm 组成。这和 Network Slimming 中用 BN gamma 做通道重要性很像但 ViT 没有 BN所以论文显式引入 soft masks 来替代 BN scaling parameter。七、搜索和重训练流程ViT-Slim 的完整 workflow 有三步。第一One-shot Searching。用预训练权重初始化 ViT把它当作 supernet然后加入 sparsity masks。搜索时同时训练模型权重和 mask。论文实验中搜索 50 epochsDeiT-S 的 batch size 为 1024DeiT-B 为 512优化器为 AdamW并沿用 DeiT 的数据增强策略。第二Budget Selection。搜索结束后根据 mask 值排序。目标 budget 决定保留多少维度或 patch。比如保留 70% MHSA / MLP 维度或者保留 80% patch。第三Re-training。把连续 mask 转成真正的压缩结构后重新训练 300 epochs。论文解释说这一步让权重从连续 mask 搜索空间适应到最终 binary / non-existing mask 的真实结构。所以 ViT-Slim 不是 training-free pruning。它需要搜索 重训练。八、ViT-SlimPS 和 ViT-SlimJS 是什么论文里有两个主要模型族。ViT-SlimPSPS 可以理解为Partial Search / Parameter Search。它只搜索MHSA 维度。MLP 维度。不做 patch selection。因此它能同时减少参数量。FLOPs。因为 MHSA 和 MLP 的权重矩阵变小了。ViT-SlimJSJS 可以理解为Joint Search。它联合搜索MHSA 维度。MLP 维度。patch selection。因此它不仅减少参数和 FLOPs还进一步通过删除 patches 降低 FLOPs。论文在 Table 8 中也说明PS 是 MHSAMLP 的参数搜索JS 是最终的 joint-dimension searchpatch selection 不影响参数量但能进一步减少 FLOPs。九、主要实验结果9.1 DeiT-S 上的结果原始DeiT-S22.0M 参数。4.6B FLOPs。79.9% Top-1。ViT-SlimPS17.7M 参数。3.7B FLOPs。80.6% Top-1。也就是说ViT-SlimPS 在参数和 FLOPs 都减少的情况下Top-1 反而比原始 DeiT-S 高。ViT-SlimJS15.7M 参数。3.1B FLOPs。79.9% Top-1。或者更激进15.7M 参数。2.8B FLOPs。79.5% Top-1。这说明加入 patch selection 后FLOPs 可以进一步降到 2.8B 左右同时精度损失很小。9.2 单独搜索 MHSA 和 MLP 的结果论文对 MHSA 和 MLP 单独做了搜索。在 DeiT-S 上MHSA 单独 70% budget 时19.9M 参数。4.1B FLOPs。80.90% Top-1。MLP 单独 70% budget 时17.8M 参数。3.8B FLOPs。80.80% Top-1。这说明 MHSA 和 MLP 都有冗余而且适度压缩后甚至能提高精度可能起到了结构正则化作用。但 MLP 压缩更能减少参数因为 ViT 中 MLP 通常占大量参数MHSA 压缩对 attention 表达能力影响更敏感。9.3 联合搜索比单独搜索更有意义论文的结论不是“只剪 MLP 最好”而是不同模块冗余不同联合搜索能找到更平衡的结构。例如 ViT-SlimPS 在 80/80 budget 下达到17.7M 参数。3.7B FLOPs。80.60% Top-1。在 70/70 budget 下15.6M 参数。3.3B FLOPs。80.03% Top-1。在 60/60 budget 下13.5M 参数。2.8B FLOPs。79.20% Top-1。这些结果说明保留 70% 左右的 MHSAMLP 维度已经可以基本保持甚至超过 DeiT-S baseline而过度压缩到 50% 会明显掉点。9.4 DeiT-B 和 Swin 上的结果论文也把方法扩展到 DeiT-B 和 Swin-T。在 DeiT-B 上作者设置 60% budget约等价于40% FLOPs 和参数下降并报告 ViT-Slim 在参数和 FLOPs 更少的情况下Top-1 相比预训练 DeiT-B 提升约0.6%。在 Swin-T 上论文也做了搜索。结果显示 80% budget 可以保持精度但 70% budget 会下降。作者解释说Swin 本身是精心设计的层级结构每层维度已经比较紧凑因此压缩空间小于 isotropic 的 DeiT。这个结论很重要ViT-Slim 对 DeiT 这类等宽 isotropic ViT 特别有效对 Swin 这种已经层级化设计的模型压缩空间相对有限。十、迁移学习结果论文还测试了迁移任务包括 CIFAR-10、CIFAR-100、iNaturalist-2018、iNaturalist-2019。结果显示ViT-SlimPS 在 ImageNet 上搜索并预训练后迁移到这些数据集时仍然优于 DeiT-S baseline。例如 CIFAR-100 从 87.80 提升到 88.16iNaturalist-2019 从 75.35 提升到 76.67。论文强调这些结构是在 ImageNet 上搜索的不是在对应下游数据集上重新搜索说明搜索出的 Slim 结构有一定迁移性。这对模型压缩很重要因为如果每个下游任务都要重新搜索成本会很高。ViT-Slim 的结果说明至少在分类迁移任务上ImageNet 搜索得到的结构可以复用。十一、搜索到的结构有什么规律论文可视化了搜索后的 MHSA head dimensions 和 MLP dimensions。一个重要现象是深层维度被压缩得更多。中间层保留维度更多。前面层有中等程度压缩。论文解释说self-attention 在中间和前面层仍然比较重要因为 patch 还比较 distinct需要信息交换而深层已经学到大量特征冗余更强因此可以保留更小维度。这个现象和你前面看的 Patch Slimming、A-ViT、DynamicViT 的发现有呼应ViT 深层通常存在更强冗余。不过 ViT-Slim 的冗余不是只体现在 token 数上也体现在 MHSA head dimension 和 MLP hidden dimension 上。十二、和 DynamicViT 的区别DynamicViT 主要做动态 token sparsification。它不减少参数量只减少推理时 token 数因此主要降低 FLOPs。ViT-Slim 不同它搜索 MHSA 和 MLP 维度可以减少参数量。它也可以搜索 patch selection进一步减少 FLOPs。它得到的是静态压缩结构而不是每张图动态不同的 token path。论文对比表中也明确指出DynamicViT 是 dynamic patch selection主要减少 FLOPs而 ViT-Slim 可以同时减少参数和 FLOPs。所以DynamicViT输入自适应 token pruning。ViT-Slim多维度结构搜索 静态 slim architecture。十三、和 Patch Slimming 的区别Patch Slimming 主要从最终输出误差出发top-down 搜索每层保留哪些 patch tokens。它同样只减少 FLOPs不减少参数。ViT-Slim 把 patch selection 作为三维搜索中的一部分同时还搜索 MHSA 和 MLP 维度。所以Patch Slimming专注 patch token slimming。ViT-Slimpatch MHSA MLP 联合 slimming。论文也指出 Patch Slimming 和 DynamicViT 在 FLOPs 方面可与 ViT-SlimJS 接近但它们没有参数效率提升因为 patch selection 不减少权重矩阵。十四、和 ToMe / TokenLearner / ATS 的区别ToMe是 token merging可以 training-free 地把相似 tokens 合并。TokenLearner是 learned tokenization从特征图中生成少量新 tokens。ATS是 parameter-free adaptive token sampling根据 class attention 和 value norm 采样 tokens。ViT-Slim则是结构搜索。它不是为了直接 patch 到已有模型做无训练加速而是通过 soft masks 搜索一个真实更小的 ViT 子结构再重新训练。所以 ViT-Slim 更接近structured architecture search / structured pruning。而不是纯 token reduction 插件。十五、它是不是剪枝可以说 ViT-Slim 是一种结构化剪枝 架构搜索方法。它有剪枝味道因为最终会删除MHSA head 内部维度。MLP hidden dimensions。某些 patch tokens。但它又不完全是传统剪枝因为它不是用单一重要性指标直接剪而是把不同维度放到连续搜索空间中用 L1 sparse masks 做 one-shot search。所以最准确的分类是multi-dimensional structured slimming / one-shot differentiable architecture search for ViTs。十六、方法优点第一压缩维度全面。不是只剪 token而是同时考虑 patch、MHSA、MLP 三个维度。第二能减少参数和 FLOPs。MHSAMLP 搜索直接让权重矩阵变小patch selection 进一步减少 token 计算。第三搜索效率较高。论文报告 DeiT-S 搜索约 43 GPU hoursDeiT-B 约 71 GPU hours相比很多 NAS 方法的数百 epoch 或大量子网评估更轻。第四搜索一次可支持多个 budget。搜索结束后mask 排序已经形成可以根据不同 budget 切出不同子模型不需要为每个预算重复搜索。第五压缩后结构是真实小模型。不像 mask-only 稀疏方法依赖稀疏 kernelViT-Slim 最终得到的是维度变小的 dense 子网络更容易实际部署。十七、方法局限第一需要重新训练。搜索后必须 retrain 300 epochs才能让最终二值结构恢复性能。因此它不是 training-free 方法。第二patch selection 主要是静态的。ViT-Slim 的 patch selection 不是像 DynamicViT / ATS 那样对每张图动态选择。论文自己也提到patch 维度上动态方式更有前景因为不同图像重要区域不同。第三搜索和训练成本仍然不低。虽然比许多 NAS 方法低但 43 GPU hours 搜索 300 epochs retraining 对普通实验仍然比较重。第四主要验证分类。ImageNet 和下游分类任务效果较强但检测、分割等 dense prediction 任务中patch selection 和维度压缩是否保持空间细节需要额外验证。第五和现代大规模预训练 ViT 的关系需要重新评估。论文主要基于 DeiT、Swin 等 2021–2022 年主流模型。对 MAE、DINOv2、CLIP ViT、SAM encoder 等更大规模表示模型是否还能直接获得同样规律需要进一步实验。十八、整体评价ViT-Slim 的核心贡献是把 ViT 压缩从“单维度剪枝”推进到“多维度联合结构搜索”。它不满足于只删 patch也不满足于只剪 head 或 MLP而是同时搜索哪些 patch tokens 可以少算。哪些 attention head 内部维度可以变窄。哪些 MLP hidden dimensions 可以删除。这种思路很符合 ViT 的结构特点因为 ViT 的冗余确实分布在多个位置。只剪一个维度容易造成某个模块过度压缩联合搜索可以在不同模块之间分配压缩预算。从你前面看的 token reduction 方法脉络看ViT-Slim 的位置很清楚DynamicViT / EViT / ATS / ToMe / TokenLearner 更关注 token 数。ViT-Slim 更关注整体结构维度包括 token、attention、MLP。因此它更适合被放在ViT structured pruning / architecture search类而不是单纯的 token pruning 类。十九、一句话总结《Vision Transformer Slimming: Multi-Dimension Searching in Continuous Optimization Space》提出 ViT-Slim通过在 ViT 的 MHSA、MLP 和 patch/token 维度上引入可学习 soft masks并用 L1 稀疏约束进行 one-shot 连续空间搜索自动找到高效子结构搜索完成后按 mask 排序删除低重要维度再重新训练得到真实小模型。它的核心价值是同时减少参数和 FLOPs而不是只减少 token 计算在 DeiT-S 上ViT-SlimPS 可从 22.0M / 4.6B / 79.9% 压缩到 17.7M / 3.7B / 80.6%ViT-SlimJS 进一步通过 patch selection 达到约 15.7M / 3.1B / 79.9%是 ViT 多维度结构化压缩的重要代表。