HiAP

alias:: HiAP title:: “HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers” method_name:: HiAP authors:: Andy Li, Aiden Durrant, Milan Markovic, Georgios Leontidis year:: 2025 venue:: arXiv tags:: structured pruning, vision transformer, Gumbel-Sigmoid, neural architecture search, knowledge distillation image_source:: online arxiv_html:: https://arxiv.org/html/2603.12222 created:: 2026-03-14

• 论文笔记：HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

• 元信息

• 一句话总结

• 提出多粒度 Gumbel-Sigmoid 门控框架 HiAP，在单阶段端到端训练中自动发现 ViT 最优子网络，无需人工重要性启发式或预定义稀疏目标

• 核心贡献

• 多层级门控体系: 统一 macro-level（整个注意力头和 FFN 块）和 micro-level（头内维度和 FFN 神经元）的 结构化剪枝 到单一可微框架

• 预算感知端到端搜索: 网络自主发现并硬化最优子架构，无需手动重要性启发式、代理排序指标或昂贵的二次微调阶段

• 自动结构可行性保障: 通过可行性惩罚防止层坍塌，保证有效的前向传播路径

• 问题背景

• 要解决的问题

• ViT 计算和内存开销大，难以部署到边缘设备

• 现有 结构化剪枝 方法通常只在单一粒度上操作，且依赖多阶段流水线和后处理阈值

• 现有方法的局限

• 单粒度局限: 仅剪枝 micro 结构（如头内维度）可降低 FLOPs，但保留了所有层和注意力头的数量，硬件仍需承担大量内存访问开销（HBM 访问）；仅剪枝 macro 结构（整个头/块）虽能绕过内存传输，但容易导致表征能力严重损失

• 后处理依赖: 现代可微搜索方法常需后处理 magnitude thresholding，需要专家知识和人工干预

• 多阶段流水线: 搜索阶段 + 微调阶段的两阶段流程计算代价高

• 本文的动机

• 将剪枝问题转化为单次预算感知学习问题，让模型自己学会该剪什么

• 同时在 macro 和 micro 两个粒度放置 Gumbel-Sigmoid 门控，通过温度退火自然收敛到离散子架构

• 方法详解

• 模型架构

• HiAP 应用于标准 ViT 架构（以 DeiT-Small 为主要实验对象）：

  • 输入: 图像 patch 序列，长度 $N$，嵌入维度 $D$
  • Backbone: $L$ 层 Transformer，每层 $H$ 个注意力头，头维度 $D_h$，FFN 隐藏宽度 $D_{\text{ffn}}$
  • 核心模块: 层级 Gumbel-Sigmoid 门控系统，分为 macro-gates 和 micro-gates
  • 输出: 分类预测

• 符号定义

• Macro-gates: $g_{l,h} \in \{0,1\}$ 控制第 $l$ 层第 $h$ 个注意力头；$b_l \in \{0,1\}$ 控制第 $l$ 层 FFN 块

• Micro-gates: $d_{l,h,j} \in \{0,1\}$ 控制头内第 $j$ 个维度；$c_{l,k} \in \{0,1\}$ 控制 FFN 第 $k$ 个神经元

• 完整门控集合记为 $\mathcal{G}$

• 核心模块

• 模块1: Macro-Level 剪枝

• 设计动机: 消除整个注意力头和 FFN 块，直接减少内存带宽瓶颈和层级计算开销

• 具体实现:

  • 注意力头门控：当 $g_{l,h}=0$ 时，第 $l$ 层第 $h$ 个注意力头被完全绕过
  • FFN 块门控：当 $b_l=0$ 时，第 $l$ 层整个 FFN 块被移除
  • 利用残差连接保证梯度流在块被剪除后仍可通过

• 模块2: Micro-Level 剪枝

• 设计动机: 在活跃的 macro 结构内部进一步精细裁剪，获得异构宽度分布

• 具体实现:

  • 对活跃注意力头的 value-path 维度施加门控 $d_{l,h} \in \{0,1\}^{D_h}$
  • 对活跃 FFN 块的中间层神经元施加门控 $c_l \in \{0,1\}^{D_{\text{ffn}}}$
  • 导出时物理截断：删除 $W_{1,l}$ 对应列和 $W_{2,l}$ 对应行

• 模块3: 单阶段训练与温度退火

• 设计动机: 消除两阶段流水线（搜索+微调）的低效，统一为单一连续过程

• 具体实现:

  • Gumbel-Sigmoid 温度 $\tau$ 从初始值 $\tau_0=2.0$ 指数衰减到 $\tau_{\min}=0.5$
  • 早期高 $\tau$：门控类似随机 dropout，迫使存活权重学习鲁棒分布式表征
  • 后期低 $\tau$：概率密度向 0 和 1 双极化，自然硬化为离散子架构
  • 训练结束后以 $\hat{z}>0.5$ 阈值确定性硬化，物理提取子网络

