ToaSt: Token Channel Selection and Structured Pruning for Efficient ViT

ToaSt 提出了一种解耦框架，通过结合多头自注意力模块的耦合头部结构化剪枝与针对前馈网络的 Token 通道选择（TCS）策略，在显著降低计算成本的同时提升了 ViT 模型的精度与效率。

要点

想象一下，你正在管理一家超级繁忙的大型咨询公司（这就是现在的 AI 模型，比如 Vision Transformer，简称 ViT）。这家公司非常聪明，能看懂各种图片、识别物体，甚至能写诗。但是，它有一个大问题：太烧钱了，而且太慢了。

为什么慢？因为这家公司的运作方式有两个“吞金兽”：

1. 全员大会（Self-Attention）：每次开会，公司里的每一个员工（Token，即图片的碎片）都要和所有其他员工一对一交流。如果公司有 1000 个员工，就要开 100 万次对话！这太浪费时间了。
2. 超级大脑（Feed-Forward Network, FFN）：每个员工在开会后，都要去一个巨大的“思考室”进行深度加工。这个思考室非常大，占据了公司 60% 以上的电费（计算量）。

以前的解决办法要么是把员工数量强行减少（Token 压缩），要么是把整个部门裁掉一部分（权重剪枝）。但前者会导致信息丢失（像把重要文件扔了），后者则需要重新培训整个公司，耗时几个月，成本极高。

ToaSt 是什么？
这篇论文提出了一种叫 ToaSt（你可以想象成“烤面包机”，把面包烤得又脆又香，去掉了多余的水分）的新方法。它的核心理念是：“对症下药，互不干扰”。它把公司的两个“吞金兽”分开处理，用两种不同的策略来裁员和提效，而且不需要重新培训（或者只需要极少量的微调）。

---

策略一：给“全员大会”做精准瘦身（MHSA 结构化剪枝）

问题：在“全员大会”环节，每个员工都有很多个“小助手”（Head，即多头注意力机制）来帮忙分析。以前大家觉得，只要把某些小助手裁掉就行，但结果发现，如果裁得不整齐，整个会议就乱套了，因为大家交流的信息对不上了。

ToaSt 的解法：
想象一下，每个小助手手里都拿着一叠卡片（权重矩阵）。ToaSt 发现，这些卡片是成对绑定的（比如“提问卡”和“回答卡”必须对应）。

• 以前的做法：乱砍，导致“提问卡”没了，“回答卡”还在，逻辑不通。
• ToaSt 的做法：它像是一个严格的纪律委员。它发现某些小助手其实是在“摸鱼”（冗余），于是它成对地把这一组小助手对应的卡片全部剪掉。
• 关键点：它只剪掉内部的小卡片，不改变员工之间的交流接口。就像把会议室的椅子数量减少了，但大家还是能正常说话，只是每个人手里的笔记变薄了。
• 效果：会议开得更快了，而且因为去掉了那些只会重复别人话的“废话小助手”，会议质量反而更高了（准确率提升）。

策略二：给“超级大脑”做智能筛选（Token 通道选择 TCS）

问题：这是最烧钱的部分（占 60% 以上）。员工进入“思考室”后，信息会被放大 4 倍（从 D 变成 4D），然后再压缩回来。研究发现，在这个放大的过程中，大部分通道其实是在制造噪音，或者是在重复已经有的信息。就像你让 100 个人写报告，结果 90 个人写的都是废话。

以前的做法：要么全删（太危险），要么重新训练（太慢）。

ToaSt 的解法：
ToaSt 发明了一种**“无需培训的智能筛选器”**。

• 观察：它发现，越往后的“思考室”（深层网络），废话越多，而且这些废话之间有很强的线性关系（你写 A，我就能猜出 B）。
• 操作：它不需要重新培训员工。在员工进入思考室的一瞬间，它会根据当前的任务（比如这张图里有猫还是狗），动态地决定保留哪些通道，扔掉哪些通道。
  • 就像在食堂打饭，ToaSt 不是把整个食堂拆了，而是根据你今天的胃口，只给你盛你爱吃的菜，把那些你不爱吃的、重复的菜（冗余通道）直接过滤掉。
• 神奇之处：它发现，扔掉这些“废话通道”后，不仅省了电费，饭的味道反而更好了（准确率提升）。因为它过滤掉了干扰项，让模型更专注于真正重要的特征。

---

为什么 ToaSt 这么厉害？

