[论文解读] DSA: More Efficient Budgeted Pruning via Differentiable Sparsity Allocation

Xuefei Ning, Tianchen Zhao+4 more|arXiv (Cornell University)|Apr 5, 2020

Advanced Image and Video Retrieval Techniques参考文献 25被引用 9

一句话总结

DSA 提出了一种用于高效预算剪枝的可微稀疏性分配方法，通过基于梯度的逐层稀疏性分配实现端到端、从零开始的剪枝优化。它在 CIFAR-10 和 ImageNet 上相较于迭代剪枝方法实现了至少 1.5 倍的加速，达到最先进性能，其核心是新颖的可微剪枝过程与在 FLOPs 约束下的 ADMM 优化方法。

ABSTRACT

Budgeted pruning is the problem of pruning under resource constraints. In budgeted pruning, how to distribute the resources across layers (i.e., sparsity allocation) is the key problem. Traditional methods solve it by discretely searching for the layer-wise pruning ratios, which lacks efficiency. In this paper, we propose Differentiable Sparsity Allocation (DSA), an efficient end-to-end budgeted pruning flow. Utilizing a novel differentiable pruning process, DSA finds the layer-wise pruning ratios with gradient-based optimization. It allocates sparsity in continuous space, which is more efficient than methods based on discrete evaluation and search. Furthermore, DSA could work in a pruning-from-scratch manner, whereas traditional budgeted pruning methods are applied to pre-trained models. Experimental results on CIFAR-10 and ImageNet show that DSA could achieve superior performance than current iterative budgeted pruning methods, and shorten the time cost of the overall pruning process by at least 1.5x in the meantime.

研究动机与目标

• 解决传统迭代预算剪枝方法中依赖离散、搜索式稀疏性分配所导致的低效问题。
• 通过实现从零开始的剪枝优化，消除对预训练的依赖。
• 开发一种可微剪枝过程，使能够对剪枝比例进行梯度计算，实现连续优化。
• 在 FLOPs 或延迟约束下，提升稀疏性分配的样本效率与计算效率。
• 通过分组策略处理跳跃连接带来的拓扑约束，确保结构完整性。

提出的方法

• 引入一种可微剪枝过程，通过在通道掩码上使用概率分布软化硬性剪枝，使梯度能够通过剪枝比例流动。
• 使用验证损失作为准确率的可微代理，以支持基于梯度的稀疏性分配优化。
• 应用交替方向乘子法（ADMM）在预算约束下联合优化权重与剪枝比例。
• 将 ADMM 更新集成到主训练循环中，每几十轮权重优化步骤更新一次稀疏性分配，以降低开销。
• 实施一种拓扑分组策略，通过将跳跃连接组视为单位进行剪枝，以在残差网络中保持连接性。
• 通过层间敏感性对齐剪枝比例的梯度幅值，实现归一化与对齐，以指导合理的稀疏性分配。

实验结果

研究问题

• RQ1 在预算剪枝中，基于梯度的连续稀疏性比例优化是否优于离散搜索？
• RQ2 可微剪枝过程是否能在无需预训练的情况下实现有效且高效的稀疏性分配？
• RQ3 任务损失相对于剪枝比例的梯度是否与实际的层间敏感性良好对齐？
• RQ4 DSA 在 CIFAR-10 和 ImageNet 上是否能实现优于迭代剪枝基线的准确率与更快收敛速度？
• RQ5 DSA 的从零开始剪枝方法是否能消除对昂贵预训练的需求，同时保持性能？

主要发现

• 在保留 ResNet-18 60% FLOPs 的情况下，DSA 在 ImageNet 上仅造成 1.11% 的准确率下降，优于所有迭代剪枝基线。
• 与迭代方法相比，DSA 将总剪枝时间减少至少 1.5 倍，在 CIFAR-10 上实现 5/300 GPU 小时，在 ImageNet 上实现 150/300 GPU 小时。
• 相对于剪枝比例的梯度幅值与实际的层间敏感性良好对齐，验证了该方法在稀疏性分配中的合理性。
• DSA 通过将拓扑相关的层（如残差块）分组，成功保持了模型连接性，维持了结构完整性。
• DSA 的从零开始剪枝方法消除了对预训练的需求，通过跳过 300 轮预训练阶段，显著降低了总计算成本。
• 在 CIFAR-10 上，DSA 在多个 FLOPs 预算下持续优于 AMC 和 SSL 等基线方法，展现出鲁棒性与高效性。