• 模块4: 结构可行性约束

• 设计动机: 防止 神经架构搜索 中常见的结构坍塌（贪婪剪除整层）

• 具体实现:

  • 对每层施加最小保留配额：若活跃头数低于阈值 $k_{\min}$，则施加 ReLU 二次惩罚
  • 类似约束保证最小比例的注意力维度（$\gamma_{\text{attn}}$）和 FFN 神经元（$\gamma_{\text{ffn}}$）存活

• 关键公式

• 公式1: Macro-Level 注意力门控

• 含义: macro-gate $g_{l,h}$ 控制第 $l$ 层第 $h$ 个注意力头的完整输出，$g_{l,h}=0$ 时整个头被跳过

• 符号说明:

  • $g_{l,h} \in \{0,1\}$: 第 $l$ 层第 $h$ 个头的 macro 门控
  • $W^Q_{l,h}, W^K_{l,h}, W^V_{l,h}$: Query / Key / Value 投影矩阵

• 公式2: Macro-Level FFN 门控

• 含义: macro-gate $b_l$ 控制第 $l$ 层整个 FFN 块的保留或移除

• 符号说明:

  • $b_l \in \{0,1\}$: 第 $l$ 层 FFN 块的 macro 门控

• 公式3: Micro-Level 注意力维度门控

• 含义: 在活跃头内部，micro-gate $d_{l,h}$ 逐维度选择性剪枝 value-path 维度

• 符号说明:

  • $d_{l,h} \in \{0,1\}^{D_h}$: 头内维度的 micro 门控向量
  • $\odot$: channel-wise 广播乘法

• 公式4: Micro-Level FFN 神经元门控

• 含义: 在活跃 FFN 块内部，micro-gate $c_l$ 逐神经元选择性剪枝中间隐藏层

• 符号说明:

  • $c_l \in \{0,1\}^{D_{\text{ffn}}}$: FFN 中间层神经元的 micro 门控向量
  • $\varphi$: 非线性激活函数（如 GELU）
  • $W_{1,l}, W_{2,l}$: FFN 的两层权重矩阵

• 公式5: 可微分 MACs 成本建模

• 含义: 将网络的可剪枝计算成本分解为关于门控变量的线性期望，使优化器能精确地将硬件惩罚归因到单个结构单元

• 符号说明:

  • $C_1 = 2ND(3D_h) + 2N^2 D_h$: 单个头的 macro 开销（Q/K/V 投影 + 注意力图计算）
  • $C_2 = 2ND + 2N^2$: 单个存活 value 维度的 micro 成本
  • $C_3 = 4ND$: 单个存活 FFN 中间神经元的 micro 成本

• 公式6: Gumbel-Sigmoid 松弛

• 含义: 将二值门控松弛为连续变量 $\hat{z} \in (0,1)$，通过添加 Logistic 噪声 $\varepsilon$ 并施加带温度 $\tau$ 的 sigmoid，实现端到端可微优化

• 符号说明:

  • $\alpha$: 可学习 logit 参数
  • $\varepsilon$: Logistic 分布噪声
  • $\tau$: 温度参数（从 $\tau_0=2.0$ 退火到 $\tau_{\min}=0.5$）
  • $\sigma(\cdot)$: sigmoid 函数

• 公式7: 总优化目标

• 含义: 联合优化任务损失（交叉熵 + 知识蒸馏）、macro/micro 成本惩罚和结构可行性约束

• 符号说明:

  • $\mathcal{L}_{\text{task}}$: 交叉熵 + KD 软目标损失（$\alpha_{\text{KD}}=0.7$, $T=4.0$）
  • $\lambda_{\text{macro}}, \lambda_{\text{micro}}$: macro/micro 成本惩罚系数（解耦控制）
  • $\mathcal{L}_{\text{macro}}$: 惩罚 $C_1$ 对应的 macro 结构成本
  • $\mathcal{L}_{\text{micro}}$: 惩罚 $C_2, C_3$ 对应的 micro 结构成本
  • $\mathcal{L}_{\text{feasibility}}$: 结构可行性约束

• 公式8: 结构可行性惩罚

其中 $\mathcal{L}_{f,\text{head}} = \sum_l \text{ReLU}(k_{\min} - \sum_h g_{l,h})^2$

• 含义: 当任一层的活跃头数低于阈值 $k_{\min}$ 时施加二次惩罚，防止结构坍塌

• 符号说明:

  • $\beta_{\text{head}}, \beta_{\text{dim}}, \beta_{\text{ffn}}$: 各约束项的权重
  • $k_{\min}$: 每层最小保留头数
  • $\gamma_{\text{attn}}, \gamma_{\text{ffn}}$: 最小存活维度/神经元比例