1. 不用“回炉重造”：以前的剪枝方法，裁掉人后，公司得花几个月重新培训才能恢复业绩。ToaSt 因为剪得很有章法（成对剪、动态选），裁掉后几乎不需要培训，或者只需要极短的时间（比如大模型只需要微调 15 个 epoch，而以前可能需要 300 个！）。
2. 越大的模型越受益：有趣的是，公司越大（模型参数越多），里面的“废话”和“冗余”就越多。所以，ToaSt 在超大的模型（如 ViT-MAE-Huge）上效果最惊人，不仅省了 40% 的计算量，准确率还提升了 1.64%。
3. 不仅限于看图：这个方法不仅让图片分类变快了，连让 AI 去找物体（目标检测）的任务也变强了。就像给侦探配了更高效的助手，破案率（mAP）反而提高了。

总结

ToaSt 就像是一个精明的公司顾问。
它告诉 AI 模型：“别搞那些虚头巴脑的全员大会和过度思考了。”

• 它把会议开得更精简（成对剪枝），保证沟通顺畅。
• 它把思考变得更聚焦（动态选通道），只保留最有用的信息。
• 最重要的是，它不需要你重新招聘和培训，直接就能让公司运转得更快、更聪明、更省钱。

这就是为什么 ToaSt 能让那些原本只能在超级计算机上跑的 AI 模型，未来有可能轻松跑在你的手机或笔记本电脑上。

技术摘要

1. 研究背景与问题 (Problem)

Vision Transformers (ViTs) 在计算机视觉任务中取得了巨大成功，但其高昂的计算成本限制了在资源受限环境（如移动设备、边缘计算）中的部署。ViT 的计算复杂度主要来源于两个方面：

1. 自注意力机制 (Self-Attention)：具有 $O(N^2)$ 的二次复杂度（$N$ 为序列长度）。
2. 前馈网络 (FFN)：包含线性投影操作，其计算量与隐藏层维度 $D$ 和 $D_{mlp}$ 相关。在标准 ViT 架构中，FFN 层贡献了约 61% 的总 FLOPs，而注意力机制仅占约 19%。

现有方法的局限性：

• 结构化权重剪枝 (Structured Weight Pruning)：通常针对注意力机制，但往往需要漫长的重新训练（Retraining），且容易忽略 FFN 层中巨大的冗余。
• Token 压缩 (Token Compression)：通过减少序列长度 $N$ 来降低注意力复杂度。然而，这类方法仅线性减少 FFN 的计算量，无法解决 FFN 内部 $O(D^2)$ 的通道冗余问题。此外，Token 压缩决策会在层间传播，导致全局依赖，增加了优化难度。

核心挑战：如何在不进行昂贵重新训练的前提下，同时有效压缩 MHSA（多头自注意力）和 FFN（前馈网络）中的冗余，并避免层间依赖带来的优化困难。

---

2. 方法论 (Methodology)

作者提出了 ToaSt，一个解耦 (Decoupled) 的压缩框架，采用“层独立压缩”理念，针对 ViT 的不同组件应用专用策略：

2.1 针对 MHSA 的耦合结构化权重剪枝 (Structured Coupled Weight Pruning)

• 目标：减少每个注意力头的内部维度 $d_k$，而非全局嵌入维度 $D$，以保持残差连接和层间接口的兼容性。
• 耦合约束：为了保持数学完整性，剪枝必须同步进行：
  • Q-K 同步：剪枝 $W_Q$ 的第 $j$ 列必须同时剪枝 $W_K$ 的第 $j$ 列。
  • V-Proj 同步：剪枝 $W_V$ 的第 $j$ 列必须同时剪枝 $W_{proj}$ 的第 $j$ 行。
• 重要性评估：不使用动态梯度，而是基于预训练权重的几何中值 (Geometric Median, GM)。计算权重向量到几何中值的欧氏距离，距离越近表示冗余度越高（可被其他维度替代）。
• 策略：采用逐头均匀剪枝 (Head-wise Uniform Pruning)，除第一层（处理关键接口）外，其余层进行 90% 的激进剪枝。

2.2 针对 FFN 的 Token 通道选择 (Token Channel Selection, TCS)

• 核心洞察：通过分析预训练 ViT 的 FFN 激活模式，发现深层存在显著冗余：
  1. 高线性重建保真度 ($R^2$)：通道间存在高度线性依赖，少量通道即可重建全局分布。
  2. 有效秩坍缩 (Collapsing Effective Rank)：深层特征的有效维度远低于扩展后的维度。
  3. 稀疏性增加：深层中大量神经元激活值接近零（“死神经元”）。
