[论文解读] Winning the Lottery Ahead of Time: Efficient Early Network Pruning

John Rachwan, Daniel Zügner+4 more|arXiv (Cornell University)|Jun 21, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出 EarlyCroP 方法，一种高效的早期神经网络剪枝技术，通过保持梯度流动以维持模型性能，同时实现结构化剪枝。该方法在高达 98.5% 参数稀疏度下实现当前最优的准确率，训练时间减少 7 倍，GPU 显存使用减少 4.9 倍，碳排放降低 70%，使大型模型可在消费级硬件上训练而不损失性能。

ABSTRACT

Pruning, the task of sparsifying deep neural networks, received increasing attention recently. Although state-of-the-art pruning methods extract highly sparse models, they neglect two main challenges: (1) the process of finding these sparse models is often very expensive; (2) unstructured pruning does not provide benefits in terms of GPU memory, training time, or carbon emissions. We propose Early Compression via Gradient Flow Preservation (EarlyCroP), which efficiently extracts state-of-the-art sparse models before or early in training addressing challenge (1), and can be applied in a structured manner addressing challenge (2). This enables us to train sparse networks on commodity GPUs whose dense versions would be too large, thereby saving costs and reducing hardware requirements. We empirically show that EarlyCroP outperforms a rich set of baselines for many tasks (incl. classification, regression) and domains (incl. computer vision, natural language processing, and reinforcment learning). EarlyCroP leads to accuracy comparable to dense training while outperforming pruning baselines.

研究动机与目标

• 解决在深度神经网络中寻找稀疏获胜彩票所面临的高计算成本问题。
• 实现结构化剪枝，以减少训练时间、GPU 显存使用和碳排放，同时不损失模型准确率。
• 通过追踪梯度流动和神经正切核（NTK）稳定性，识别剪枝的最佳早期训练时机。
• 通过保持梯度流动来保留学习动态，使大型稀疏模型在消费级 GPU 上的训练成为可能。
• 在视觉、自然语言处理和强化学习任务中，全面超越现有的非结构化与结构化剪枝基线方法。

提出的方法

• 提出一种基于梯度流动（GF）的剪枝准则，以最小化对模型学习动态的干扰。
• 利用梯度流动与神经正切核（NTK）之间的联系，识别适合早期剪枝的稳定训练阶段。
• 在懒惰核阶段应用剪枝，此时 NTK 近似保持恒定，从而确保性能下降最小化。
• 通过移除整个神经元或卷积核，将方法扩展至结构化剪枝，同时保持梯度流动和模型性能。
• 仅通过一次训练运行即可识别并剪枝稀疏子网络，无需多次密集训练循环。
• 基于 NTK 稳定性和 GF 保持设计系统性剪枝调度，以确定最早可行的剪枝时间点。

实验结果

研究问题

• RQ1 我们能否识别出一个剪枝对模型性能影响最小的早期训练阶段，从而实现高效子网络发现？
• RQ2 梯度流动保持能否作为结构化剪枝的合理准则，以维持模型准确率和效率？
• RQ3 与现有结构化及非结构化剪枝基线相比，早期结构化剪枝在准确率、推理速度和碳足迹方面是否表现更优？
• RQ4 通过早期剪枝训练的稀疏模型是否能实现优于同等规模密集模型的性能？
• RQ5 与密集训练相比，早期剪枝在多大程度上减少了训练时间、GPU 显存使用和碳排放？

主要发现

• EarlyCroP-U 在高稀疏度下测试准确率与密集模型相当或更优，并全面超越所有非结构化基线方法，包括 LTH 和 GraSP。
• 与密集训练相比，EarlyCroP-S 将训练时间减少 7 倍，GPU 显存使用减少 4.9 倍，碳排放降低高达 70%。
• 在 CIFAR-100 上使用 VGG16 模型，EarlyCroP-S 在 98.5% 参数稀疏度和 89.9% 节点稀疏度下实现 92.5% 的测试准确率，优于密集模型。
• 对于 ResNeXt101_32x48d 模型，EarlyCroP-S 在 80 个周期后实现 93.2% 的准确率，稀疏度达 98.5%，训练速度提升 6.2 倍，碳排放仅为相似大小密集模型的 1/9.5。
• 在自然语言处理任务中，EarlyCroP-S 在 89% 稀疏度下保持性能，且在更高稀疏度下全面超越所有基线方法，而 LTR 因层特定权重衰减而失败。
• 在强化学习任务中，EarlyCroP-S 超过 LTR 和非结构化基线方法，非结构化模型仅在更长训练周期后才逐渐追平，原因在于梯度计算效率低下。