• 关键图表

• Figure 1: Overview / HiAP 框架概览

{:width 600}

• 说明: HiAP 框架应用于标准 ViT 块的整体架构。可学习 Gumbel-Sigmoid 门控在两个粒度上操作：macro-gates 控制整个注意力头（$g_{l,h}$）和 FFN 块（$b_l$）的保留，micro-gates 控制头内维度（$d_{l,h}$）和 FFN 神经元（$c_l$）的精细裁剪。

• Figure 2: Temperature Annealing / Gumbel-Sigmoid 温度退火

{:width 600}

• 说明: Gumbel-Sigmoid 温度退火过程可视化。早期高温（$\tau=2.0$）时分布近似高斯，作为软连续正则化器；随训练推进温度衰减，概率密度向 0 和 1 双极化，自然硬化为离散子架构而不引起梯度震荡。

• Figure 3: Discovered Architecture / 训练结束后的架构拓扑

{:width 600}

• 说明: HiAP 在 ImageNet 上训练结束后自动发现的架构拓扑。展示了各层注意力头和 FFN 神经元的异构保留模式——浅层保留更多容量，深层（尤其最后一层 $L=12$）被大幅剪枝甚至整个 FFN 块被移除。

• Figure 4: Pareto Frontier / 不同惩罚配置的 Pareto 前沿

{:width 600}

• 说明: 早期训练阶段不同 $\lambda_{\text{macro}}:\lambda_{\text{micro}}$ 配比下的 Top-1 准确率 vs GFLOPs Pareto 前沿。$2:1$ 比率在 DeiT-Small 上表现最优。

• Figure 5: Penalty Ratio 2:1 / $\lambda_{\text{macro}}=0.9, \lambda_{\text{micro}}=0.45$

{:width 600}

• 说明: 推荐的 $2:1$ 比率配置下各层结构演化。深色表示保留更多神经元/头，浅色表示更多剪枝。展现了稳定的层级剪枝轨迹。

• Figure 6: Penalty Ratio 5:1 / $\lambda_{\text{macro}}=1.0, \lambda_{\text{micro}}=0.2$

{:width 600}

• 说明: 偏重 macro 的 $5:1$ 配置。优先丢弃注意力头，但弱 micro 惩罚导致剩余神经元过度膨胀。

• Figure 7: Penalty Ratio 1.5:1 / $\lambda_{\text{macro}}=1.2, \lambda_{\text{micro}}=0.8$

{:width 600}

• 说明: $1.5:1$ 比率产生更均衡的拓扑，与 $2:1$ 配置竞争力相当。

• Figure 8: Macro-Only / $\lambda_{\text{macro}}=1.5, \lambda_{\text{micro}}=0.0$

{:width 600}

• 说明: 纯 macro 惩罚配置。网络在训练早期即激进绕过多个块，MLP 块保留近 100% 神经元，导致 FLOP 预算分配不均衡。

• Figure 9: Micro-Only / $\lambda_{\text{macro}}=0.0, \lambda_{\text{micro}}=1.5$

{:width 600}

• 说明: 纯 micro 惩罚配置。有趣的是，不仅直接减少神经元数量，还间接导致整个 MLP 块被消除。所有注意力头被保留但 FFN 中间层被大幅压缩。

• Table 1: 结构化剪枝框架对比

• Table 2: ImageNet-1K DeiT-Small 剪枝结果

• Table 3: CIFAR-10 ViT-Tiny 消融实验

• 实验结果

• 数据集

• 实现细节

• Backbone: DeiT-Small（ImageNet）/ 自定义 6 层 ViT-Tiny（CIFAR-10）

• 优化器: AdamW, 学习率 $5 \times 10^{-5}$

• Batch Size: 256（全局）

• 训练轮数: 200 epochs（单阶段，搜索+训练同时进行）

• 温度退火: $\tau$ 从 2.0 指数衰减到 0.5

• 知识蒸馏: 使用预训练 Dense DeiT-Small 作为教师（$\alpha_{\text{KD}}=0.7$, $T=4.0$）

• 硬化阈值: $\hat{z}>0.5$

• 主要结论

• ImageNet: HiAP 将 DeiT-Small 压缩至 3.1G MACs（约 33% 计算量削减），Top-1 准确率 79.10%（仅降 0.75%）；更激进的 2.5G MACs 配置下准确率 77.95%

• CIFAR-10: 在 33% 压缩下比 Uniform-Ratio 高 +0.93%，在 50% 压缩下比 $\ell_1$-Structured 高 +0.45%

