[论文解读] Full deep neural network training on a pruned weight budget

Maximilian Golub, Guy Lemieux+1 more|arXiv (Cornell University)|Jun 11, 2018

Advanced Neural Network Applications参考文献 52被引用 10

一句话总结

该论文提出 Dropback，一种训练时剪枝技术，可在反向传播过程中仅动态追踪最具梯度信息的权重，实时重新计算未追踪的权重，从而大幅减少内存访问。该方法在 ResNet18 上实现高达 24.4× 的权重压缩，且精度损失极小，在 Densenet 和 WRN 等复杂架构上达到最先进性能，优于无需微调的训练后剪枝方法。

ABSTRACT

We introduce a DNN training technique that learns only a fraction of the full parameter set without incurring an accuracy penalty. To do this, our algorithm constrains the total number of weights updated during backpropagation to those with the highest total gradients. The remaining weights are not tracked, and their initial value is regenerated at every access to avoid storing them in memory. This can dramatically reduce the number of off-chip memory accesses during both training and inference, a key component of the energy needs of DNN accelerators. By ensuring that the total weight diffusion remains close to that of baseline unpruned SGD, networks pruned using our technique are able to retain state-of-the-art accuracy across network architectures -- including networks previously identified as difficult to compress, such as Densenet and WRN. With ResNet18 on ImageNet, we observe an 11.7$ imes$ weight reduction with no accuracy loss, and up to 24.4$ imes$ with a small accuracy impact.

研究动机与目标

• 为解决 DNN 训练期间，尤其是在移动和嵌入式设备上，高能耗的片外内存访问问题。
• 在不牺牲模型精度的前提下，减少训练时的内存占用和带宽需求，尤其针对密集且难以剪枝的网络。
• 通过将剪枝直接集成到训练过程中，消除剪枝后重新训练的需求。
• 通过最小化持久权重存储，实现更大模型的训练或在低容量硬件上进行微调。

提出的方法

• Dropback 在反向传播期间仅选择并追踪累积梯度最高的权重，其余权重从内存中剪除。
• 未追踪的权重通过其初始值实时再生，避免持久存储和片外内存访问。
• 该算法初始化所有权重，但仅更新在整个训练步骤中总梯度幅值最高的那些权重。
• 其核心洞察在于：高梯度权重主导学习动态，从而确保权重扩散与全量 SGD 接近。
• 该方法与批量归一化和参数化 ReLU 层兼容，这些层通常被排除在剪枝之外。
• 该方法完全在训练阶段运行，避免训练后剪枝的微调过程，且与量化和压缩技术正交。

实验结果

研究问题

• RQ1 我们能否在训练过程中剪枝权重，从而减少 DNN 训练的内存占用，且不造成精度损失？
• RQ2 仅剪枝高梯度权重对权重扩散和模型收敛的影响如何，相较于全量 SGD？
• RQ3 Dropback 在 Densenet 和 WRN 等密集架构上能否实现优于训练后剪枝方法的压缩率和精度？
• RQ4 在低功耗训练场景中，Dropback 能在多大程度上减少片外内存访问和能耗？
• RQ5 该技术是否可应用于通常被排除在剪枝之外的层（如批量归一化和 PReLu）？

主要发现

• 在 ImageNet 上的 ResNet18 上，Dropback 实现 11.7× 的权重压缩且无精度损失，最高可达 24.4× 压缩率，仅造成轻微精度影响。
• 在 Densenet 上，Dropback 实现 5.86% 的验证误差，权重压缩 4.5×，优于先前方法（5.65% / 2.9×）。
• 在 WRN-28-10 上，Dropback 实现 3.85–4.20% 的误差，权重压缩 4.5×–7.3×，优于最佳先前剪枝结果（16.6% 误差 / 4× 压缩）。
• Dropback 保持了与全量 SGD 接近的权重扩散特性，从而在剪枝后无需微调即可实现高精度。
• 该方法通过避免存储未追踪权重，显著降低了内存访问成本，使在低容量设备上训练成为可能。
• Dropback 与量化和压缩技术正交，可与之结合使用以实现进一步的内存压缩。