• 训练免费 (Training-free) 策略：
  • 统计采样：利用高 $R^2$ 特性，仅随机采样少量 Token（2%-20%）来估算通道重要性，大幅降低分析开销。
  • 注意力引导的统一重要性：结合 CLS Token 的全局激活和 Patch Token 的局部激活（加权计算），确定哪些通道需要保留。
  • 层自适应剪枝：
    • FC1 (扩展层)：保守剪枝，保留特征多样性。
    • FC2 (压缩层)：激进剪枝（最高 90%），直接利用深层的高冗余特性。
• 优势：直接减少 $D^2$ 复杂度，且由于是结构化剪枝（整列/整行移除），无需特殊稀疏库，可直接在标准 GPU 上加速。

---

3. 关键贡献 (Key Contributions)

1. 解耦框架设计：提出了 ToaSt，将 MHSA 的权重剪枝与 FFN 的通道选择解耦，避免了全局传播问题，简化了优化景观。
2. MHSA 耦合剪枝：提出了一种基于几何中值的耦合权重剪枝方法，通过同步 Q-K 和 V-Proj 索引，实现了层独立的维度压缩，无需修改 Transformer 块接口。
3. FFN 通道选择 (TCS)：
   • 首次通过实证分析揭示了 FFN 深层的冗余特征（高 $R^2$、低有效秩、高稀疏性）。
   • 提出了一种无需重新训练 (Training-free) 的通道选择方法，利用统计采样和注意力引导机制，有效过滤冗余噪声。
4. 性能与效率的双重提升：证明了该方法不仅能大幅降低计算量，还能通过去除噪声提升模型泛化能力（即“剪枝即正则化”）。

---

4. 实验结果 (Results)

实验在 ImageNet-1K（分类）和 COCO（目标检测）上进行了广泛评估，涵盖 DeiT、ViT-MAE 和 Swin Transformer 等 9 种模型。

• ImageNet-1K 分类性能：
  • ViT-MAE-Huge：在 FLOPs 减少 39.4% 的情况下，Top-1 准确率从 86.88% 提升至 88.52% (+1.64%)。
  • DeiT-Small：FLOPs 减少 45.7%，准确率从 79.82% 提升至 83.40% (+3.58%)，吞吐量提升 2.07 倍。
  • 对比 SOTA：在同等 FLOPs 预算下，ToaSt 的准确率比 ToMe、DiffRate 等 Token 压缩方法高出 1% - 4%。
• 微调效率：
  • 大模型（如 ViT-MAE-Huge）仅需 ~15 个 epoch 的微调即可恢复并超越基线性能，而中小模型（如 DeiT-Base）通常需要 ~290 个 epoch。这表明大模型具有更高的内在冗余，对 ToaSt 更友好。
• 下游任务 (COCO 目标检测)：
  • 使用压缩后的 Swin-Base 作为骨干网络，mAP 从 51.9 提升至 52.2，证明了压缩后的模型在密集预测任务中依然有效，且去除了引入噪声的冗余特征。
• 硬件加速：
  • 在 NVIDIA H100 GPU 上，由于保持了稠密矩阵结构，实现了显著的吞吐量提升（最高 2.07 倍），无需特殊稀疏硬件支持。

---

5. 意义与总结 (Significance)

• 突破 FFN 瓶颈：现有方法多关注注意力机制或序列长度，ToaSt 首次系统性地解决了占计算量 60% 以上的 FFN 通道冗余问题。
• 无需昂贵重训：通过“训练免费”的通道选择和基于统计的剪枝，大幅降低了压缩大模型的门槛和时间成本。
• 正则化效应：实验表明，ToaSt 不仅没有损害精度，反而通过去除冗余噪声提升了模型的泛化能力，这在压缩领域是一个反直觉但极具价值的发现。
• 通用性与可扩展性：该方法适用于多种架构（DeiT, MAE, Swin）和任务（分类、检测），且大模型规模越大，收益越明显，为未来基础模型的高效部署提供了新思路。

总结：ToaSt 通过解耦 MHSA 和 FFN 的压缩策略，利用结构化剪枝和统计通道选择，在显著降低计算成本的同时提升了模型精度，为 ViT 在边缘设备上的高效部署提供了强有力的解决方案。