Deep neural networks (DNNs) often have to be compressed, via pruning and/or quantization, before they can be deployed in practical settings. In this work we propose a new compression-aware minimizer dubbed CrAM that modifies the optimization step in a principled way, in order to produce models whose local loss behavior is stable under compression operations such as pruning. Thus, dense models trained via CrAM should be compressible post-training, in a single step, without significant accuracy loss. Experimental results on standard benchmarks, such as residual networks for ImageNet classification and BERT models for language modelling, show that CrAM produces dense models that can be more accurate than the standard SGD/Adam-based baselines, but which are stable under weight pruning: specifically, we can prune models in one-shot to 70-80% sparsity with almost no accuracy loss, and to 90% with reasonable ($\sim 1\%$) accuracy loss, which is competitive with gradual compression methods. Additionally, CrAM can produce sparse models which perform well for transfer learning, and it also works for semi-structured 2:4 pruning patterns supported by GPU hardware. The code for reproducing the results is available at https://github.com/IST-DASLab/CrAM .

翻译：深度神经网络（DNN）在部署到实际场景前通常需通过剪枝和/或量化进行压缩。本文提出一种名为CrAM的新型压缩感知优化器，以原则性方式修正优化步骤，使得模型在剪枝等压缩操作下具有稳定的局部损失行为。因此，经CrAM训练的稠密模型可在训练后通过单步压缩实现显著精度保持。在标准基准测试（如用于ImageNet分类的残差网络及用于语言建模的BERT模型）上的实验表明：CrAM生成的稠密模型不仅准确率优于基于标准SGD/Adam的基线方法，且在权重剪枝下保持稳定——具体而言，通过单步剪枝即可在几乎无精度损失下实现70-80%稀疏度，在适度（约1%）精度损失下达到90%稀疏度，性能与渐进压缩方法相当。此外，CrAM可生成适用于迁移学习的稀疏模型，并支持GPU硬件实现的半结构化2:4剪枝模式。复现结果的代码见https://github.com/IST-DASLab/CrAM。