[论文解读] Only Train Once: A One-Shot Neural Network Training And Pruning Framework

Tianyi Chen, Bo Ji+8 more|arXiv (Cornell University)|Jul 15, 2021

Advanced Neural Network Applications参考文献 81被引用 43

一句话总结

OTO 以一次性训练并从头压缩完整的神经网络，使用 zero-invariant groups 和 一种 新颖的 Half-Space Stochastic Projected Gradient 在不进行微调的情况下进行裁剪。

ABSTRACT

Structured pruning is a commonly used technique in deploying deep neural networks (DNNs) onto resource-constrained devices. However, the existing pruning methods are usually heuristic, task-specified, and require an extra fine-tuning procedure. To overcome these limitations, we propose a framework that compresses DNNs into slimmer architectures with competitive performances and significant FLOPs reductions by Only-Train-Once (OTO). OTO contains two keys: (i) we partition the parameters of DNNs into zero-invariant groups, enabling us to prune zero groups without affecting the output; and (ii) to promote zero groups, we then formulate a structured-sparsity optimization problem and propose a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to solve it, which outperforms the standard proximal methods on group sparsity exploration and maintains comparable convergence. To demonstrate the effectiveness of OTO, we train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction, and achieve state-of-the-art results on VGG16 for CIFAR10, ResNet50 for CIFAR10 and Bert for SQuAD and competitive result on ResNet50 for ImageNet. The source code is available at https://github.com/tianyic/only_train_once.

研究动机与目标

• 在剪枝工作流中，说明在无需多阶段微调的情况下实现高效模型压缩的必要性。
• 引入一个从零开始即可同时训练和裁剪完整模型的框架。
• 定义 zero-invariant groups (ZIGs) 以实现安全、保持输出的裁剪。
• 开发一种新颖的优化算法，Half-Space Stochastic Projected Gradient (HSPG)，以促进结构化稀疏性。
• 在标准基准上展示对 CNN 和 NLP 模型的最前沿或具竞争力的压缩与准确率。

提出的方法

• 将模型参数划分为 zero-invariant groups (ZIGs)，以便在不影响输出的情况下裁剪零值组。
• 对 ZIGs 采用混合 l1/l2 范数的结构化稀疏正则化来诱导组稀疏。
• 使用 Half-Space Stochastic Projected Gradient (HSPG) 求解带非光滑正则化的问题，其包含 Initialize Stage (SGD) 和 Group-Sparsity Stage (Half-Space Projection)。
• 引入一个 Half-Space Projection 运算符，当其下降条件满足时将整组置零，从而实现强力稀疏。
• 通过移除零 ZIGs 对完整模型进行裁剪，以在无需额外微调的情况下获得更瘦的架构。
• 为常见架构（Conv-BN、残差块、全连接和多头注意力）提供自动或预设的 ZIG 划分指导。

实验结果

研究问题

• RQ1 单次从零开始的训练能否在不进行微调的情况下产生一个压缩且性能较高的模型？
• RQ2 如何将参数划分为组（ZIGs），使得零值组在不改变输出的情况下可以被裁剪？
• RQ3 基于 HSPG 的结构化稀疏方法是否比标准的近端方法更有效地促进零组？
• RQ4 在应用 OTO 时，标准基准测试在 FLOPs、参数量和准确率方面的经验提升是多少？
• RQ5 该方法是否与架构无关且可扩展到 CNN 和像 BERT 这样的 NLP 模型？

主要发现

• OTO 在保持竞争力准确率的同时实现了显著的 FLOPs 和参数量下降。
• 在 CIFAR-10 上的 VGG16 和 VGG16-BN 显示了在保持 Top-1 准确率的前提下，显著的 FLOPs 和参数量减少。
• ResNet-50 在 CIFAR-10 和 ImageNet 上实现了显著的 FLOPs 和参数减少，同时准确率损失很小。
• BERT 在 SQuAD 上在单次剪枝框架下展示了有竞争力的结果。
• HSPG 在促进组稀疏性方面优于标准近端方法，并保持收敛性。
• 该剪枝过程得到的压缩模型在相同输入下计算出的输出与完整模型相同且无需微调。