HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers¶

会议: CVPR 2026
arXiv: 2603.12222
代码: 待确认
领域: 模型压缩 / Vision Transformer剪枝 / 神经架构搜索
关键词: Vision Transformer剪枝, 多粒度结构化剪枝, Gumbel-Sigmoid门控, 端到端子网络搜索, 边缘部署

一句话总结¶

提出HiAP——统一宏观（整头/FFN块）和微观（头内维度/FFN神经元）的层级Gumbel-Sigmoid门控框架，在单次端到端训练中自动发现满足算力预算的高效ViT子网络，无需手动重要性排序或多阶段流程。

背景与动机¶

ViT计算和内存开销大，结构化剪枝是主流压缩手段。但现有方法存在两大痛点：(1) 通常只在单一粒度操作——纯微观剪枝（如ViT-Slim修剪头内维度）虽降FLOP但仍需加载所有层的权重矩阵，内存带宽瓶颈未解；纯宏观剪枝（如UPDP丢弃整个block）会大幅损失表示能力。(2) 依赖复杂多阶段流程：先用手工启发式（如Taylor重要性、图排序）确定剪枝掩码，再单独微调恢复精度，流程繁琐且需专家知识调参。

核心问题¶

能否在单次训练中，让网络自己学会在多粒度上该剪什么、保留什么，无需人为设定每层剪枝比例或重要性度量？

方法详解¶

整体框架¶

HiAP在ViT的每个Transformer block内引入两级可学习的Gumbel-Sigmoid随机门控。训练时，门控logit和网络权重联合优化，温度退火使门控从软连续逐步逼近硬二值决策，训练结束后直接提取物理子网络，无需二次微调。

关键设计¶

层级门控机制: 宏观门控 $g_{l,h}$ 和 $b_l$ 控制整个注意力头和FFN块的保留/移除；微观门控 $d_{l,h,j}$ 和 $c_{l,k}$ 在存活的宏观结构内选择性修剪头内维度和FFN神经元。微观门控受宏观门控约束——当宏观门关闭时，其下所有微观门自动失效
可微算力建模: 将网络的MAC开销线性分解为 $\mathbb{E}[C(\mathcal{G})] = \sum_l \sum_h (C_1 \cdot \mathbb{E}[g_{l,h}] + C_2 \sum_j \mathbb{E}[g_{l,h} \cdot d_{l,h,j}]) + \sum_l \sum_k C_3 \cdot \mathbb{E}[b_l \cdot c_{l,k}]$，其中 $C_1$ 为宏观注意力开销，$C_2$、$C_3$ 为单维度/神经元的微观开销，使优化器能精确归因每个结构的硬件代价
解耦代价惩罚: 将宏观($\mathcal{L}_{macro}$)和微观($\mathcal{L}_{micro}$)的算力惩罚分离，用独立超参数控制，允许显式管理粗细粒度稀疏性的权衡
结构可行性约束: 引入 $\mathcal{L}_{feasibility}$ 防止层坍塌——对每层强制保留最低数量的注意力头和FFN神经元比例，用 ReLU 阈值惩罚实现

损失函数 / 训练策略¶

$$\mathcal{L}_{total} = \mathcal{L}_{task} + \lambda_{macro}\mathcal{L}_{macro} + \lambda_{micro}\mathcal{L}_{micro} + \mathcal{L}_{feasibility}$$ - $\mathcal{L}_{task}$: 交叉熵 + 知识蒸馏（用预训练DeiT-Small做teacher, $\alpha_{KD}=0.7$, $T=4.0$） - Gumbel-Sigmoid温度从 $\tau_0=2.0$ 指数退火至 $\tau_{min}=0.5$，200 epoch训练，AdamW优化器 lr=5e-5

实验关键数据¶

方法	训练epoch	MACs(G)	Top-1 Acc(%)	Δ Acc
DeiT-Small (dense)	-	4.6	79.85	-
ViT-Slim	15.6	3.1	79.90	+0.05
GOHSP	14.4	3.0	79.98	+0.13
S2ViT	15.3	3.1	79.22	-0.63
WDPruning	15.0	3.1	78.55	-1.30
HiAP	15.0	3.1	79.10	-0.75
HiAP (aggressive)	12.3	2.5	77.95	-1.90

HiAP在3.1G MACs（约33%压缩）下达79.1%精度，通过单阶段训练无需多阶段流程
CIFAR-10上：33%压缩下87.56%（优于Uniform-Ratio的86.63%和ℓ1-Structured的87.15%）
实测推理延迟从5.57ms降至3.86ms（1.44×加速），证明物理子网络直接可用

消融实验要点¶

宏微观惩罚比例消融（2:1, 5:1, 1.5:1, 纯宏观, 纯微观）：2:1为DeiT-Small的最佳Pareto均衡
纯宏观penalty导致大量头被移除但FFN神经元几乎未减，分布不均
纯微观penalty间接导致整个MLP block被消除，是一种意外的depth-pruning效果
网络自动发现最后一层FFN block完全冗余（$b_{12}=0$），无需人工指定

亮点¶

框架设计非常干净：层级门控 + 可微预算 + 可行性约束，三板斧解决自动剪枝
单阶段训练的简洁性是核心卖点——对比GOHSP/NViT等需图排序+多轮评估的方法，工程复杂度大幅降低
解耦代价惩罚的设计让剪枝行为可控可分析，不同penalty比例下的结构演化可视化非常有说明力
Proposition 1证明可微预算约束的线性分解不需要门控独立性假设，理论上严谨

局限性¶

目标函数只优化MACs，未建模实际延迟/能耗，论文自己承认MACs-to-latency gap是主要限制
ImageNet上精度与GOHSP/ViT-Slim相比仍有差距（79.1 vs 79.98），代价是换来流程简化
仅在分类任务上验证，检测/分割等密集预测任务适用性未知
温度退火schedule需要调参（$\tau_0$, $\tau_{min}$），对不同模型/任务的敏感性未充分分析

与相关工作的对比¶

ViT-Slim: 仅微观剪枝（头内维度+FFN），需稀疏性ranking+阈值确定，HiAP在多粒度统一框架中端到端生成决策
GOHSP: 使用图排序确定头重要性+优化剪枝组合，流程复杂且需专家调参；HiAP通过Gumbel门控让网络自学
UPDP: 仅宏观层面（FFN block级）剪枝，用遗传算法搜索，与HiAP的可微搜索范式不同

启发与关联¶

Gumbel-Sigmoid门控+温度退火的范式可复用到几乎任何需要学习离散结构选择的场景
解耦宏微观代价的思想可推广到NAS中对不同资源维度的独立控制

评分¶

新颖性: ⭐⭐⭐⭐
实验充分度: ⭐⭐⭐
写作质量: ⭐⭐⭐⭐
价值: ⭐⭐⭐