• 硬件加速: 33.1% 压缩的模型延迟从 5.57ms 降至 3.86ms，实现约 1.44x 推理加速

• 搜索动态分析

• 早期 Macro 优先: 前 10 epochs 内，网络激进地将活跃注意力头从 6 减至平均 2-4 个/层；算法一致性地识别出最后一层（$L=12$）的 FFN 块完全冗余（$b_{12}=0$）

• 后期 Micro 精调: macro 门稳定后，网络转向利用 micro 稀疏性满足剩余 MACs 预算。浅层保留近满容量（约 1400/1536 活跃神经元），深层激进压缩（约 1200 活跃神经元）

• 异构宽度: 存活注意力头内维度从 64 动态截断至 32 或更少

• 理论备注

• Lemma 1: 表达性严格超集

令 $\mathcal{A}_{\text{head}}$ 为仅用 macro 门可达的架构集，$\mathcal{A}_{\text{hiap}}$ 为同时使用 micro 门的架构集。若任一层 $D_h>1$ 或 $D_{\text{ffn}}>1$，则 $\mathcal{A}_{\text{head}} \subsetneq \mathcal{A}_{\text{hiap}}$。

• Proposition 1: 预算线性分解

存在非负权重 $\{w_i\}$ 使得 $\mathbb{E}[C] = \sum_i w_i \mathbb{E}[z_i]$，其中 $z_i \in \{g_{l,h},\; g_{l,h}d_{l,h,j},\; b_l c_{l,k}\}$。无需独立性假设，因线性期望对任意相关性成立。

• Proposition 2: 软-硬预算对齐

当 $\tau \to 0$ 时，连续期望成本收敛至离散成本，使用固定阈值 $\theta=0.5$ 即可找到满足 $|f(0.5) - C_{\text{target}}| \leq \varepsilon$ 的子网络。

• 批判性思考

• 优点

• 统一且简洁: 将 macro+micro 剪枝统一到单一可微框架，无需多阶段流水线或人工启发式

• 理论完备: 可微 MACs 成本的线性分解有严格数学保证（Proposition 1-2），温度退火有信息论支撑

• 实用性强: 物理提取的子网络无需稀疏卷积引擎，直接在标准硬件上获得真实加速

• 自动异构: 网络自主学习不同层的异构剪枝策略，浅层保留多、深层激进压缩，符合直觉

• 局限性

• 准确率差距: 在 3.1G MACs 下 HiAP（79.10%）相比 GOHSP（79.98%）和 ViT-Slim（79.90%）仍有约 0.8-0.9% 的差距，说明简洁性与极致性能之间存在 trade-off

• MACs vs 真实延迟: 优化目标为期望 MACs 而非校准后的延迟/能耗，实际加速因硬件/kernel 而异

• 仅验证分类: 未在检测、分割等下游任务验证泛化性

• 搜索成本: 虽为单阶段但仍需 200 epochs 完整训练，对大规模模型的可扩展性需进一步验证

• 潜在改进方向

• 引入平台校准的延迟/能耗信号替代 MACs 代理

• 与 token 剪枝（如 Token Merging）、量化、编译器优化组合

• 扩展至 LLM 或 multi-modal Transformer 的自动结构搜索

• 探索 $\lambda_{\text{macro}}:\lambda_{\text{micro}}$ 的自适应调度策略

• 可复现性评估

• 代码开源（论文未提供）

• 预训练模型（未提供）

• 训练细节完整（优化器、学习率、batch size、温度退火策略均详细说明）

• 数据集可获取（ImageNet-1K、CIFAR-10）

• 关联笔记

• 基于

• DeiT: 主要实验骨干网络

• ViT: 基础架构

• 对比

• ViT-Slim: 仅 micro 粒度剪枝，$\ell_1$ 稀疏+秩阈值

• GOHSP: 图排序的头和维度联合剪枝

• WDPruning: 多粒度但基于显著性分数

• S2ViTE: 动态稀疏训练

• 方法相关

• Gumbel-Sigmoid: 核心门控松弛技术

• 结构化剪枝: 剪枝范式

• STE: 反向传播中的 Straight-Through Estimator

• 知识蒸馏: 训练中使用预训练教师模型的软目标

• 神经架构搜索: 可微分架构搜索范式

• 硬件/数据相关

• ImageNet: 主要评估数据集

• 速查卡片

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

  • 核心: 多粒度 Gumbel-Sigmoid 门控实现 ViT 单阶段自动剪枝
  • 方法: Macro-gates（头+FFN块）+ Micro-gates（维度+神经元）+ 可微 MACs 成本 + 温度退火
  • 结果: DeiT-Small 在 3.1G MACs（-33%）下 Top-1 79.10%，1.44x 真实加速
  • 代码: 未公开

笔记创建时间: 2026-03-14