Original: C R AM: A C OMPRESSION -A WARE M INIMIZER
Revised: C R AM: A C OMPRESSION -A WARE M INIMIZER

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep neural networks (DNNs) often have to be compressed, via pruning and/or quantization, before they can be deployed in practical settings.
Revised: Deep neural networks (DNNs) often have to be compressed, via pruning and/or quantization, before they can be deployed in practical settings.

Original: 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.
Revised: 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.

Original: Thus, dense models trained via CrAM should be compressible post-training, in a single step, without significant accuracy loss.
Revised: Thus, dense models trained via CrAM should be compressible post-training, in a single step, without significant accuracy loss.

Original: 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 reasonable ( ≤ 1% ) accuracy loss, which is competitive with gradual compression methods.
Revised: 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 reasonable ( ≤ 1% ) accuracy loss, which is competitive with gradual compression methods.

Original: Additionally, we show that CrAM produces sparse models which perform well for transfer learning, and that it also works for semi-structured pruning patterns supported by GPU hardware.
Revised: Additionally, we show that CrAM produces sparse models which perform well for transfer learning, and that it also works for semi-structured pruning patterns supported by GPU hardware.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The massive recent progress of deep learning models has been accompanied by an increase in computational costs (Thompson et al., 2020).
Revised: The massive recent progress of deep learning models has been accompanied by an increase in computational costs (Thompson et al., 2020).

Original: In turn, this has led to significant interest in model compression techniques in order to reduce these costs.
Revised: In turn, this has led to significant interest in model compression techniques in order to reduce these costs.

Original: For many existing models, compression techniques such as distillation (Hinton et al., 2015), pruning (Hoefler et al., 2021) and quantization (Gholami et al., 2021) can usually reduce the number of parameters or FLOPs of a given model by up to an order of magnitude with relatively little accuracy loss.
Revised: For many existing models, compression techniques such as distillation (Hinton et al., 2015), pruning (Hoefler et al., 2021) and quantization (Gholami et al., 2021) can usually reduce the number of parameters or FLOPs of a given model by up to an order of magnitude with relatively little accuracy loss.

Original: However, performant compression still usually requires re-training or fine-tuning the model separately for each compression target, provided by the user as a target sparsity and/or quantization level.
Revised: However, performant compression still usually requires re-training or fine-tuning the model separately for each compression target, provided by the user as a target sparsity and/or quantization level.

Original: In turn, this compression process can be cumbersome and error-prone, as it requires additional computation and hyper-parameter tuning for each run.
Revised: In turn, this compression process can be cumbersome and error-prone, as it requires additional computation and hyper-parameter tuning for each run.

Original: In this work, we propose Compression-Aware Minimization (CrAM) , a method for training neural networks, which results in models that are easily compressible one-shot , while still being highlyaccurate.
Revised: In this work, we propose Compression-Aware Minimization (CrAM) , a method for training neural networks, which results in models that are easily compressible one-shot , while still being highlyaccurate.

Original: Specifically, CrAM enables training a single (dense) model, which can later be compressed to different target levels, with minimal or no recalibration.
Revised: Specifically, CrAM enables training a single (dense) model, which can later be compressed to different target levels, with minimal or no recalibration.

Original: Such flexibility is desirable, as models can be trained once, and then deployed on multiple devices, with different specifications.
Revised: Such flexibility is desirable, as models can be trained once, and then deployed on multiple devices, with different specifications.

Original: Having a single model that can be easily configured to meet the computational requirements of a specific device can both reduce the overall computational cost, and also allow easier customization to individual devices.
Revised: Having a single model that can be easily configured to meet the computational requirements of a specific device can both reduce the overall computational cost, and also allow easier customization to individual devices.

Original: CrAM is loosely-inspired by the recently-introduced sharpness-aware minimizer (SAM) (Foret et al., 2021), which trains models that potentially converge to flatter minima, leading to better generalization compared to SGD-type baselines, by biasing the process towards minima of uniformly low loss .
Revised: CrAM is loosely-inspired by the recently-introduced sharpness-aware minimizer (SAM) (Foret et al., 2021), which trains models that potentially converge to flatter minima, leading to better generalization compared to SGD-type baselines, by biasing the process towards minima of uniformly low loss .

Original: Multiple subsequent works have investigated and improved upon the original SAM algorithm, by either obtaining better generalization (Kwon et al., 2021), or by reducing the computational costs of
Revised: Multiple subsequent works have investigated and improved upon the original SAM algorithm, by either obtaining better generalization (Kwon et al., 2021), or by reducing the computational costs of

Original: SAM training (Liu et al., 2020; Du et al., 2022a).
Revised: SAM training (Liu et al., 2020; Du et al., 2022a).

Original: We are the first to carry over this idea to the task of obtaining compressible models.
Revised: We are the first to carry over this idea to the task of obtaining compressible models.

Original: Roughly speaking, CrAM works by optimizing not against the original “dense” model, but over a compression projection applied to the intermediate model iterate, at every optimization step.
Revised: Roughly speaking, CrAM works by optimizing not against the original “dense” model, but over a compression projection applied to the intermediate model iterate, at every optimization step.

Original: Thus, the CrAM update aims to bias optimization towards iterates that have both low loss and are robust under one-shot compression .
Revised: Thus, the CrAM update aims to bias optimization towards iterates that have both low loss and are robust under one-shot compression .

Original: Similarly to SAM, CrAM is simple to implement as part of a regular training loop and has a single scaling hyper-parameter, for which we provide a well-performing default value.
Revised: Similarly to SAM, CrAM is simple to implement as part of a regular training loop and has a single scaling hyper-parameter, for which we provide a well-performing default value.

Original: We detail the CrAM algorithm and provide a theoretical motivation leveraging fundamental results in robust optimization (Danskin, 2012) in Section 3.
Revised: We detail the CrAM algorithm and provide a theoretical motivation leveraging fundamental results in robust optimization (Danskin, 2012) in Section 3.

Original: To complement our algorithmic contribution, we perform an extensive experimental analysis of CrAM.
Revised: To complement our algorithmic contribution, we perform an extensive experimental analysis of CrAM.

Original: We mainly focus on compression via weight pruning, but we also show that CrAM is compatible with weight quantization.
Revised: We mainly focus on compression via weight pruning, but we also show that CrAM is compatible with weight quantization.

Original: Generally, CrAM models trained on large-scale image classification or language modelling tasks can improve over the dense baseline performance, while being very robust to one-shot pruning, at different sparsity levels.
Revised: Generally, CrAM models trained on large-scale image classification or language modelling tasks can improve over the dense baseline performance, while being very robust to one-shot pruning, at different sparsity levels.

Original: For image classification, CrAM can train a highly-accurate dense ResNet50 model on ImageNet, that can be pruned in one-shot to 80% and 90% sparsity, and is competitive in terms of accuracy relative to state-of-the-art gradual pruning methods , following an inexpensive Batch Normalization re-tuning step on a small calibration set.
Revised: For image classification, CrAM can train a highly-accurate dense ResNet50 model on ImageNet, that can be pruned in one-shot to 80% and 90% sparsity, and is competitive in terms of accuracy relative to state-of-the-art gradual pruning methods , following an inexpensive Batch Normalization re-tuning step on a small calibration set.

Original: Moreover, we show that full CrAM training is not necessary for good performance: specifically, a short CrAM finetuning period is sufficient to substantially improve one-shot pruning accuracy.
Revised: Moreover, we show that full CrAM training is not necessary for good performance: specifically, a short CrAM finetuning period is sufficient to substantially improve one-shot pruning accuracy.

Original: For instance, we used CrAM to transfer the standard BERT-base model (Devlin et al., 2019) on the SQuADv1.1 question-answering task (Rajpurkar et al., 2016), and obtained models that are both more accurate and more compressible than those obtained with standard optimizers, such as
Revised: For instance, we used CrAM to transfer the standard BERT-base model (Devlin et al., 2019) on the SQuADv1.1 question-answering task (Rajpurkar et al., 2016), and obtained models that are both more accurate and more compressible than those obtained with standard optimizers, such as

Original: Adam (Kingma & Ba, 2015) or SAM (Foret et al., 2021).
Revised: Adam (Kingma & Ba, 2015) or SAM (Foret et al., 2021).

Original: In addition, we noticed that a short ( ≤epochs) finetuning of the sparse model can provide substantial additional improvements: on the above task, the 80%-sparse CrAM finetuned model reaches higher accuracy than the highly-competitive gradual pruning methods PLATON (Zhang et al., 2022) and Movement Pruning (Sanh et al., 2020), at a fraction of the training budget.
Revised: In addition, we noticed that a short ( ≤epochs) finetuning of the sparse model can provide substantial additional improvements: on the above task, the 80%-sparse CrAM finetuned model reaches higher accuracy than the highly-competitive gradual pruning methods PLATON (Zhang et al., 2022) and Movement Pruning (Sanh et al., 2020), at a fraction of the training budget.

Original: Further, CrAM lends itself to several extensions: it can be used with different layer-wise sparsity distributions, semi-structured N:M sparsity patterns, and one-shot pruning techniques.
Revised: Further, CrAM lends itself to several extensions: it can be used with different layer-wise sparsity distributions, semi-structured N:M sparsity patterns, and one-shot pruning techniques.

Original: Sparse CrAM models can be successfully used for sparse transfer learning, where they can perform better on a wide range of “downstream” target tasks, even when compared to pruning methods which adapt to the downstream task (Chen et al., 2021a).
Revised: Sparse CrAM models can be successfully used for sparse transfer learning, where they can perform better on a wide range of “downstream” target tasks, even when compared to pruning methods which adapt to the downstream task (Chen et al., 2021a).

Original: Lastly, we also provide evidence that the CrAM update can produce models that are robust to quantization.
Revised: Lastly, we also provide evidence that the CrAM update can produce models that are robust to quantization.

Original: Similar to SAM (Foret et al., 2021), one limitation of our method is added computational cost, as it requires an additional backwards pass for the model perturbation.
Revised: Similar to SAM (Foret et al., 2021), one limitation of our method is added computational cost, as it requires an additional backwards pass for the model perturbation.

Original: This can be addressed by only performing limited finetuning via CrAM instead of full retraining, or by only performing a regular optimization step for a fraction of the time, both of which we show to have a limited impact on accuracy.
Revised: This can be addressed by only performing limited finetuning via CrAM instead of full retraining, or by only performing a regular optimization step for a fraction of the time, both of which we show to have a limited impact on accuracy.

Original: Moreover, our approach is also compatible with efficient SAM-type updates (Liu et al., 2020; Du et al., 2022a).
Revised: Moreover, our approach is also compatible with efficient SAM-type updates (Liu et al., 2020; Du et al., 2022a).

Original: We also provide a well-performing variant of CrAM that uses sparse gradients, which could be leveraged by frameworks with support for sparse back-propagation.
Revised: We also provide a well-performing variant of CrAM that uses sparse gradients, which could be leveraged by frameworks with support for sparse back-propagation.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: We describe in this section some of the recent research directions that have inspired the development of our method, together with existing literature focused on solving similar problems.
Revised: We describe in this section some of the recent research directions that have inspired the development of our method, together with existing literature focused on solving similar problems.

Original: Sharpness-Aware Minimization (SAM).
Revised: Sharpness-Aware Minimization (SAM).

Original: The recently introduced SAM optimizer (Foret et al., 2021) aims to improve the generalization of deep neural networks, by encouraging the minimization of loss sharpness; this in turn should lead to flatter local minima, with better generalization properties.
Revised: The recently introduced SAM optimizer (Foret et al., 2021) aims to improve the generalization of deep neural networks, by encouraging the minimization of loss sharpness; this in turn should lead to flatter local minima, with better generalization properties.

Original: The authors show that SAM-trained models have higher validation accuracy compared to vanilla
Revised: The authors show that SAM-trained models have higher validation accuracy compared to vanilla

Original: SGD-type baselines,
Revised: SGD-type baselines,

Original: that their performance continues to improve with prolonged training, and that they can also be successfully used for transfer learning.
Revised: that their performance continues to improve with prolonged training, and that they can also be successfully used for transfer learning.

Original: One important drawback of SAM is its computational overhead, as it requires twice as many forward-backward passes through the network.
Revised: One important drawback of SAM is its computational overhead, as it requires twice as many forward-backward passes through the network.

Original: Subsequent work has focused on reducing computational cost by, for example, reducing the frequency of the extra gradient steps (Liu et al., 2022), computing the perturbations on a subset of the parameters (Du et al., 2022a), or by proposing a new trajectory loss to replace the sharpness definition (Du et al., 2022b).
Revised: Subsequent work has focused on reducing computational cost by, for example, reducing the frequency of the extra gradient steps (Liu et al., 2022), computing the perturbations on a subset of the parameters (Du et al., 2022a), or by proposing a new trajectory loss to replace the sharpness definition (Du et al., 2022b).

Original: We draw inspiration from properties of the initial SAM method proposed by Foret et al.
Revised: We draw inspiration from properties of the initial SAM method proposed by Foret et al.

Original: Instead of attempting to minimize the maximum local increase loss (sharpness), our goal is to minimize the maximum local increase in loss due to compression.
Revised: Instead of attempting to minimize the maximum local increase loss (sharpness), our goal is to minimize the maximum local increase in loss due to compression.

Original: Training prunable networks.
Revised: Training prunable networks.

Original: The increasing scale of deep neural networks have made their deployment to edge devices dependent on compression techniques, such as quantization and/or pruning.
Revised: The increasing scale of deep neural networks have made their deployment to edge devices dependent on compression techniques, such as quantization and/or pruning.

Original: While post-training quantization can be an efficient and successful technique for quantizing models without any retraining (Frantar & Alistarh, 2022), in the case of pruning the gold standard is still training a separate model for every target sparsity level (Zhu & Gupta, 2017; Singh & 
Revised: While models can be quantized post-training efficiently and successfully without any retraining (Frantar & Alistarh, 2022), in the case of pruning the gold standard is still training a separate model for every target sparsity level (Zhu & Gupta, 2017; Singh & Alistarh, 2020;

Original: Alistarh, 2020; Evci et al., 2020; Peste et al., 2021); the latter can be an expensive procedure, which would still rely on powerful computational resources to obtain the sparse models in the first place.
Revised: Evci et al., 2020; Peste et al., 2021); the latter can be an expensive procedure, which would still rely on powerful computational resources to obtain the sparse models in the first place.

Original: A potential solution would be training a single dense model, which either contains multiple smaller ones that can be easily deployed, or which is itself prunable at multiple sparsity levels, without additional retraining.
Revised: A potential solution would be training a single dense model, which either contains multiple smaller ones that can be easily deployed, or which is itself prunable at multiple sparsity levels, without additional retraining.

Original: For example, the “once-for-all” (OFA) framework (Cai et al., 2019) can train a large network that contains multiple specialized sub-nets, adapted to different resource constraint devices.
Revised: For example, the “once-for-all” (OFA) framework (Cai et al., 2019) can train a large network that contains multiple specialized sub-nets, adapted to different resource constraint devices.

Original: However, obtaining the large OFA network is extremely expensive, and requires intensive finetuning to ensure a good performance for the sub-nets.
Revised: However, obtaining the large OFA network is extremely expensive, and requires intensive finetuning to ensure a good performance for the sub-nets.

Original: A similar idea thatalso requires extensive finetuning has been explored for automatic speech recognition (Wu et al., 2021).
Revised: A similar idea that requires extensive finetuning has been explored for automatic speech recognition (Wu et al., 2021).

Original: An orthogonal direction is to obtain “slimmable neural networks” (Yu et al., 2019; Yu & Huang, 2019b;a), by training a single model that can be executed at different widths; this is usually achieved by performing multiple backpropagations using all the predefined widths, at each optimization step, and by carefully considering the Batch Normalization layers.
Revised: An orthogonal direction is to obtain “slimmable neural networks” (Yu et al., 2019; Yu & Huang, 2019b;a), by training a single model that can be executed at different widths; this is usually achieved by performing multiple backpropagations using all the predefined widths, at each optimization step, and by carefully considering the Batch Normalization layers.

Revised: Additionally, BigNAS (Yu et al., 2020) is inspired by these approaches, and achieves a similar goal to OFA, without additional finetuning.

Original: Related to one-shot pruning, Only
Revised: Related to one-shot pruning, Only

Original: Train Once (OTO) (Chen et al., 2021b) has been proposed for structured pruning, to train a large model that is easily slimmable one-shot.
Revised: Train Once (OTO) (Chen et al., 2021b) has been proposed for structured pruning, to train a large model that is easily slimmable one-shot.

Original: While we obtain better results than OTO for the same sparsity level, the two methods are not directly comparable, since we focus on unstructured sparsity .
Revised: While we obtain better results than OTO for the same sparsity level, the two methods are not directly comparable, since we focus on unstructured sparsity .

Original: Our work is more closely related to Miao et al.
Revised: Our work is more closely related to Miao et al.

Original: ; Zimmer et al.
Revised: ; Zimmer et al.

Original: (2022), which propose leveraging
Revised: (2022), which propose leveraging

Original: Stochastic Frank-Wolfe (SFW) (Reddi et al., 2016) to encourage the weights to lie in a convex hull spanned by sparse vectors; this would make the model prunable one-shot, without any finetuning.
Revised: Stochastic Frank-Wolfe (SFW) (Reddi et al., 2016) to encourage the weights to lie in a convex hull spanned by sparse vectors; this would make the model prunable one-shot, without any finetuning.

Original: The methods proposed in Miao et al.
Revised: The methods proposed in Miao et al.

Original: ; Zimmer et al.
Revised: ; Zimmer et al.

Original: (2022) result in highly-prunable models on relatively-small tasks; specifically, their experimental analysis is limited to image classification on small datasets and architectures with many redundancies (e.g. VGG-16 (Simonyan & Zisserman, 2014) on CIFAR-10).
Revised: (2022) result in highly-prunable models on relatively-small tasks; specifically, their experimental analysis is limited to image classification on small datasets and architectures with many redundancies (e.g. VGG-16 (Simonyan & Zisserman, 2014) on CIFAR-10).

Original: CrAM is able to match or outperform these methods in the same setting: for instance, CrAM can prune VGG-16 trained on CIFAR-10 in one-shot to 95% sparsity without accuracy loss, outperforming SFW by more than 2% Top-1 acuracy.
Revised: CrAM is able to match or outperform these methods in the same setting: for instance, CrAM can prune VGG-16 trained on CIFAR-10 in one-shot to 95% sparsity without accuracy loss, outperforming SFW by more than 2% Top-1 acuracy.

Original: More importantly, we show that CrAM produces models compressible in one-shot at both ImageNet scale and BERT language modeling scale.
Revised: More importantly, we show that CrAM produces models compressible in one-shot at both ImageNet scale and BERT language modeling scale.

Original: Remarkably, with one-shot pruning CrAM can offer competitive performance to gradual pruning methods, whether they are designed for CNNs (Kusupati et al., 2020; Lin et al., 2020) or for language modelling (Sanh et al., 2020; Zhang et al., 2022).
Revised: Remarkably, with one-shot pruning CrAM can offer competitive performance to gradual pruning methods, whether they are designed for CNNs (Kusupati et al., 2020; Lin et al., 2020) or for language modelling (Sanh et al., 2020; Zhang et al., 2022).

Original:  T HE C OMPRESSION -A WARE M INIMIZER (C R AM) 3. B ACKGROUND
Revised:  T HE C OMPRESSION -A WARE M INIMIZER (C R AM) 3. B ACKGROUND

Original: We now give an overview of our method, together with the corresponding algorithm and generalizations.
Revised: We now give an overview of our method, together with the corresponding algorithm and generalizations.

Original: One of the main goals of CrAM is to train models that are “compressible” in one-shot, following training, via sparsity or quantization.
Revised: One of the main goals of CrAM is to train models that are “compressible” in one-shot, following training, via sparsity or quantization.

Original: We denote C a compression operator, for example Top-K, where only the highest K absolute values of a tensor are kept, while the rest are set to 0.
Revised: We denote C a compression operator, for example Top-K, where only the highest K absolute values of a tensor are kept, while the rest are set to 0.

Original: We say that a model is easily compressible if small perturbations do not affect its performance after compression.
Revised: We say that a model is easily compressible if small perturbations do not affect its performance after compression.

Original: To enforce this during training, we optimize against the perturbation which has the most significant impact on the compressed model.
Revised: To enforce this during training, we optimize against the perturbation which has the most significant impact on the compressed model.

Original: We want to minimize the “compression-aware” (CrAM) loss, defined as:
Revised: We want to minimize the “compression-aware” (CrAM) loss, defined as:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where θ is the vector of model parameters, L is the regular cross-entropy loss and δ is a norm-bounded perturbation.
Revised: where θ is the vector of model parameters, L is the regular cross-entropy loss and δ is a norm-bounded perturbation.

Original: Unless otherwise stated, we employ the ℓ 2 -norm throughout the rest of the paper.
Revised: Unless otherwise stated, we employ the ℓ 2 -norm throughout the rest of the paper.

Original: We approximate max δ L ( C ( θ + δ )) by taking a gradient ascent step in the direction of the current update, followed by a projection using the compression operator.
Revised: We approximate max δ L ( C ( θ + δ )) by taking a gradient ascent step in the direction of the current update, followed by a projection using the compression operator.

Original: This is inspired by the iterative hard thresholding (IHT) algorithm used for optimizing functions under sparse constraints (Blumensath &
Revised: This is inspired by the iterative hard thresholding (IHT) algorithm used for optimizing functions under sparse constraints (Blumensath &

Original: Davies, 2008; Foucart, 2011; 2012).
Revised: Davies, 2008; Foucart, 2011; 2012).

Original: To obtain the gradient with respect to the parameters, we employ a straight-through estimator, by using instead the gradient under the perturbation.
Revised: To obtain the gradient with respect to the parameters, we employ a straight-through estimator, by using instead the gradient under the perturbation.

Original: This gives us the following update for minimizing the CrAM loss:
Revised: This gives us the following update for minimizing the CrAM loss:

Original: [Equation]
Revised: [Equation]

Original: We note that solely optimizing the CrAM loss cannot offer guarantees for the performance of the dense model.
Revised: We note that solely optimizing the CrAM loss cannot offer guarantees for the performance of the dense model.

Original: Alongside improving robustness to compression, maintaining the quality of the dense model is one of the prerequisites of our method; therefore, we propose to also explicitly optimize for the performance of the dense model.
Revised: Alongside improving robustness to compression, maintaining the quality of the dense model is one of the prerequisites of our method; therefore, we propose to also explicitly optimize for the performance of the dense model.

Original: Specifically, we optimize instead the following composite
Revised: Specifically, we optimize instead the following composite

Original: CrAM + loss function :
Revised: CrAM + loss function :

Original: [Equation]
Revised: [Equation]

Original: This can be achieved with a simple modification to the CrAM update, at no extra cost, by simply adding the gradient ∇ L ( θ t ) , before the next update of the parameters θ t +1 .
Revised: This can be achieved with a simple modification to the CrAM update, at no extra cost, by simply adding the gradient ∇ L ( θ t ) , before the next update of the parameters θ t +1 .

Original: For (cid:101) θ t =
Revised: For (cid:101) θ t =

Original: C ( θ t + ρ ∇ L ( θ t )) , the CrAM + update is the following:
Revised: C ( θ t + ρ ∇ L ( θ t )) , the CrAM + update is the following:

Original: [Equation]
Revised: [Equation]

Original: We note that we can add different regularization terms to the objective in Equation 3; for example, in our experiments we use weight decay, as it is standard for training image classification models.
Revised: We note that we can add different regularization terms to the objective in Equation 3; for example, in our experiments we use weight decay, as it is standard for training image classification models.

Original: Require: Compression methods C
Revised: Require: Compression methods C

Original: = { C 1 , C 2 , . . .
Revised: = { C 1 , C 2 , . . .

Original: , C M } , training data S , training iterations T , learning rate η , perturbation step size ρ 1: Initialize the weights θ 0 2: while t ≤ T do 3:
Revised: , C M } , training data S , training iterations T , learning rate η , perturbation step size ρ 1: Initialize the weights θ 0 2: while t ≤ T do 3:

Original: Sample batch x ∈ S 4:
Revised: Sample batch x ∈ S 4:

Original: Compute loss L ( θ t ; x ) and gradient g t = ∇ L ( θ t ; x ) 5:
Revised: Compute loss L ( θ t ; x ) and gradient g t = ∇ L ( θ t ; x ) 5:

Original: Uniformly choose a compression method C ∈ C 6:
Revised: Uniformly choose a compression method C ∈ C 6:

Original: Get perturbed weights (cid:101) θ t
Revised: Get perturbed weights (cid:101) θ t

Original: = C ( θ t + ρ g t ) 7:
Revised: = C ( θ t + ρ g t ) 7:

Original: if C = Top-K then 8:
Revised: if C = Top-K then 8:

Original: Let M t be the linear projection operator onto the support of the largest K coordinates of | θ t + ρ g t | , such that (cid:101) θ t =
Revised: Let M t be the linear projection operator onto the support of the largest K coordinates of | θ t + ρ g t | , such that (cid:101) θ t =

Original: M t ( θ t + ρ g t ) 9: (cid:101) g t =
Revised: M t ( θ t + ρ g t ) 9: (cid:101) g t =

Original: M t ∇ L ( (cid:101) θ t ; x ) 10: else 11: (cid:101) g t = ∇ L ( (cid:101) θ t ; x ) 12: end if 13: if use CrAM + then 14: (cid:101) g t ← (cid:101) g
Revised: M t ∇ L ( (cid:101) θ t ; x ) 10: else 11: (cid:101) g t = ∇ L ( (cid:101) θ t ; x ) 12: end if 13: if use CrAM + then 14: (cid:101) g t ← (cid:101) g

Original: t + g t 15: end if 16:
Revised: t + g t 15: end if 16:

Original: Update the weights using a gradient descent step: θ t +1 =
Revised: Update the weights using a gradient descent step: θ t +1 =

Original: θ t − η · (cid:101) g
Revised: θ t − η · (cid:101) g

Original: t 17: end while 18: return θ T
Revised: t 17: end while 18: return θ T

Original: 3. T HEORETICAL J USTIFICATION OF THE C R AM U PDATE
Revised: 3. T HEORETICAL J USTIFICATION OF THE C R AM U PDATE

Original: To derive the CrAM update, and justify the choices made in designing our training method, we start from the optimization objective defined in Equation 1.
Revised: To derive the CrAM update, and justify the choices made in designing our training method, we start from the optimization objective defined in Equation 1.

Original: As our goal is to minimize the CrAM loss
Revised: As our goal is to minimize the CrAM loss

Original: L CrAM , we use gradient descent.
Revised: L CrAM , we use gradient descent.

Original: Using this loss complicates our objective, as it now includes an inner maximization problem, together with a potentially problematic compression operator.
Revised: Using this loss complicates our objective, as it now includes an inner maximization problem, together with a potentially problematic compression operator.

Original: However, under mild assumptions we can efficiently estimate a “fake” gradient which gives a descent direction.
Revised: However, under mild assumptions we can efficiently estimate a “fake” gradient which gives a descent direction.

Original: To do so we rely on a well-known theorem from robust optimization (Danskin, 2012), which allows one to obtain descent directions for min-max objectives under a broad range of assumptions.
Revised: To do so we rely on a well-known theorem from robust optimization (Danskin, 2012), which allows one to obtain descent directions for min-max objectives under a broad range of assumptions.

Original: Using
Revised: Using

Original: Danskin’s theorem (Theorem 1 from Appendix E.1) we obtain that by computing the maximizer of the inner problem
Revised: Danskin’s theorem (Theorem 1 from Appendix E.1) we obtain that by computing the maximizer of the inner problem

Original: [Equation]
Revised: [Equation]

Original: and letting ϕ
Revised: and letting ϕ

Original: = θ + δ ∗ , which compresses to the extrapolated iterate (cid:101) θ = C ( ϕ ) , we obtain a descent direction −∇ L ( C ( ϕ )) .
Revised: = θ + δ ∗ , which compresses to the extrapolated iterate (cid:101) θ = C ( ϕ ) , we obtain a descent direction −∇ L ( C ( ϕ )) .

Original: Implementing this approach faces two difficulties – first, to compute the gradient we must back-propagate through the composition of functions L ( C ( · )) , which may cause trouble since C is not necessarily differentiable; second, and more importantly, it is unclear how to solve the inner maximization problem.
Revised: Implementing this approach faces two difficulties – first, to compute the gradient we must back-propagate through the composition of functions L ( C ( · )) , which may cause trouble since C is not necessarily differentiable; second, and more importantly, it is unclear how to solve the inner maximization problem.

Original: To address the first issue, we may choose to use a straight-through gradient estimator (Bengio et al., 2013), which permits us to only backpropagate through L , and use ∇ L ( (cid:101) θ ) instead of the true gradient.
Revised: To address the first issue, we may choose to use a straight-through gradient estimator (Bengio et al., 2013), which permits us to only backpropagate through L , and use ∇ L ( (cid:101) θ ) instead of the true gradient.

Original: To increase precision, in the case where compression is performed via Top-K we interpret C as a “mask” operator M which zeroes out a subset of coordinates dependent on ϕ .
Revised: To increase precision, in the case where compression is performed via Top-K we interpret C as a “mask” operator M which zeroes out a subset of coordinates dependent on ϕ .

Original: Since except for articulation points, M t is constant and does not change as the argument varies, we approximate ∇ L ( C ( ϕ ))
Revised: Since except for articulation points, M t is constant and does not change as the argument varies, we approximate ∇ L ( C ( ϕ ))

Original: ≈ M ∇ L ( M ϕ ) =
Revised: ≈ M ∇ L ( M ϕ ) =

Original: M ∇ L ( (cid:101) θ ) .
Revised: M ∇ L ( (cid:101) θ ) .

Original: To address the second issue, rather than exactly maximizing the inner problem, we instead seek a good enough maximizer using a standard iterative method.
Revised: To address the second issue, rather than exactly maximizing the inner problem, we instead seek a good enough maximizer using a standard iterative method.

Original: For this, we choose projected gradient ascent, which provides theoretical guarantees, even when the projection is performed onto non-convex domains (Peste et al., 2021).
Revised: For this, we choose projected gradient ascent, which provides theoretical guarantees, even when the projection is performed onto non-convex domains (Peste et al., 2021).

Original: For instance, if the compression operator is magnitude pruning, this becomes the iterative hard thresholding (IHT) method, frequently employed in the sparse recovery literature (Blumensath & Davies, 2008).
Revised: For instance, if the compression operator is magnitude pruning, this becomes the iterative hard thresholding (IHT) method, frequently employed in the sparse recovery literature (Blumensath & Davies, 2008).

Original: Thus, to reach a good iterate within this specific domain, in practice we perform a single step of (projected) gradient ascent, which matches the IHT iteration:
Revised: Thus, to reach a good iterate within this specific domain, in practice we perform a single step of (projected) gradient ascent, which matches the IHT iteration:

Original: [Equation]
Revised: [Equation]

Original: In Appendix E, we provide a full re-derivation of the CrAM update in Equation 2 under fairly reasonable assumptions on the objective function, along with a detailed discussion on the necessity of these assumptions.
Revised: In Appendix E, we provide a full re-derivation of the CrAM update in Equation 2 under fairly reasonable assumptions on the objective function, along with a detailed discussion on the necessity of these assumptions.

Original: As a side result, we also obtain a simple re-derivation of the SAM update.
Revised: As a side result, we also obtain a simple re-derivation of the SAM update.

Original: 3. I MPLEMENTATION D ETAILS AND E XTENSIONS
Revised: 3. I MPLEMENTATION D ETAILS AND E XTENSIONS

Original: Multiple compression types.
Revised: Multiple compression types.

Original: CrAM can be used to train models that are robust to multiple types of compression operators.
Revised: CrAM can be used to train models that are robust to multiple types of compression operators.

Original: This can be enforced by choosing between multiple compression projections at each CrAM optimization step.
Revised: This can be enforced by choosing between multiple compression projections at each CrAM optimization step.

Original: Examples include pruning using different sparsity levels or quantizing at different precisions.
Revised: Examples include pruning using different sparsity levels or quantizing at different precisions.

Original: We illustrate the general CrAM algorithm which handles multiple compression operators, and includes the explicit optimization of the dense model, in Algorithm 1.
Revised: We illustrate the general CrAM algorithm which handles multiple compression operators, and includes the explicit optimization of the dense model, in Algorithm 1.

Original: In our experiments, we found that applying the CrAM + update with a different randomly chosen compression at each optimization step typically achieves a good trade-off between a high dense model accuracy and robustness to multiple one-shot compression schemes post-training.
Revised: In our experiments, we found that applying the CrAM + update with a different randomly chosen compression at each optimization step typically achieves a good trade-off between a high dense model accuracy and robustness to multiple one-shot compression schemes post-training.

Original: When optimizing for robustness against sparse perturbations, we use the Top-K operator at each step, and choose the sparsity level uniformly at random among a set of predefined values.
Revised: When optimizing for robustness against sparse perturbations, we use the Top-K operator at each step, and choose the sparsity level uniformly at random among a set of predefined values.

Original: Addressing the computational overhead of CrAM.
Revised: Addressing the computational overhead of CrAM.

Original: Similar to the original SAM update, CrAM requires twice as many forward-backward passes, compared to a regular training cycle.
Revised: Similar to the original SAM update, CrAM requires twice as many forward-backward passes, compared to a regular training cycle.

Original: In the case of
Revised: In the case of

Original: TopK-CrAM, we can reduce this overhead, by making use of the sparsity in the intermediate updates.
Revised: TopK-CrAM, we can reduce this overhead, by making use of the sparsity in the intermediate updates.

Original: Furthermore, we found that using only the gradients from the support of (cid:101) θ t in ∇ L ( (cid:101) θ t ) improves both the resulting dense model obtained with TopK-CrAM, as well as its robustness to one-shot pruning.
Revised: Furthermore, we found that using only the gradients from the support of (cid:101) θ t in ∇ L ( (cid:101) θ t ) improves both the resulting dense model obtained with TopK-CrAM, as well as its robustness to one-shot pruning.

Original: This observation is motivated by the fact that the straight-through (Bengio et al., 2013) gradient estimator using the identity function is often times suboptimal (Yin et al., 2019), and better straight-through estimators can be defined using different functions.
Revised: This observation is motivated by the fact that the straight-through (Bengio et al., 2013) gradient estimator using the identity function is often times suboptimal (Yin et al., 2019), and better straight-through estimators can be defined using different functions.

Original: As seen in Section 3.2, we can assume, via Danskin’s theorem (Danskin, 2012), that we can obtain descent directions for L CrAM ( θ t ) by evaluating ∇ L ( C ( ϕ t )) , where ϕ t is the extrapolated point ϕ t
Revised: As seen in Section 3.2, we can assume, via Danskin’s theorem (Danskin, 2012), that we can obtain descent directions for L CrAM ( θ t ) by evaluating ∇ L ( C ( ϕ t )) , where ϕ t is the extrapolated point ϕ t

Original: = θ t + ρ ∇ L ( θ t ) .
Revised: = θ t + ρ ∇ L ( θ t ) .

Original: To evaluate the gradient, we may use a straight-through estimator.
Revised: To evaluate the gradient, we may use a straight-through estimator.

Original: For Top-K, C is as an operator M t which zeroes out a subset of coordinates dependent on ϕ t .
Revised: For Top-K, C is as an operator M t which zeroes out a subset of coordinates dependent on ϕ t .

Original: Provided that M t is constant and does not change as the argument varies, we can approximate ∇ L ( C ( ϕ t )) ∼ M t ∇ L ( M t ϕ t ) .
Revised: Provided that M t is constant and does not change as the argument varies, we can approximate ∇ L ( C ( ϕ t )) ∼ M t ∇ L ( M t ϕ t ) .

Original: As both the iterate and gradient estimator are sparse, this implies a theoretical speed-up.
Revised: As both the iterate and gradient estimator are sparse, this implies a theoretical speed-up.

Original: Alternative Updates.
Revised: Alternative Updates.

Original: We note that alternative compression-aware updates can be derived.
Revised: We note that alternative compression-aware updates can be derived.

Original: For example, by following similar derivations to those developed for SAM (Foret et al., 2021), we get (cid:101) 
Revised: One such derivation is (cid:101) θ t = C (cid:16) θ t + ρ ∇ L ( C ( θ t ))

Original: [Equation] = [Equation]
Revised: ∥∇ L ( C ( θ t ))

Original: + [Equation] [Equation] L ( C ( θ t )) ∥ (cid:17) .
Revised: ∥ (cid:17) , which we obtain by following similar steps to those developed for SAM (Foret et al., 2021).

Original: We call this update Compressed-SAM (C-SAM) .
Revised: We call this update Compressed-SAM (C-SAM) .

Original: We observed that training with C-SAM can also result in models that are robust to one-shot pruning, but typically the accuracy of the resulting dense models is lower, compared to training with CrAM.
Revised: We observed that training with C-SAM can also result in models that are robust to one-shot pruning, but typically the accuracy of the resulting dense models is lower, compared to training with CrAM.

Original: Moreover, training with C-SAM cannot offer guarantees for the performance of the dense model.
Revised: Moreover, training with C-SAM cannot offer guarantees for the performance of the dense model.

Original: While with CrAM we can optimize the dense model loss for free (i.e. using CrAM + ), with C-SAM optimizing for the dense model explicitly would require a third forward-backward pass at each training step.
Revised: While with CrAM we can optimize the dense model loss for free (i.e. using CrAM + ), with C-SAM optimizing for the dense model explicitly would require a third forward-backward pass at each training step.

Original: Additionally, we examine the importance of the extra gradient step in CrAM, by comparing against simply applying the Top-K operator to the parameters.
Revised: Additionally, we examine the importance of the extra gradient step in CrAM, by comparing against simply applying the Top-K operator to the parameters.

Original: We provide an ablation study in Appendix C.1.
Revised: We provide an ablation study in Appendix C.1.

Original: Statistics Correction.
Revised: Statistics Correction.

Original: It is well-known (Hubara et al., 2021; Frantar & Alistarh, 2022) that pruning weights in a single step at high sparsity levels can have a large negative effect on normalization layers, due to a mismatch between layer statistics, e.g. the running mean and variance of BatchNorm layers, computed during training, and those of the pruned model.
Revised: It is well-known (Hubara et al., 2021; Frantar & Alistarh, 2022) that pruning weights in a single step at high sparsity levels can have a large negative effect on normalization layers, due to a mismatch between layer statistics, e.g. the running mean and variance of BatchNorm layers, computed during training, and those of the pruned model.

Original: To correct for this, following prior work, we keep a subset of randomly chosen 1000 training samples (e.g. for ImageNet one sample per class), to which we apply standard training augmentations, and which are used post-pruning for resetting the Batch Norm statistics of the pruned model.
Revised: To correct for this, following prior work, we keep a subset of randomly chosen 1000 training samples (e.g. for ImageNet one sample per class), to which we apply standard training augmentations, and which are used post-pruning for resetting the Batch Norm statistics of the pruned model.

Original: We note that this procedure, which we refer to as BatchNorm Tuning (BNT) is very inexpensive, and does not finetune any other parameters of the model.
Revised: We note that this procedure, which we refer to as BatchNorm Tuning (BNT) is very inexpensive, and does not finetune any other parameters of the model.

Original: Furthermore, during CrAM training on image classification models we only track the
Revised: Furthermore, during CrAM training on image classification models we only track the

Original: BatchNorm statistics on the dense model, before applying the compression perturbation.
Revised: BatchNorm statistics on the dense model, before applying the compression perturbation.

Original: In the case of BERT models, we do not apply any statistics corrections.
Revised: In the case of BERT models, we do not apply any statistics corrections.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our experimental validation mainly focuses on sparsity, obtained by applying the Top-K operator, in the context of CrAM (i.e. TopK-CrAM).
Revised: Our experimental validation mainly focuses on sparsity, obtained by applying the Top-K operator, in the context of CrAM (i.e. TopK-CrAM).

Original: We denote the CrAM runs by the sparsity level used during training.
Revised: We denote the CrAM runs by the sparsity level used during training.

Original: For example, “CrAM-k50” indicates that the Top-K operator with k=50% was used at each step, while “CrAM-Multi” indicates that the sparsity level is chosen uniformly at random, at each step, from a set of given values (e.g. CrAM-k { 50 , 70 , 90 } ) .
Revised: For example, “CrAM-k50” indicates that the Top-K operator with k=50% was used at each step, while “CrAM-Multi” indicates that the sparsity level is chosen uniformly at random, at each step, from a set of given values (e.g. CrAM-k { 50 , 70 , 90 } ) .

Original: For image classification experiments, all one-shot pruning results are presented after Batch Norm tuning (BNT) on a subset of 1000 training samples, i.e. 100 inference steps on batches of size 128, using standard random augmentations.
Revised: For image classification experiments, all one-shot pruning results are presented after Batch Norm tuning (BNT) on a subset of 1000 training samples, i.e. 100 inference steps on batches of size 128, using standard random augmentations.

Original: 4. I MAGE N ET E XPERIMENTS
Revised: 4. I MAGE N ET E XPERIMENTS

Original: General Setup.
Revised: General Setup.

Original: We use a standard setup for training our ImageNet/ResNet50 models, similar to Foret et al.
Revised: We use a standard setup for training our ImageNet/ResNet50 models, similar to Foret et al.

Original: (2021), which we describe in the Appendix B.
Revised: (2021), which we describe in the Appendix B.

Original: To match the number of backpropagation steps of CrAM, we additionally train the dense baseline for twice as many epochs.
Revised: To match the number of backpropagation steps of CrAM, we additionally train the dense baseline for twice as many epochs.

Original: We have found that ρ = 0 .
Revised: We have found that ρ = 0 .

Original: 05 recommended by the authors of SAM (Foret et al., 2021) is a good value for CrAM, and we have kept it for all our ImageNet experiments.
Revised: 05 recommended by the authors of SAM (Foret et al., 2021) is a good value for CrAM, and we have kept it for all our ImageNet experiments.

Original: As stated, after one-shot pruning, we perform BNT on a subset of 1000 training samples (e.g. one per class), with standard augmentations.
Revised: As stated, after one-shot pruning, we perform BNT on a subset of 1000 training samples (e.g. one per class), with standard augmentations.

Original: We show in the
Revised: We show in the

Original: Appendix C.3 that the accuracy after BNT is extremely stable, w.r.t.
Revised: Appendix C.3 that the accuracy after BNT is extremely stable, w.r.t.

Original: the choice of calibration set.
Revised: the choice of calibration set.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: Results for one-shot pruning.
Revised: Results for one-shot pruning.

Original: We validate the robustness to post-training compression of models trained through different versions of CrAM, by testing their accuracy after one-shot pruning, at different sparsity levels.
Revised: We validate the robustness to post-training compression of models trained through different versions of CrAM, by testing their accuracy after one-shot pruning, at different sparsity levels.

Original: We train models using CrAM-k50, CrAM + -k70 and CrAM + -Multi, where for the latter we choose uniformly at random, at each step, among 50%, 70% or 90% global sparsity levels.
Revised: We train models using CrAM-k50, CrAM + -k70 and CrAM + -Multi, where for the latter we choose uniformly at random, at each step, among 50%, 70% or 90% global sparsity levels.

Original: Using CrAM + is crucial for preserving the dense model accuracy, at higher sparsities ( ≥ 70%) during training; however, when training with low sparsities (e.g. CrAM-k50), the resulting dense model is slightly better than the baseline.
Revised: Using CrAM + is crucial for preserving the dense model accuracy, at higher sparsities ( ≥ 70%) during training; however, when training with low sparsities (e.g. CrAM-k50), the resulting dense model is slightly better than the baseline.

Original: For both CrAM + -k70 and CrAM + -Multi we use sparse gradients for the Top-K model perturbation, as described in Section 3.2.
Revised: For both CrAM + -k70 and CrAM + -Multi we use sparse gradients for the Top-K model perturbation, as described in Section 3.2.

Original: This improved substantially the accuracy after one-shot pruning, as well as the resulting dense model.
Revised: This improved substantially the accuracy after one-shot pruning, as well as the resulting dense model.

Original: Additionally, this could offer trainingtime speed-up compared to, for example, using dense gradients or training with SAM with the right framework support.
Revised: Additionally, this could offer trainingtime speed-up compared to, for example, using dense gradients or training with SAM with the right framework support.

Original: We include an ablation study on the effects of sparse gradients in Appendix C.2.
Revised: We include an ablation study on the effects of sparse gradients in Appendix C.2.

Original: The results from Figure 1 show that CrAM models are substantially more robust to one-shot pruning, compared to standard SGD or SAM training.
Revised: The results from Figure 1 show that CrAM models are substantially more robust to one-shot pruning, compared to standard SGD or SAM training.

Original: CrAM models do not lose accuracy at lower sparsity (e.g. at 50% for all or 70% for CrAM + -k70 and CrAM + -Multi).
Revised: CrAM models do not lose accuracy at lower sparsity (e.g. at 50% for all or 70% for CrAM + -k70 and CrAM + -Multi).

Original: Moreover, as shown in Table 1, the results at higher sparsity (80% and 90%) levels are competitive with those obtained by gradual pruning, such as WoodFisher (Singh & Alistarh, 2020) or methods that prune during training– STR (Kusupati et al., 2020), DPF (Lin et al., 2020), or AC/DC (Peste et al., 2021).
Revised: Moreover, as shown in Table 1, the results at higher sparsity (80% and 90%) levels are competitive with those obtained by gradual pruning, such as WoodFisher (Singh & Alistarh, 2020) or methods that prune during training– STR (Kusupati et al., 2020), DPF (Lin et al., 2020), or AC/DC (Peste et al., 2021).

Original: We emphasize that CrAM requires a single round of training (albeit with twice as many forward-backward passes, compared to regular SGD), while standard pruning methods require training separately for each target sparsity, sometimes from a pretrained model (e.g. WoodFisher).
Revised: We emphasize that CrAM requires a single round of training (albeit with twice as many forward-backward passes, compared to regular SGD), while standard pruning methods require training separately for each target sparsity, sometimes from a pretrained model (e.g. WoodFisher).

Original: In addition to global magnitude, CrAM can be used successfully with uniform magnitude pruning; we show additional results in the Appendix C.4, as well as evidence that CrAM models are robust to sparse distributions different from those used during training.
Revised: In addition to global magnitude, CrAM can be used successfully with uniform magnitude pruning; we show additional results in the Appendix C.4, as well as evidence that CrAM models are robust to sparse distributions different from those used during training.

Original: Results for N:M sparsity patterns.
Revised: Results for N:M sparsity patterns.

Original: We show the robustness of CrAM on semi-structured N:M sparsity patterns, which can provide practical speed-ups (Mishra et al., 2021)
Revised: We show the robustness of CrAM on semi-structured N:M sparsity patterns, which can provide practical speed-ups (Mishra et al., 2021)

Original: CrAM + trained using
Revised: CrAM + trained using

Original: N:M sparsity preserve the dense model’s accuracy (77.3%), and do not lose accuracy after pruning oneshot (+BNT) to 2:4 (77.0% ) or 4:8 (77.2%) patterns.
Revised: N:M sparsity preserve the dense model’s accuracy (77.3%), and do not lose accuracy after pruning oneshot (+BNT) to 2:4 (77.0% ) or 4:8 (77.2%) patterns.

Original: This is competitive with state-of-the-art methods for training N:M sparse models Zhou et al.
Revised: This is competitive with state-of-the-art methods for training N:M sparse models Zhou et al.

Original: We provide a full discussion in Appendix C.5.
Revised: We provide a full discussion in Appendix C.5.

Original: Finetuning with CrAM.
Revised: Finetuning with CrAM.

Original: To reduce the computational overhead of CrAM-training, we investigate whether a dense model’s robustness to pruning can be improved with only a short finetuning using
Revised: To reduce the computational overhead of CrAM-training, we investigate whether a dense model’s robustness to pruning can be improved with only a short finetuning using

Original: CrAM.
Revised: CrAM.

Original: This approach is inspired by Andriushchenko & Flammarion (2022), who showed that similar benefits to full SAM training can be obtained when SAM is used only in the final training phase.
Revised: This approach is inspired by Andriushchenko & Flammarion (2022), who showed that similar benefits to full SAM training can be obtained when SAM is used only in the final training phase.

Original: We finetune pretrained ImageNet ResNet18 and ResNet50 models, from the Torchvision library, for 10 epochs, using CrAM + -k70 and CrAM + -Multi, both with sparse gradients for the pruned perturbation.
Revised: We finetune pretrained ImageNet ResNet18 and ResNet50 models, from the Torchvision library, for 10 epochs, using CrAM + -k70 and CrAM + -Multi, both with sparse gradients for the pruned perturbation.

Original: For CrAM
Revised: For CrAM

Original: + -Multi we randomly select at each step a sparsity level in the range 30%- 90%.
Revised: + -Multi we randomly select at each step a sparsity level in the range 30%- 90%.

Original: For comparison, we also finetuned using SGD with momentum or using SAM, under the same hyperparameters.
Revised: For comparison, we also finetuned using SGD with momentum or using SAM, under the same hyperparameters.

Original: We report in Tables 2 and 3 the validation accuracy for the dense models, and after one-shot pruning at 50%-80% sparsity levels.
Revised: We report in Tables 2 and 3 the validation accuracy for the dense models, and after one-shot pruning at 50%-80% sparsity levels.

Original: Finetuning with CrAM + preserves or outperforms the baseline accuracy
Revised: Finetuning with CrAM + preserves or outperforms the baseline accuracy

Original: [Table]
Revised: [Table]

Original: and results in good sparse models after one-shot pruning, at moderate sparsity levels (up to 70%).
Revised: and results in good sparse models after one-shot pruning, at moderate sparsity levels (up to 70%).

Original: Results improve with longer finetuning: after 20 epochs, the CrAM + -Multi model can be pruned one-shot to 70% sparsity, with ≤ 1% drop in accuracy, compared to the baseline
Revised: Results improve with longer finetuning: after 20 epochs, the CrAM + -Multi model can be pruned one-shot to 70% sparsity, with ≤ 1% drop in accuracy, compared to the baseline

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Sparse Transfer Experiments.
Revised: Sparse Transfer Experiments.

Original: We additionally test how well the sparse models obtained through one-shot pruning after CrAM training on ImageNet transfer across different tasks.
Revised: We additionally test how well the sparse models obtained through one-shot pruning after CrAM training on ImageNet transfer across different tasks.

Original: The setup is very similar to Salman et al.
Revised: The setup is very similar to Salman et al.

Original: ; Kornblith et al.
Revised: ; Kornblith et al.

Original: (2019), where transfer is performed acrossbenchmark tasks.
Revised: (2019), where transfer is performed acrossbenchmark tasks.

Original: Following Iofinova et al.
Revised: Following Iofinova et al.

Original: (2022), we use full finetuning of the non-zero weights, with fixed masks, and reinitialize the dense classification layer.
Revised: (2022), we use full finetuning of the non-zero weights, with fixed masks, and reinitialize the dense classification layer.

Original: We compare dense and the one-shot pruned CrAM-k50 and CrAM + -Multi models trained on ImageNet/ResNet50 to the corresponding ones obtained from SGD or SAM.
Revised: We compare dense and the one-shot pruned CrAM-k50 and CrAM + -Multi models trained on ImageNet/ResNet50 to the corresponding ones obtained from SGD or SAM.

Original: We consider the one-shot pruned models before BNT.
Revised: We consider the one-shot pruned models before BNT.

Original: In addition, we compare the transfer performance of these one-shot pruned models with standard pruning methods used in the literature, such as lottery-tickets (LTH-T) (Chen et al., 2021a),
Revised: In addition, we compare the transfer performance of these one-shot pruned models with standard pruning methods used in the literature, such as lottery-tickets (LTH-T) (Chen et al., 2021a),

Original: AC/DC (Peste et al., 2021), STR (Kusupati et al., 2020) or WoodFisher (Singh & Alistarh, 2020), using the same hyperparameters as Iofinova et al.
Revised: AC/DC (Peste et al., 2021), STR (Kusupati et al., 2020) or WoodFisher (Singh & Alistarh, 2020), using the same hyperparameters as Iofinova et al.

Original: For these pruning methods, we use public finetuned models on the downstream tasks Iofinova et al.
Revised: For these pruning methods, we use public finetuned models on the downstream tasks Iofinova et al.

Original: For each model and sparsity level, we aggregate the results over all tasks, measuring the relative increase in error of the sparse models, compared to the dense baseline Iofinova et al.
Revised: For each model and sparsity level, we aggregate the results over all tasks, measuring the relative increase in error of the sparse models, compared to the dense baseline Iofinova et al.

Original: The results in Figure 2 show that one-shot pruned CrAM models transfer well.
Revised: The results in Figure 2 show that one-shot pruned CrAM models transfer well.

Original: In fact, both CrAM-k50 and CrAM + -Multi models transfer better than LTH-T at 80% sparsity, although pruning is performed in one-shot, and LTH-T can adapt pruning to the downstream task.
Revised: In fact, both CrAM-k50 and CrAM + -Multi models transfer better than LTH-T at 80% sparsity, although pruning is performed in one-shot, and LTH-T can adapt pruning to the downstream task.

Original: Also, CrAM + -Multi at 90% sparsity has a similar transfer performance to AC/DC models, and gives better results compared to the other pruning methods used for comparison (LTH-T or STR), with the exception of the second-order WoodFisher method, which is the best performing method across both 80% and 90% sparsity.
Revised: Also, CrAM + -Multi at 90% sparsity has a similar transfer performance to AC/DC models, and gives better results compared to the other pruning methods used for comparison (LTH-T or STR), with the exception of the second-order WoodFisher method, which is the best performing method across both 80% and 90% sparsity.

Original: Compared to the standard pruning methods used for comparison, CrAM has the added advantage that it produces an accurate dense model, and both 80% and 90% models with a single ImageNet run.
Revised: Compared to the standard pruning methods used for comparison, CrAM has the added advantage that it produces an accurate dense model, and both 80% and 90% models with a single ImageNet run.

Original: [Figure]
Revised: [Figure]

Original: Quantization.
Revised: Quantization.

Original: In addition to pruning, we provide evidence that CrAM can be adapted to quantization.
Revised: In addition to pruning, we provide evidence that CrAM can be adapted to quantization.

Original: Namely, a short finetuning using a quantization version of CrAM on pretrained ImageNet models can preserve their accuracy with respect to the baseline, after symmetric per-channel 4 bits quantization, while also boosting the accuracy of the dense model.
Revised: Namely, a short finetuning using a quantization version of CrAM on pretrained ImageNet models can preserve their accuracy with respect to the baseline, after symmetric per-channel 4 bits quantization, while also boosting the accuracy of the dense model.

Original: More details can be found in Appendix C.6.
Revised: More details can be found in Appendix C.6.

Original: 4. E XPERIMENTS ON L ANGUAGE M ODELLING
Revised: 4. E XPERIMENTS ON L ANGUAGE M ODELLING

Original: In addition to image classification, we also successfully apply CrAM for language models.
Revised: In addition to image classification, we also successfully apply CrAM for language models.

Original: We demonstrate that CrAM produces models that are more compressible and accurate than the ones obtained with standard optimizers like Adam (Kingma & Ba, 2015) and SAM (Foret et al., 2021).
Revised: We demonstrate that CrAM produces models that are more compressible and accurate than the ones obtained with standard optimizers like Adam (Kingma & Ba, 2015) and SAM (Foret et al., 2021).

Original: Also, we show that CrAM models are even competitive with gradual pruning methods, which usually require a higher computational budget to produce accurate sparse models for each sparsity target independently.
Revised: Also, we show that CrAM models are even competitive with gradual pruning methods, which usually require a higher computational budget to produce accurate sparse models for each sparsity target independently.

Original: General setup.
Revised: General setup.

Original: We focus on the standard benchmark for compression methods: the BERT-base (Devlin et al., 2019) model on the span-based question-answering task SQuADv1.1 (Rajpurkar et al., 2016).
Revised: We focus on the standard benchmark for compression methods: the BERT-base (Devlin et al., 2019) model on the span-based question-answering task SQuADv1.1 (Rajpurkar et al., 2016).

Original: We consider the short fine-tuning setup (1-3 epochs) of the pretrained BERT-base model on a downstream task.
Revised: We consider the short fine-tuning setup (1-3 epochs) of the pretrained BERT-base model on a downstream task.

Original: Following the community standards (e.g. Sanh et al.
Revised: Following the community standards (e.g. Sanh et al.

Original: (2020), Kurtic et al.
Revised: (2020), Kurtic et al.

Original: we sparsify weights of the encoder part, and make use of the Top-K operator at each step to impose uniform sparsity distribution over all layers.
Revised: we sparsify weights of the encoder part, and make use of the Top-K operator at each step to impose uniform sparsity distribution over all layers.

Original: Robustness to one-shot pruning.
Revised: Robustness to one-shot pruning.

Original: We fine-tune the model with Adam, SAM, and several variants of
Revised: We fine-tune the model with Adam, SAM, and several variants of

Original: CrAM and test their robustness to one-shot pruning with the standard magnitude pruner.
Revised: CrAM and test their robustness to one-shot pruning with the standard magnitude pruner.

Original: To identify the optimal set of hyper-parameters we run a grid search (please see Appendix D for more details) and pick the one with the best one-shot performance at 50% sparsity target.
Revised: To identify the optimal set of hyper-parameters we run a grid search (please see Appendix D for more details) and pick the one with the best one-shot performance at 50% sparsity target.

Original: For a fair comparison, we allow Adam to fine-tune for twice as many epochs as SAM and CrAM.
Revised: For a fair comparison, we allow Adam to fine-tune for twice as many epochs as SAM and CrAM.

Original: The results presented in Table 4 suggest that CrAM models are more robust to one-shot pruning while still being able to match or even outperform the dense accuracy obtained with other optimizers.
Revised: The results presented in Table 4 suggest that CrAM models are more robust to one-shot pruning while still being able to match or even outperform the dense accuracy obtained with other optimizers.

Original: Comparison with gradual pruning methods.
Revised: Comparison with gradual pruning methods.

Original: We investigate whether CrAM models can be competitive with models produced by gradual pruning methods, which progressively prune smaller fractions of weights and fine-tune the model for many epochs.
Revised: We investigate whether CrAM models can be competitive with models produced by gradual pruning methods, which progressively prune smaller fractions of weights and fine-tune the model for many epochs.

Original: We adopt the CrAM + -Multi model from Table 4 and prune it in one-shot with the standard magnitude pruner, but also with the stateof-the-art BERT-pruning method called oBERT (Kurtic et al., 2022).
Revised: We adopt the CrAM + -Multi model from Table 4 and prune it in one-shot with the standard magnitude pruner, but also with the stateof-the-art BERT-pruning method called oBERT (Kurtic et al., 2022).

Original: Since one-shot pruning to high sparsity targets can severely impact the model’s performance, we also investigate whether short fine-tuning (for at most 2 epochs) on top of it can bridge the gap towards full accuracy recovery.
Revised: Since one-shot pruning to high sparsity targets can severely impact the model’s performance, we also investigate whether short fine-tuning (for at most 2 epochs) on top of it can bridge the gap towards full accuracy recovery.

Original: In Table 5 we present the results and compare against the following gradual pruning methods: ℓregularization (Louizos et al., 2018), Magnitude (Zhu & Gupta, 2017), Movement (Sanh et al., 2020), Soft-Movement (Sanh et al., 2020) and PLATON (Zhang et al., 2022).
Revised: In Table 5 we present the results and compare against the following gradual pruning methods: ℓregularization (Louizos et al., 2018), Magnitude (Zhu & Gupta, 2017), Movement (Sanh et al., 2020), Soft-Movement (Sanh et al., 2020) and PLATON (Zhang et al., 2022).

Original: For details regarding hyper-parameters, please see Appendix D. As can be seen from the results, one-shot pruned CrAM models are competitive with gradual pruning methods, which they outperform by huge margins when additionally fine-tuned for a few epochs.
Revised: For details regarding hyper-parameters, please see Appendix D. As can be seen from the results, one-shot pruned CrAM models are competitive with gradual pruning methods, which they outperform by huge margins when additionally fine-tuned for a few epochs.

Original: It is worth emphasizing that competitive results obtained with two different one-shot pruners, magnitude and oBERT, suggest that CrAM models are indeed robust and compatible with different pruning techniques than the ones they have been trained with.
Revised: It is worth emphasizing that competitive results obtained with two different one-shot pruners, magnitude and oBERT, suggest that CrAM models are indeed robust and compatible with different pruning techniques than the ones they have been trained with.

Original: We provide in Appendix D Tables 19 and 18 speed-ups for the sparse BERT models, and evidence that models become more robust to pruning even when CrAM is not used at every optimization step.
Revised: We provide in Appendix D Tables 19 and 18 speed-ups for the sparse BERT models, and evidence that models become more robust to pruning even when CrAM is not used at every optimization step.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: 4. D ETAILED C OMPARISONS WITH O THER M ETHODS
Revised: 4. D ETAILED C OMPARISONS WITH O THER M ETHODS

Original: We now perform a detailed comparison with other methods.
Revised: We now perform a detailed comparison with other methods.

Original: Since most of the other methods present experiments and are tuned on CIFAR-10 (Krizhevsky et al., 2009), we perform the comparison on this dataset as well.
Revised: Since most of the other methods present experiments and are tuned on CIFAR-10 (Krizhevsky et al., 2009), we perform the comparison on this dataset as well.

Original: Specifically, we compare CrAM with existing state-of-the-art gradual pruning methods (Lin et al., 2020) on ResNet20 (He et al., 2016), or with similar methods that train prunable networks (Miao et al., 2022) on VGG-16 (Simonyan & Zisserman, 2014).
Revised: Specifically, we compare CrAM with existing state-of-the-art gradual pruning methods (Lin et al., 2020) on ResNet20 (He et al., 2016), or with similar methods that train prunable networks (Miao et al., 2022) on VGG-16 (Simonyan & Zisserman, 2014).

Original: All hyperparameters for
Revised: All hyperparameters for

Original: CrAM are discussed in Appendix B, together with a comparison between CrAM and one-shot pruning from dense baselines in Appendix C.7.
Revised: CrAM are discussed in Appendix B, together with a comparison between CrAM and one-shot pruning from dense baselines in Appendix C.7.

Original: Comparison with Gradual Methods.
Revised: Comparison with Gradual Methods.

Original: We present results for CrAM + -Multi, trained with sparse intermediate gradients, where at each optimization step we select the sparsity level uniformly at random amongvalues in the interval [30% − 90%] .
Revised: We present results for CrAM + -Multi, trained with sparse intermediate gradients, where at each optimization step we select the sparsity level uniformly at random among the values { 50% , 70% , 80% , 90% , 95% } .

Original: The results in Table 6 show that the dense model obtained by training with CrAM + -Multi is highly accurate, and also very robust to one-shot pruning, even at high sparsities (e.g. 90%).
Revised: The results in Table 6 show that the dense model obtained by training with CrAM + -Multi is highly accurate, and also very robust to oneshot pruning, even at high sparsities (e.g. 90% and 95%).

Original: Remarkably, our results for one-shot pruning+BNT are usually competitive with those obtained by other methods which train sparse models separately, for each sparsity target, for example through DPF (Lin et al., 2020).
Revised: Remarkably, our results for one-shot pruning (+BNT) are usually competitive with those obtained by other methods which train sparse models separately, for each sparsity target, for example through DPF (Lin et al., 2020).

Original: The exception is at 95% sparsity, where DPF substantially outperforms CrAM; however, this is expected, since CrAM was not explicitly trained for such high sparsity, unlike DPF.
Revised: Notably, we substantially outperform DPF at 95% sparsity, after BNT. [Table]

Original: [Table]
Revised: [Table]

Original: Comparison with One-shot Methods.
Revised: Comparison with One-shot Methods.

Original: We compare against other methods for training prunable models, such as Miao et al.
Revised: We compare against other methods for training prunable models, such as Miao et al.

Original: ; Zimmer et al.
Revised: ; Zimmer et al.

Original: based on Stochastic Frank-Wolfe (SFW) (Reddi et al., 2016), which is used to encourage the parameters to lie in the convex hull spanned by sparse vectors, with directions given by the gradients.
Revised: based on Stochastic Frank-Wolfe (SFW) (Reddi et al., 2016), which is used to encourage the parameters to lie in the convex hull spanned by sparse vectors, with directions given by the gradients.

Original: We compare CrAM against SFW for
Revised: We compare CrAM against SFW for

Original: CIFAR10, using ResNet18 (He et al., 2016) and VGG-16 models, which is the same experimental setup employed in Miao et al.
Revised: CIFAR10, using ResNet18 (He et al., 2016) and VGG-16 models, which is the same experimental setup employed in Miao et al.

Original: (2022) and Zimmer et al.
Revised: (2022) and Zimmer et al.

Original: Note that Miao et al.
Revised: Note that Miao et al.

Original: ; Zimmer et al.
Revised: ; Zimmer et al.

Original: (2022) use the “ImageNet” variant of the ResNet18 model, which is substantially larger than the ResNet20 used for our previous experiments (11M vs. 0.3M parameters).
Revised: (2022) use the “ImageNet” variant of the ResNet18 model, which is substantially larger than the ResNet20 used for our previous experiments (11M vs. 0.3M parameters).

Original: We train CrAM + -k with sparse intermediate gradients, for 180 epochs (same as Miao et al.
Revised: We train CrAM + -k with sparse intermediate gradients, for 180 epochs (same as Miao et al.

Original: , using SGD with momentum and weight decay, and a cosine learning rate schedule.
Revised: , using SGD with momentum and weight decay, and a cosine learning rate schedule.

Original: We use BNT after pruning.
Revised: We use BNT after pruning.

Original: On both ResNet18 and VGG16, we obtain dense models that do not lose accuracy compared to the baseline: 94.2% (CrAM + -k95) vs. 93.9% (dense) on VGG16 and 95.7% (CrAM + -k95) vs. 95.4% (dense) on ResNet18.
Revised: On both ResNet18 and VGG16, we obtain dense models that do not lose accuracy compared to the baseline: 94.2% (CrAM + -k95) vs. 93.9% (dense) on VGG-16 and 95.7% (CrAM + -k95) vs. 95.4% (dense) on ResNet18.

Original: Furthermore, on ResNet18 we maintain the model accuracy after pruning one-shot at 96% sparsity (95.4%, after BNT) and have a 1.3% drop at 98% sparsity (94.4% Top-1), which is higher than Miao et al.
Revised: Furthermore, on ResNet18 we maintain the model accuracy after pruning one-shot at 96% sparsity (95.4%, after BNT) and have a 1.3% drop at 98% sparsity (94.4% Top-1), which is higher than Miao et al.

Original: (2022) and Zimmer et al.
Revised: (2022) and Zimmer et al.

Original: (2022), which obtain ≤ 93% accuracy at 95% sparsity.
Revised: (2022), which obtain ≤ 93% accuracy at 95% sparsity.

Original: We show a numerical comparison for VGG-16 in Table 6: CrAM
Revised: We show a numerical comparison for VGG-16 in Table 6: CrAM

Original: + -k95 preserves model accuracy even at 95% sparsity, which is competitive with the DPF gradual method, while SFW produces models that have lower accuracy even at higher density.
Revised: + -k95 preserves model accuracy even at 95% sparsity, which is competitive with the DPF gradual method, while SFW produces models that have lower accuracy even at higher density.

Original: We note that CrAM hasa higher training cost than SFW.
Revised: CrAM has higher training costs than SFW, but requires much less hyper-parameter tuning, and leads to higher accuracies.

Original: Also, our method executes BNT to preserve sparse accuracy, while Miao et al.
Revised: Moreover, our method uses BNT, while Miao et al.

Original: (2022); Zimmer et al.
Revised: ; Zimmer et al.

Original: (2022) do not employ it.
Revised: (2022) do not employ it.

Original: However, the cost of BNT is minimal; even without BNT , our method preserves accuracy at up to 80% sparsity (see Appendix C.7), therefore obtaining better resultsat these lower sparsities thanreported in Miao et al.
Revised: However, the cost of BNT is minimal; even without BNT , our method preserves accuracy at up to 80% sparsity (Appendix C.7), leading to better results than Miao et al.

Original: ; Zimmer et al.
Revised: ; Zimmer et al.

Original:  C ONCLUSIONS AND F UTURE W ORK
Revised:  C ONCLUSIONS AND F UTURE W ORK

Original: In this work, we proposed a new method for training neural networks, CrAM, which results in models that are both highly accurate, and easily-compressible.
Revised: In this work, we proposed a new method for training neural networks, CrAM, which results in models that are both highly accurate, and easily-compressible.

Original: Our extensive experimental analysis on large scale image classification (ImageNet/ResNets) and language modelling (SQuADv1.1/BERT-base) focuses on compression methods based on pruning, and shows that CrAM models can be pruned one-shot at a wide range of sparsity levels, while resulting in sparse models that are competitive with existing gradual pruning methods.
Revised: Our extensive experimental analysis on large scale image classification (ImageNet/ResNets) and language modelling (SQuADv1.1/BERT-base) focuses on compression methods based on pruning, and shows that CrAM models can be pruned one-shot at a wide range of sparsity levels, while resulting in sparse models that are competitive with existing gradual pruning methods.

Original: Furthermore, we show that one-shot pruned CrAM models can transfer better to downstream tasks, compared to some of the existing pruning methods.
Revised: Furthermore, we show that one-shot pruned CrAM models can transfer better to downstream tasks, compared to some of the existing pruning methods.

Original: While we focus on pruning as the main compression operator, we also give encouraging evidence that
Revised: While we focus on pruning as the main compression operator, we also give encouraging evidence that

Original: the
Revised: the

Original: CrAM update can be successfully adapted to other compression projections, such as quantization, and we plan to investigate this more closely in future work.
Revised: CrAM update can be successfully adapted to other compression projections, such as quantization, and we plan to investigate this more closely in future work.

Original: Furthermore, we would like to explore whether prolonged CrAM-training would further enhance both the performance of the resulting dense model, as well as its robustness to one-shot compression.
Revised: Furthermore, we would like to explore whether prolonged CrAM-training would further enhance both the performance of the resulting dense model, as well as its robustness to one-shot compression.

Original: Finally, we are interested in leveraging in the CrAM update different methods developed for reducing the computational complexity of SAM, in order to improve the efficiency of our method.
Revised: Finally, we are interested in leveraging in the CrAM update different methods developed for reducing the computational complexity of SAM, in order to improve the efficiency of our method.

Original: A PPENDIX A E THICS AND R EPRODUCIBILITY S TATEMENTS
Revised: A PPENDIX A E THICS AND R EPRODUCIBILITY S TATEMENTS

Original: Our work is technical in nature, considering existing models and datasets, and thus does not pose any direct ethical concerns.
Revised: Our work is technical in nature, considering existing models and datasets, and thus does not pose any direct ethical concerns.

Original: With regards to reproducibility, we plan to release an implementation of our optimizer and an example experimental harness during the discussion period, as suggested in the
Revised: With regards to reproducibility, we plan to release an implementation of our optimizer and an example experimental harness during the discussion period, as suggested in the

Original: ICLR guide to authors.
Revised: ICLR guide to authors.

Original: B I MAGE C LASSIFICATION H YPERPARAMETERS
Revised: B I MAGE C LASSIFICATION H YPERPARAMETERS

Original: Hyperparameters for CIFAR10 experiments We train ResNet20 models for 200 epochs, using SGD with momentum and weight decay, and a cosine learning rate scheduler, with a learning rate warm-up of 5 epochs.
Revised: Hyperparameters for CIFAR10 experiments We train ResNet20 models for 200 epochs, using SGD with momentum and weight decay, and a cosine learning rate scheduler, with a learning rate warm-up of 5 epochs.

Original: Additionally, we trained the baseline model for twice as many epochs, to match the number of backpropagation steps of SAM and CrAM.
Revised: Additionally, we trained the baseline model for twice as many epochs, to match the number of backpropagation steps of SAM and CrAM.

Original: To determine the value of the hyperparameter ρ , we performed a grid search over values in the range 0 .
Revised: To determine the value of the hyperparameter ρ , we performed a grid search over values in the range 0 .

Original: 01 − 0 .
Revised: 01 − 0 .

Original: 2 , using a 90% − 10% train-validation split and found 0.1 and 0.2 to be the best values for SAM and CrAM, respectively (achieving highest validation accuracy).
Revised: 2 , using a 90% − 10% train-validation split and found 0.1 and 0.15 to be the best values for SAM and CrAM + -Multi, respectively (achieving highest validation accuracy).

Original: After finding the best value of ρ for each model configuration, we retrained using the entire training set, and starting from 3 different random seeds, and report the final accuracy after 200 epochs of training.
Revised: After finding the best value of ρ , we retrained using the entire training set, and starting from 3 different random seeds, and report the final accuracy after 200 epochs of training.

Original: We follow a very similar training recipe and hyperparameter search for ResNet18 and VGG experiments, but train instead of 180 epochs.
Revised: We follow a very similar training recipe and hyperparameter search for ResNet18 and VGG experiments, but train instead of 180 epochs.

Original: Hyperparameters for ImageNet experiments For our ImageNet experiments, we use standard data augmentation, and we train the models using SGD for 100 epochs, with batch size 512, momentum 0.9, and weight decay 0.0001.
Revised: Hyperparameters for ImageNet experiments For our ImageNet experiments, we use standard data augmentation, and we train the models using SGD for 100 epochs, with batch size 512, momentum 0.9, and weight decay 0.0001.

Original: The learning rate is linearly increased for the first 5 epochs until it reaches a maximum value of 0.2, after which it is decreased at each epoch, using a cosine scheduler.
Revised: The learning rate is linearly increased for the first 5 epochs until it reaches a maximum value of 0.2, after which it is decreased at each epoch, using a cosine scheduler.

Original: To determine the value of the hyperparameter ρ , we search over a small grid, by training 90% of
Revised: To determine the value of the hyperparameter ρ , we search over a small grid, by training 90% of

Original: ImageNet using TopK-CrAM-k50, and using the remaining 10% of the dataset for validation.
Revised: ImageNet using CrAM-k50, and using the remaining 10% of the dataset for validation.

Original: We have found ρ = 0 .
Revised: We have found ρ = 0 .

Original: 05 to give good results for CrAM-k50, and we have kept this value for all our other
Revised: 05 to give good results for CrAM-k50, and we have kept this value for all our other

Original: CrAM experiments.
Revised: CrAM experiments.

Original: For SAM, we also use ρ = 0 .
Revised: For SAM, we also use ρ = 0 .

Original: 05 , which is the standard value recommended by the authors of Foret et al.
Revised: 05 , which is the standard value recommended by the authors of Foret et al.

Original: C A DDITIONAL I MAGE C LASSIFICATION E XPERIMENTS C. A BLATION S TUDY FOR A LTERNATIVE U PDATES
Revised: C A DDITIONAL I MAGE C LASSIFICATION E XPERIMENTS C. A BLATION S TUDY FOR A LTERNATIVE U PDATES

Original: Comparison between CrAM and C-SAM.
Revised: Comparison between CrAM and C-SAM.

Original: We investigate the importance of individual components from the CrAM update by comparing against other similar updates, on the CIFAR10 dataset, using a
Revised: We investigate the importance of individual components from the CrAM update by comparing against other similar updates, on the CIFAR10 dataset, using a

Original: ResNet20 model.
Revised: ResNet20 model.

Original: One such update can be obtained by following closely the derivations for SAM (Foret et al., 2021).
Revised: One such update can be obtained by following closely the derivations for SAM (Foret et al., 2021).

Original: We assume ∥ δ ∥ ≤ ρ , define h ( x ) :
Revised: We assume ∥ δ ∥ ≤ ρ , define h ( x ) :

Original: =
Revised: =

Original: L ( C ( x ))
Revised: L ( C ( x ))

Original: (with C the Top-K operator), and the loss max ∥ δ ∥≤ ρ h ( θ + δ ) .
Revised: (with C the Top-K operator), and the loss max ∥ δ ∥≤ ρ h ( θ + δ ) .

Original: By using a first-order Taylor approximation of h ( θ + δ ) around θ , together with the quadratic constraint for δ , we obtain δ = ρ ∇ L ( C ( θ ))
Revised: By using a first-order Taylor approximation of h ( θ + δ ) around θ , together with the quadratic constraint for δ , we obtain δ = ρ ∇ L ( C ( θ ))

Original: ∥∇ L ( C ( θ ))
Revised: ∥∇ L ( C ( θ ))

Original: ∥ .
Revised: ∥ .

Original: This enables us to define the compressed-SAM (C-SAM) update as:
Revised: This enables us to define the compressed-SAM (C-SAM) update as:

Original: [Equation]
Revised: [Equation]

Original: We observed that C-SAM training benefits from using sparsified gradients, for both the intermediate interpolation step, as well as in the final weight update.
Revised: We observed that C-SAM training benefits from using sparsified gradients, for both the intermediate interpolation step, as well as in the final weight update.

Original: Namely, we always use the approximation: ∇ L ( C ( ϕ ))
Revised: Namely, we always use the approximation: ∇ L ( C ( ϕ ))

Original: ≈ M ϕ · ∇ L ( M ϕ · ϕ ) for parameter ϕ , where M ϕ is the Top-K mask of ϕ .
Revised: ≈ M ϕ · ∇ L ( M ϕ · ϕ ) for parameter ϕ , where M ϕ is the Top-K mask of ϕ .

Original: To determine the value of the interpolation step ρ in C-SAM and CrAM, we perform a grid search over a 90-10% train/validation of CIFAR10; we re-run from 3 different seeds, using the best configurations and the entire training set, and report the test accuracy after 200 epochs of training.
Revised: To determine the value of the interpolation step ρ in C-SAM and CrAM, we perform a grid search over a 90-10% train/validation of CIFAR10; we re-run from 3 different seeds, using the best configurations and the entire training set, and report the test accuracy after 200 epochs of training.

Original: For fairness, we compare C-SAM with CrAM, and not CrAM + , and for both we use sparsified gradients.
Revised: For fairness, we compare C-SAM with CrAM, and not CrAM + , and for both we use sparsified gradients.

Revised: In the case of C-SAM-Multi and CrAM Multi, we choose the sparsity levels during training uniformly at random among values in the interval 30% − 90% .

Original: The results in Table 7 show that C-SAM can be more robust at higher sparsity levels, but with the cost of an accuracy drop for the dense models.
Revised: The results in Table 7 show that C-SAM can be more robust at higher sparsity levels, but with the cost of an accuracy drop for the dense models.

Original: Moreover, the dense model can be improved using CrAM + with no additional cost, whereas for C-SAM such a modification would require a computational overhead.
Revised: Moreover, the dense model can be improved using CrAM + with no additional cost, whereas for C-SAM such a modification would require a computational overhead.

Original: [Table]
Revised: [Table]

Original: Importance of extra-gradient step.
Revised: Importance of extra-gradient step.

Original: Furthermore, we explore the importance of the extra-gradient step in the CrAM update.
Revised: Furthermore, we explore the importance of the extra-gradient step in the CrAM update.

Original: Notably, we investigate whether not using the extra-gradient achieves a similar effect to CrAM training.
Revised: Notably, we investigate whether not using the extra-gradient achieves a similar effect to CrAM training.

Original: The removal of the extra gradient step would correspond to the following equation:
Revised: The removal of the extra gradient step would correspond to the following equation:

Original: [Equation]
Revised: [Equation]

Revised: ∇

Original:  Since the compression we use in our experiments is the Top-K sparsification, we simply call this update “Top-K”.
Revised: [Equation] Since the compression we use in our experiments is the Top-K sparsification [Table] we simply call this update “Top-K”.

Original: This update has been previously studied in Lin et al.
Revised: This update has been previously studied in Lin et al.

Original: (2020), where the dense gradients, computed with respect to the sparse parameters, are used in the model updated.
Revised: [Table] where the dense gradients, computed with respect to the sparse parameters, are used in the model updated.

Original: Generally, we have training instability using this update, particularly at high sparsity.
Revised: Generally, we have training instability using this update, particularly at high sparsity.

Original: However, incorporating the optimization of the dense model, as well as sparsified gradients for the compressed parameters, greatly improved the stability and overall quality of the resulting models.
Revised: However, incorporating the optimization of the dense model, as well as sparsified gradients for the compressed parameters, greatly improved the stability and overall quality of the resulting models.

Original: These changes resulted in an update close to CrAM + and of the same computational complexity, which will be referred to as
Revised: These changes resulted in an update close to CrAM + and of the same computational complexity, which will be referred to as

Original: Top-K + :
Revised: Top-K + :

Original: [Equation]
Revised: [Equation]

Original: where (cid:101) θ t =
Revised: where (cid:101) θ t =

Original: C ( θ t ) and M (cid:101) θ t is its mask after applying Top-K
Revised: C ( θ t ) and M (cid:101) θ t is its mask after applying Top-K

Original: [Table]
Revised: [Table]

Original: The results of the comparison between CrAM + and Top-K + are presented in Table 8 and show that
Revised: The results of the comparison between CrAM + and Top-K + are presented in Table 8 and show that

Original: CrAM models tend to have higher accuracy for the dense models, as well as to be more robust to one-shot pruning at high sparsity (e.g. CrAM + and Top-K + trained with 50% or 70% sparsity, and one-shot pruned to 80% and 90% sparsity).
Revised: CrAM models tend to have higher accuracy for the dense models, as well as to be more robust to one-shot pruning at high sparsity (e.g. CrAM + and Top-K + trained with 50% or 70% sparsity, and one-shot pruned to 80% and 90% sparsity).

Original: The comparison between CrAM and C-SAM or Top-K shows that, although all methods can achieve good results with one-shot pruning, CrAM-trained models have the best trade-off between preserving (or improving) the dense model accuracy, while having good performance after one-shot pruning, at a range of sparsity levels.
Revised: The comparison between CrAM and C-SAM or Top-K shows that, although all methods can achieve good results with one-shot pruning, CrAM-trained models have the best trade-off between preserving (or improving) the dense model accuracy, while having good performance after one-shot pruning, at a range of sparsity levels.

Original: C. I MPORTANCE OF S PARSE G RADIENTS
Revised: C. I MPORTANCE OF S PARSE G RADIENTS

Original: For all our image classification experiments, we observed an improvement in the robustness to posttraining one-shot pruning, when using a different straight-through estimator for the gradient ∇ θ L ( (cid:101) θ t ) , where (cid:101) θ t =
Revised: For all our image classification experiments, we observed an improvement in the robustness to posttraining one-shot pruning, when using a different straight-through estimator for the gradient ∇ θ L ( (cid:101) θ t ) , where (cid:101) θ t =

Original: C ( θ t + ρ ∇ L ( θ t )) .
Revised: C ( θ t + ρ ∇ L ( θ t )) .

Original: Namely, instead of by-passing the Top-K operator in the gradient, and estimating ∇ θ t L ( (cid:101) θ t ) ≈ ∇ (cid:101) θ t L ( (cid:101) θ t ) , we can assume instead that the masks M t of (cid:101) θ t change very little during training.
Revised: Namely, instead of by-passing the Top-K operator in the gradient, and estimating ∇ θ t L ( (cid:101) θ t ) ≈ ∇ (cid:101) θ t L ( (cid:101) θ t ) , we can assume instead that the masks M t of (cid:101) θ t change very little during training.

Original: This would allow us to use the approximation ∇ θ t L ( (cid:101) θ t )
Revised: This would allow us to use the approximation ∇ θ t L ( (cid:101) θ t )

Original: ≈ M t ·∇ (cid:101) θ t L ( (cid:101) θ t ) .
Revised: ≈ M t ·∇ (cid:101) θ t L ( (cid:101) θ t ) .

Original: Please see
Revised: Please see

Original: Section 3.2 and Appendix Section E for more details.
Revised: Section 3.2 and Appendix Section E for more details.

Original: Training both CrAM and CrAM + with this new approximation for the CrAM loss gradient (which will be referred to as “sparse gradient”) led to an improvement in the robustness to one-shot pruning, particularly at higher sparsity levels.
Revised: Training both CrAM and CrAM + with this new approximation for the CrAM loss gradient (which will be referred to as “sparse gradient”) led to an improvement in the robustness to one-shot pruning, particularly at higher sparsity levels.

Original: Furthermore, our intuition that the masks are fairly constant during training is also confirmed experimentally: on CIFAR10/ResNet20, trained with CrAM + -k70, the difference between consecutive θ t masks was lower than 0.6%.
Revised: Furthermore, our intuition that the masks are fairly constant during training is also confirmed experimentally: on CIFAR10/ResNet20, trained with CrAM + -k70, the difference between consecutive θ t masks was lower than 0.6%.

Original: Interestingly, using sparse gradient under the same setup, encouraged more diversity in the masks, with the difference between them at later training stages increasing to around 2%.
Revised: Interestingly, using sparse gradient under the same setup, encouraged more diversity in the masks, with the difference between them at later training stages increasing to around 2%.

Original: We speculate this could be a potential reason for the improved robustness to pruning.
Revised: We speculate this could be a potential reason for the improved robustness to pruning.

Original: Another aspect observed on CIFAR10 experiments is that using sparse gradients tends to decrease the dense model accuracy, when training CrAM at lower sparsity; for example, the dense model for CrAM-k reached 93.4% accuracy, which decreased to 92.8% when using sparse gradients.
Revised: Another aspect observed on CIFAR10 experiments is that using sparse gradients tends to decrease the dense model accuracy, when training CrAM at lower sparsity; for example, the dense model for CrAM-k reached 93.4% accuracy, which decreased to 92.8% when using sparse gradients.

Original: For this reason, on
Revised: For this reason, on

Original: ImageNet we only experimented with the dense version of CrAM-k50.
Revised: ImageNet we only experimented with the dense version of CrAM-k50.

Original: Nonetheless, using sparse gradients improved the robustness to pruning in all cases.
Revised: Nonetheless, using sparse gradients improved the robustness to pruning in all cases.

Original: For a better illustration of all these effects, we provide the results obtained with these different versions of CrAM on ImageNet, in Table 9 for one-shot unstructured global magnitude pruning and in Table 10 for semi-structured N:M pruning.
Revised: For a better illustration of all these effects, we provide the results obtained with these different versions of CrAM on ImageNet, in Table 9 for one-shot unstructured global magnitude pruning and in Table 10 for semi-structured N:M pruning.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: C. V ARIABILITY OF B ATCH N ORM T UNING R ESULTS
Revised: C. V ARIABILITY OF B ATCH N ORM T UNING R ESULTS

Original: We emphasize that CrAM relies on a small calibration set of training samples to correct the Batch
Revised: We emphasize that CrAM relies on a small calibration set of training samples to correct the Batch

Original: Norm statistics, namely running mean and variance, after pruning, particularly at high sparsity.
Revised: Norm statistics, namely running mean and variance, after pruning, particularly at high sparsity.

Original: We call this procedure Batch Norm Tuning (BNT).
Revised: We call this procedure Batch Norm Tuning (BNT).

Original: To ensure that the accuracies we report for sparse models are stable under the choice of the calibration set, we perform 10 independent trials of BNT, on 10 randomly chosen subsets of 1000 training samples, for each model and for different sparsity levels.
Revised: To ensure that the accuracies we report for sparse models are stable under the choice of the calibration set, we perform 10 independent trials of BNT, on 10 randomly chosen subsets of 1000 training samples, for each model and for different sparsity levels.

Original: The results of this experiment are presented in Table [Table] which also contains the raw numbers used in
Revised: The results of this experiment are presented in Table 11, which also contains the raw numbers used in

Original: Figure 1.
Revised: Figure 1.

Original: Notice that the accuracy after one-shot pruning and BNT is very stable, with respect to the choice of the calibration set.
Revised: Notice that the accuracy after one-shot pruning and BNT is very stable, with respect to the choice of the calibration set.

Original: In particular, for the CrAM + model, the standard deviation is ≤ 0 .
Revised: In particular, for the CrAM + model, the standard deviation is ≤ 0 .

Original: 1% across all sparsity levels considered.
Revised: 1% across all sparsity levels considered.

Original: We also report the raw one-shot pruning accuracy, before BNT, for CrAM models in Table
Revised: We also report the raw one-shot pruning accuracy, before BNT, for CrAM models in Table

Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: C. R ESULTS WITH U NIFORM S PARSITY
Revised: C. R ESULTS WITH U NIFORM S PARSITY

Original: In this section we show that CrAM models can be trained to be robust to different sparsity distributions, such as uniform.
Revised: In this section we show that CrAM models can be trained to be robust to different sparsity distributions, such as uniform.

Original: We train CrAM + -Multi models for ImageNet/ResNet50 under the same setupas in Section 4.
Revised: We train CrAM + -Multi models for ImageNet/ResNet50 under the same setup

Original: [Table]
Revised: [Table]

Original:  but applying instead the Top-K operator at uniform sparsity across all prunable parameters (i [Table] e. excluding BatchNorm and biases), while keeping the first and last layers dense [Table]
Revised: as in Section 4.1, but applying instead the Top-K operator at uniform sparsity across all prunable parameters (i.e. excluding BatchNorm and biases), while keeping the first and last layers dense.

Original: [Table] The resulting dense model achieves 77.1% accuracy, while one-shot uniform pruning at 80% and 90% sparsities gives, after BNT, 75.6% and 75.1% accuracy, respectively.
Revised: The resulting dense model achieves 77.1% accuracy, while one-shot uniform pruning at 80% and 90% sparsities gives, after BNT, 75.6% and 75.1% accuracy, respectively.

Original: Moreover, this model is also robust to one-shot pruning using global magnitude (e.g. 75.4% accuracy at 80% sparsity).
Revised: Moreover, this model is also robust to one-shot pruning using global magnitude (e.g. 75.4% accuracy at 80% sparsity).

Original: Conversely, CrAM + -Multi trained with global magnitude is robust to one-shot pruning using uniform magnitude (e.g. 77.0% and 75.9% accuracy at 70% and 80% sparsity, respectively).
Revised: Conversely, CrAM + -Multi trained with global magnitude is robust to one-shot pruning using uniform magnitude (e.g. 77.0% and 75.9% accuracy at 70% and 80% sparsity, respectively).

Original: This suggests that CrAM-trained models can be robust to one-shot pruning using sparsity distributions different from the ones used during training.
Revised: This suggests that CrAM-trained models can be robust to one-shot pruning using sparsity distributions different from the ones used during training.

Original: C. C R AM FOR N:M S PARSITY P ATTERNS
Revised: C. C R AM FOR N:M S PARSITY P ATTERNS

Original: In this section, we show our full results regarding the robustness of CrAM models against semistructured N:M sparsity patterns, where out of each block of M weights, N are sparse.
Revised: In this section, we show our full results regarding the robustness of CrAM models against semistructured N:M sparsity patterns, where out of each block of M weights, N are sparse.

Original: In particular, the 2:4 pattern is supported on modern Ampere NVIDIA architectures, where it has been shown to provide speed-ups (Mishra et al., 2021).
Revised: In particular, the 2:4 pattern is supported on modern Ampere NVIDIA architectures, where it has been shown to provide speed-ups (Mishra et al., 2021).

Original: We train CrAM + models with the N:M pattern, by randomly choosing at each optimization step between the 2:4 or 4:8 projections; this model will be referred to as “CrAM + -N:M”.
Revised: We train CrAM + models with the N:M pattern, by randomly choosing at each optimization step between the 2:4 or 4:8 projections; this model will be referred to as “CrAM + -N:M”.

Original: Similar to the previous experiments, we also use sparse gradients for the pruned model perturbation and have found this to have a positive impact on the one-shot pruned models.
Revised: Similar to the previous experiments, we also use sparse gradients for the pruned model perturbation and have found this to have a positive impact on the one-shot pruned models.

Original: In Table 13 we show the one-shot pruning results (after BNT with 1000 samples).
Revised: In Table 13 we show the one-shot pruning results (after BNT with 1000 samples).

Original: Note that CrAM + -N:
Revised: Note that CrAM + -N:

Original: M models do not lose accuracy when they are pruned one-shot using the 2: and 4:8 patterns, which is competitve with state-of-the-art methods for training N:M sparse models, such as SR-STE (Zhou et al., 2021); however, SR-STE requires a different training run for each sparsity profile.
Revised: M models do not lose accuracy when they are pruned one-shot using the 2: and 4:8 patterns, which is competitve with state-of-the-art methods for training N:M sparse models, such as SR-STE (Zhou et al., 2021); however, SR-STE requires a different training run for each sparsity profile.

Original: CrAM + -N:M model is also robust to one-shot pruning using unstructured patterns, at moderate sparsity levels; for example, the results in Table 14 show that models trained with
Revised: CrAM + -N:M model is also robust to one-shot pruning using unstructured patterns, at moderate sparsity levels; for example, the results in Table 14 show that models trained with

Original: CrAM + -N:M can be pruned one-shot to 70% sparsity with a minor accuracy loss, compared to the dense baseline
Revised: CrAM + -N:M can be pruned one-shot to 70% sparsity with a minor accuracy loss, compared to the dense baseline.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: C. R ESULTS ON Q UANTIZATION
Revised: C. R ESULTS ON Q UANTIZATION

Original: We have shown through previous experiments that CrAM can be successfully used with the TopK operator to obtain models that preserve or improve the dense baseline’s accuracy, while also being robust to post-training one-shot pruning, to multiple sparsity levels.
Revised: We have shown through previous experiments that CrAM can be successfully used with the TopK operator to obtain models that preserve or improve the dense baseline’s accuracy, while also being robust to post-training one-shot pruning, to multiple sparsity levels.

Original: In this section we show encouraging evidence that CrAM can also be adapted to work with quantization.
Revised: In this section we show encouraging evidence that CrAM can also be adapted to work with quantization.

Original: Specifically, we use
Revised: Specifically, we use

Original: CrAM + where the compression operator C is the symmetric per-channel weight quantization tobits (round to nearest integer), and finetune pretrained Torchivison ResNet18 and ResNet50 ImageNet models, for 10 epochs, using the same hyperparameters as in the previous experiments for sparsity.
Revised: CrAM + where the compression operator C is the symmetric per-channel weight quantization tobits (round to nearest integer), and finetune pretrained Torchivison ResNet18 and ResNet50 ImageNet models, for 10 epochs, using the same hyperparameters as in the previous experiments for sparsity.

Original: [Table]
Revised: [Table]

Original: C. C OMPARISON W ITH O THER M ETHODS ON CIFAR 
Revised: C. C OMPARISON W ITH O THER M ETHODS ON CIFAR

Original: In this section we provide additional results accompanying those presented in Section 4.3.
Revised: In this section we provide additional results accompanying those presented in Section 4.3.

Original: Namely, we provide comparison between one-shot pruning CrAM + -Multi vs. standard dense baselines (SGD,
Revised: Namely, we provide comparison between one-shot pruning CrAM + -Multi vs. standard dense baselines (SGD,

Original: SAM), and we provide numbers before and after BNT for sparse models on VGG-16 and ResNet18.
Revised: SAM), and we provide numbers before and after BNT for sparse models on VGG-16 and ResNet18.

Original: In Table 16 we show the accuracy of the dense baseline, SAM and CrAM + -Multi (the same from Section 4.3) on ResNet20, before and after one-shot pruning at different sparsities.
Revised: In Table 16 we show the accuracy of the dense baseline, SAM and CrAM + -Multi, trained using sparsity levels sampled uniformly at random in the range 30% − 90% on ResNet20, before and after one-shot pruning at different sparsities.

Original: The results after one-shot pruning are presented after BNT over a random subset of 1000 train samples, overbatches.
Revised: The results after one-shot pruning are presented after BNT over a random subset of 1000 train samples, over 100 batches.

Original: We note there are small variations in the results after BNT, due to the choice of the random calibration set.
Revised: We note there are small variations in the results after BNT, due to the choice of the random calibration set.

Original: These variations are small ( ± 0.1/0.2%) for CrAM + -Multi models, across all sparsity levels considered, but they are larger for the one-shot pruned dense baselines at high sparsity (e.g. 80% and 90%).
Revised: These variations are small ( ± 0.1/0.2%) for CrAM + -Multi models, across all sparsity levels considered, but they are larger for the one-shot pruned dense baselines at high sparsity (e.g. 80% and 90%).

Original: Moreover, the accuracy before BNT is still high for CrAM at lower sparsity levels (e.g. 91.9% at 70% sparsity), but it degrades at high sparsity (e.g. 50.5% at 90% sparsity).
Revised: Moreover, the accuracy before BNT is still high for CrAM at lower sparsity levels (e.g. 91.9% at 70% sparsity), but it degrades at high sparsity (e.g. 50.5% at 90% sparsity).

Original: We believe that this is only due to the BatchNorm statistics, which are adapted to the dense model during training, but they no longer reflect the distribution shift after weight pruning.
Revised: We believe that this is only due to the BatchNorm statistics, which are adapted to the dense model during training, but they no longer reflect the distribution shift after weight pruning.

Original: This is confirmed by the fact that 90% sparse models improve to over 90% test accuracy after only a few iterations of BNT, and are very robust to the choice of the calibration set.
Revised: This is confirmed by the fact that 90% sparse models improve to over 90% test accuracy after only a few iterations of BNT, and are very robust to the choice of the calibration set.

Original: [Table]
Revised: [Table]

Original: Moreover
Revised: Moreover

Original: [Table]
Revised: [Table]

Original: we show the extended results of CrAM + -k95 discussed in Section 4.
Revised: we show the extended results of CrAM + -k95 discussed in Section 4.

Original: [Table]
Revised: [Table]

Original: before and after
Revised: before and after

Original: BNT [Table] on ResNet18 and VGG-16.
Revised: BNT, on ResNet18 and VGG-16.

Original: From Table 17 we can see that one-shot pruning CrAM + -k without BNT preserves accuracy up to 80% sparsity, after which BNT is required to correct the
Revised: From Table 17 we can see that one-shot pruning CrAM + -k without BNT preserves accuracy up to 80% sparsity, after which BNT is required to correct the

Original: BatchNorm statistics.
Revised: BatchNorm statistics.

Original: Remarkably, the VGG-16 models at 97% and 98% sparsity have very low accuracy, which is improved greatly by BNT.
Revised: Remarkably, the VGG-16 models at 97% and 98% sparsity have very low accuracy, which is improved greatly by BNT.

Original: Furthermore, also in this highly sparse regimes the accuracy is very robust with respect to the choice of the calibration set for BNT
Revised: Furthermore, also in this highly sparse regimes the accuracy is very robust with respect to the choice of the calibration set for BNT

Original: [Table]
Revised: [Table]

Original: D L ANGUAGE M ODELS - REPRODUCIBILITY AND HYPERPARAMETERS
Revised: D L ANGUAGE M ODELS - REPRODUCIBILITY AND HYPERPARAMETERS

Original: To ease reproducibility of our results
Revised: To ease reproducibility of our results

Original: [Table]
Revised: [Table]

Original: we conduct all of our experiments with the popular open-source libraries: Transformers (Wolf et al., 2020)
Revised: we conduct all of our experiments with the popular open-source libraries: Transformers (Wolf et al., 2020)

Original: [Table]
Revised: [Table]

Original: and SparseML (Kurtz et al., 2020).
Revised: and SparseML (Kurtz et al., 2020).

Original: We use the publicly available datasets via Lhoest et al.
Revised: We use the publicly available datasets via Lhoest et al.

Original: [Table]
Revised: [Table]

Original: and focus on the BERT-base (Devlin et al., 2019) as it is one of the most commonly used language models.
Revised: and focus on the BERT-base (Devlin et al., 2019) as it is one of the most commonly used language models.

Original: It is composed of 12 identical transformer layers with 110M parameters.
Revised: It is composed of 12 identical transformer layers with 110M parameters.

Original: Following community standards [Table] we prune all weights of the encoder part (85M) and report sparsities relative to this number.
Revised: Following community standards, we prune all weights of the encoder part (85M) and report sparsities relative to this number.

Original: General setup.
Revised: General setup.

Original: Our SQuADv1.1 fine-tuning recipe with Adam, SAM and CrAM mostly follows the already established hyper-parameters (Devlin et al., 2019; Wolf et al., 2020): start from the pretrained bert-base-uncased (available for download at https://huggingface.co/bert-baseuncased), batch-size=16 , max-sequence-length=384 , doc-stride=128 .
Revised: Our SQuADv1.1 fine-tuning recipe with Adam, SAM and CrAM mostly follows the already established hyper-parameters (Devlin et al., 2019; Wolf et al., 2020): start from the pretrained bert-base-uncased (available for download at https://huggingface.co/bert-baseuncased), batch-size=16 , max-sequence-length=384 , doc-stride=128 .

Original: Adam, SAM, and CrAM optimization.
Revised: Adam, SAM, and CrAM optimization.

Original: For other hyper-parameters we conduct a grid search for each optimizer independently over the following values: learning-rate ∈ { 3e − 5 , 5e − 5 , 8e − 5 } ; num-train-epochs ∈ { 2 , 3 } for SAM and CrAM, and num-train-epochs ∈ { 2 , 3 , 4 , 6 } for
Revised: For other hyper-parameters we conduct a grid search for each optimizer independently over the following values: learning-rate ∈ { 3e − 5 , 5e − 5 , 8e − 5 } ; num-train-epochs ∈ { 2 , 3 } for SAM and CrAM, and num-train-epochs ∈ { 2 , 3 , 4 , 6 } for

Original: Adam (we allow 2x more epochs for fairness to SAM and CrAM); label-smoothing-factor ∈ { 0 . 0 , 0 . 1 , 0 . 2 } .
Revised: Adam (we allow 2x more epochs for fairness to SAM and CrAM); label-smoothing-factor ∈ { 0 . 0 , 0 . 1 , 0 . 2 } .

Original: We freeze the embedding layer in all experiments.
Revised: We freeze the embedding layer in all experiments.

Original: To determine the value of the hyperparameter ρ , we performed a grid search over values in the range 1e − 4 to 1e − 1 .
Revised: To determine the value of the hyperparameter ρ , we performed a grid search over values in the range 1e − 4 to 1e − 1 .

Original: For each optimizer we pick the set of hyperparameters that produces the best results after one-shot magnitude pruning to 50% sparsity, and they are as follows.
Revised: For each optimizer we pick the set of hyperparameters that produces the best results after one-shot magnitude pruning to 50% sparsity, and they are as follows.

Original: Adam: num-train-epochs=2 , learning-rate= 8e − 5 , label-smoothing-ratio=0.1 .
Revised: Adam: num-train-epochs=2 , learning-rate= 8e − 5 , label-smoothing-ratio=0.1 .

Original: SAM: num-train-epochs=2 , learning-rate= 8e − 5 , label-smoothing-ratio=0.0 , ρ =0.01 .
Revised: SAM: num-train-epochs=2 , learning-rate= 8e − 5 , label-smoothing-ratio=0.0 , ρ =0.01 .

Original: CrAM (all CrAM runs use the same set of hyperparameters): num-train-epochs=3 , learning-rate= 8e − 5 , label-smoothing-ratio=0.2 , ρ =0.005 .
Revised: CrAM (all CrAM runs use the same set of hyperparameters): num-train-epochs=3 , learning-rate= 8e − 5 , label-smoothing-ratio=0.2 , ρ =0.005 .

Original: At each CrAM optimization step we apply Top-K (i.e. magnitude) sparsification over all layers uniformly.
Revised: At each CrAM optimization step we apply Top-K (i.e. magnitude) sparsification over all layers uniformly.

Original: One-shot pruning.
Revised: One-shot pruning.

Original: We apply one-shot pruning with two different pruners: magnitude and oBERT.
Revised: We apply one-shot pruning with two different pruners: magnitude and oBERT.

Original: For one-shot magnitude pruning we impose uniform sparsity distribution over all layers.
Revised: For one-shot magnitude pruning we impose uniform sparsity distribution over all layers.

Original: For one-shot oBERT pruning we adopt the suggested set of hyper-parameters by authors, which we briefly describe here for completeness: 1024 gradients, dampening 1e − 7 , block-size 50, 4 recomputations, global sparsity distribution over all layers.
Revised: For one-shot oBERT pruning we adopt the suggested set of hyper-parameters by authors, which we briefly describe here for completeness: 1024 gradients, dampening 1e − 7 , block-size 50, 4 recomputations, global sparsity distribution over all layers.

Original: For more details please refer to the oBERT paper (Kurtic et al., 2022).
Revised: For more details please refer to the oBERT paper (Kurtic et al., 2022).

Original: Sparse fine-tuning of one-shot pruned models.
Revised: Sparse fine-tuning of one-shot pruned models.

Original: We fine-tune one-shot oBERT-pruned models with the fixed sparsity mask and Adam optimizer.
Revised: We fine-tune one-shot oBERT-pruned models with the fixed sparsity mask and Adam optimizer.

Original: To identify the best set of hyperparameters for fine-tuning of the sparse model, we conduct a grid search over the following parameters: learning-rate ∈ { 3e − 5 , 5e − 5 , 8e − 5 , 1e − 4 } , num-train-epochs ∈ { 1 , 2 } , label-smoothing-ratio ∈ { 0 . 0 , 0 . 2 } , warmup-ratio ∈ { 0 . 0 , 0 . 1 } .
Revised: To identify the best set of hyperparameters for fine-tuning of the sparse model, we conduct a grid search over the following parameters: learning-rate ∈ { 3e − 5 , 5e − 5 , 8e − 5 , 1e − 4 } , num-train-epochs ∈ { 1 , 2 } , label-smoothing-ratio ∈ { 0 . 0 , 0 . 2 } , warmup-ratio ∈ { 0 . 0 , 0 . 1 } .

Original: We freeze the embedding layer and employ earlystopping technique to prevent overfitting.
Revised: We freeze the embedding layer and employ earlystopping technique to prevent overfitting.

Original: Speed-ups of pruned BERT-base models.
Revised: Speed-ups of pruned BERT-base models.

Original: In Table 19 we present speed-ups of our pruned models in the sparsity-aware CPU inference engine DeepSparse (Kurtz et al., 2020; NeuralMagic, 2021) (version 1.0.2).
Revised: In Table 19 we present speed-ups of our pruned models in the sparsity-aware CPU inference engine DeepSparse (Kurtz et al., 2020; NeuralMagic, 2021) (version 1.0.2).

Original: We consider two different scenarios and report speed-ups relative to the dense model benchmarked in the same environment.
Revised: We consider two different scenarios and report speed-ups relative to the dense model benchmarked in the same environment.

Original: Robustness to one-shot pruning.
Revised: Robustness to one-shot pruning.

Original: In Table 18 we present results for runs where CrAM is not used at every optimization step, and demonstrate that even in this setup the obtained models are still more robust to pruning compared to the models fully fine-tuned either with Adam or SAM optimizers reported in Table 4.
Revised: In Table 18 we present results for runs where CrAM is not used at every optimization step, and demonstrate that even in this setup the obtained models are still more robust to pruning compared to the models fully fine-tuned either with Adam or SAM optimizers reported in Table 4.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: E T HEORETICAL S UPPORT FOR THE C R AM U PDATE
Revised: E T HEORETICAL S UPPORT FOR THE C R AM U PDATE

Original: In this section [Table] we attempt to formally derive a generic training method whose purpose is to provide compressible models, which perform well even after being compressed.
Revised: In this section, we attempt to formally derive a generic training method whose purpose is to provide compressible models, which perform well even after being compressed.

Original: To understand why we can hope to achieve such guarantees, we first take a brief detour to the area of robust optimization
Revised: To understand why we can hope to achieve such guarantees, we first take a brief detour to the area of robust optimization.

Original: [Table]

Original: E. R OBUST O PTIMIZATION
Revised: E. R OBUST O PTIMIZATION

Original: Generally, practical training methods are based on versions of stochastic gradient descent attempting to minimize a loss function L ( θ ) .
Revised: Generally, practical training methods are based on versions of stochastic gradient descent attempting to minimize a loss function L ( θ ) .

Original: However, θ might turn out to be a bad solution as the landscape of
Revised: However, θ might turn out to be a bad solution as the landscape of

Original: L in its neighborhood could contain large changes in value.
Revised: L in its neighborhood could contain large changes in value.

Original: To address this issue, one may attempt to flatten L such that it is less sensitive to sharp drops in value localized around a very small region.
Revised: To address this issue, one may attempt to flatten L such that it is less sensitive to sharp drops in value localized around a very small region.

Original: To this extent, a standard robustification can be defined by
Revised: To this extent, a standard robustification can be defined by

Original: [Equation]
Revised: [Equation]

Original: which makes the value of (cid:101) L ( θ ) take that of the largest value of L given by perturbation of θ within a ball of radius ρ .
Revised: which makes the value of (cid:101) L ( θ ) take that of the largest value of L given by perturbation of θ within a ball of radius ρ .

Original: While this robustified function may seem well suited to generic training tasks, it is a priori unclear that it is amenable to optimization.
Revised: While this robustified function may seem well suited to generic training tasks, it is a priori unclear that it is amenable to optimization.

Original: However, under certain conditions, we can efficiently optimize (cid:101) L by using a classical theorem in robust optimization due to Danskin (Danskin, 2012).
Revised: However, under certain conditions, we can efficiently optimize (cid:101) L by using a classical theorem in robust optimization due to Danskin (Danskin, 2012).

Original: Theorem 1.
Revised: Theorem 1.

Original: (Danskin) Let C ⊆
Revised: (Danskin) Let C ⊆

Original: R m be a compact set, let a function ϕ : R n × C
Revised: R m be a compact set, let a function ϕ : R n × C

Original: → R such that ϕ ( · , y ) is continuously differentiable for every fixed y ∈ C and ∇ x ϕ ( x , y ) is continuous on R n × C , and let ψ : R n → R be defined as
Revised: → R such that ϕ ( · , y ) is continuously differentiable for every fixed y ∈ C and ∇ x ϕ ( x , y ) is continuous on R n × C , and let ψ : R n → R be defined as

Original: [Equation]
Revised: [Equation]

Original: Then ψ is locally Lipschitz continuous, directionally differentiable, and its directional derivatives satisfy
Revised: Then ψ is locally Lipschitz continuous, directionally differentiable, and its directional derivatives satisfy

Original: [Equation]
Revised: [Equation]

Original: where C ∗ ( x ) is the set of maximizers
Revised: where C ∗ ( x ) is the set of maximizers

Original: [Equation]
Revised: [Equation]

Original: In particular, if for some x ∈ R n the set C ∗ ( x ) = { y ∗ x } is a singleton, then ψ is differentiable at x and
Revised: In particular, if for some x ∈ R n the set C ∗ ( x ) = { y ∗ x } is a singleton, then ψ is differentiable at x and

Original: [Equation]
Revised: [Equation]

Original: This shows that, under certain assumptions, we can obtain directional derivatives for (cid:101) L ( θ ) by simply maximizing L ( θ + δ ) over δ ∈ B 2 ( ρ ) .
Revised: This shows that, under certain assumptions, we can obtain directional derivatives for (cid:101) L ( θ ) by simply maximizing L ( θ + δ ) over δ ∈ B 2 ( ρ ) .

Original: Corollary 2.
Revised: Corollary 2.

Original: Let (cid:101)
Revised: Let (cid:101)

Original: L be defined as in Equation (10), and define
Revised: L be defined as in Equation (10), and define

Original: [Equation]
Revised: [Equation]

Original: and let δ ∈ C ∗ ( θ ) .
Revised: and let δ ∈ C ∗ ( θ ) .

Original: Provided that L ( θ ) is continuously differentiable, and θ is not an articulation point for (cid:101) L , −∇ L ( θ + δ ) is a descent direction for (cid:101) L ( θ ) as long as it is nonzero.
Revised: Provided that L ( θ ) is continuously differentiable, and θ is not an articulation point for (cid:101) L , −∇ L ( θ + δ ) is a descent direction for (cid:101) L ( θ ) as long as it is nonzero.

Original: Proof.
Revised: Proof.

Original: Let h = ∇ L ( θ + δ ) .
Revised: Let h = ∇ L ( θ + δ ) .

Original: We apply Danskin’s theorem for ϕ ( θ , δ ) =
Revised: We apply Danskin’s theorem for ϕ ( θ , δ ) =

Original: L ( θ + δ ) and C = B 2 ( ρ ) .
Revised: L ( θ + δ ) and C = B 2 ( ρ ) .

Original: This shows that
Revised: This shows that

Original: [Equation]
Revised: [Equation]

Original: Provided that θ is not an articulation point for (cid:101) L , we also have that d (cid:101) L ( θ ; − h )
Revised: Provided that θ is not an articulation point for (cid:101) L , we also have that d (cid:101) L ( θ ; − h )

Original: = − d (
Revised: = − d (

Original: cid:101) L ( θ ; h ) ≤ 0 , which concludes the proof.
Revised: cid:101) L ( θ ; h ) ≤ 0 , which concludes the proof.

Original: E.1. F ROM R OBUST O PTIMIZATION TO SAM
Revised: E.1. F ROM R OBUST O PTIMIZATION TO SAM

Original: Per Corollary 2, to obtain a descent direction it suffices to maximize L ( θ + δ ) over the set of perturbations satifying ∥ δ ∥ ≤ ρ .
Revised: Per Corollary 2, to obtain a descent direction it suffices to maximize L ( θ + δ ) over the set of perturbations satifying ∥ δ ∥ ≤ ρ .

Original: In general, even when the underlying function L is convex, this may be a difficult problem.
Revised: In general, even when the underlying function L is convex, this may be a difficult problem.

Original: Instead, one may simply attempt to obtain a good local maximizer of L in a bounded region around θ .
Revised: Instead, one may simply attempt to obtain a good local maximizer of L in a bounded region around θ .

Original: The simplest possible way to do so is by performing a step of gradient ascent , which can be regarded as a proxy for the maximization subproblem.
Revised: The simplest possible way to do so is by performing a step of gradient ascent , which can be regarded as a proxy for the maximization subproblem.

Original: Using this step, we immediately obtain the iteration:
Revised: Using this step, we immediately obtain the iteration:

Original: [Equation]
Revised: [Equation]

Revised: which recovers the extrapolated SAM gradient step from Foret et al.

Original: There is exhaustive research that has previously been done on robust optimization methods.
Revised: There is exhaustive research that has previously been done on robust optimization methods.

Original: For a comprehensive reference, we point the reader to Teo’s PhD thesis (Teo, 2007).
Revised: For a comprehensive reference, we point the reader to Teo’s PhD thesis (Teo, 2007).

Original: E. R OBUST O PTIMIZATION FOR C OMPRESSIBLE M ODELS
Revised: E. R OBUST O PTIMIZATION FOR C OMPRESSIBLE M ODELS

Original: With the robust optimization framework in mind, we are ready to attempt implementing a similar scheme which exhibits robustness to compression.
Revised: With the robust optimization framework in mind, we are ready to attempt implementing a similar scheme which exhibits robustness to compression.

Original: To motivate the method, let us consider the post-training compression.
Revised: To motivate the method, let us consider the post-training compression.

Original: After training the model to weights to θ
Revised: After training the model to weights to θ

Original: T
Revised: T

Original: we apply a one-shot compression method C over some perturbation θ T
Revised: we apply a one-shot compression method C over some perturbation θ T

Original: + δ of the weights.
Revised: + δ of the weights.

Original: This captures several iterative methods for compression, such as iterative pruning, where changes in weights are alternated with one-shot pruning methods.
Revised: This captures several iterative methods for compression, such as iterative pruning, where changes in weights are alternated with one-shot pruning methods.

Original: If our goal is to make the loss after compression robust within a small neighborhood of perturbations δ , we can establish as a formal objective to minimize the robustified loss
Revised: If our goal is to make the loss after compression robust within a small neighborhood of perturbations δ , we can establish as a formal objective to minimize the robustified loss

Original: [Equation]
Revised: [Equation]

Original:  for some magnitude ρ of allowed perturbations.
Revised: (12) for some magnitude ρ of allowed perturbations.

Original: In our case we will focus on the case where these are bounded in ℓ 2 norm, but this can be easily extended to other choices of the domain.
Revised: In our case we will focus on the case where these are bounded in ℓ 2 norm, but this can be easily extended to other choices of the domain.

Original: Just as before, we can now attempt to minimize L CrAM , or find a near-stationary point, by gradient descent.
Revised: Just as before, we can now attempt to minimize L CrAM , or find a near-stationary point, by gradient descent.

Original: Using the robust optimization framework we may attempt to optimize it using Corollary 2 after replacing
Revised: Using the robust optimization framework we may attempt to optimize it using Corollary 2 after replacing

Original: L ( · ) with L ( C ( · )) .
Revised: L ( · ) with L ( C ( · )) .

Original: Naturally, this poses some obstacles in our case.
Revised: Naturally, this poses some obstacles in our case.

Original: The main one is the fact that it is not true that
Revised: The main one is the fact that it is not true that

Original: L ( C ( θ )) will generally be continuously differentiable, so the conditions required to obtain descent directions via an inner maximization loop are not satisfied.
Revised: L ( C ( θ )) will generally be continuously differentiable, so the conditions required to obtain descent directions via an inner maximization loop are not satisfied.

Original: However, we can show that under certain conditions, continuous differentiability fails only at a set of points of measure 0 .
Revised: However, we can show that under certain conditions, continuous differentiability fails only at a set of points of measure 0 .

Original: Definition 1.
Revised: Definition 1.

Original: Let S be a countable set, let { P i } i ∈ S be a covering of R n with convex sets, and let
Revised: Let S be a countable set, let { P i } i ∈ S be a covering of R n with convex sets, and let

Original: S ( x ) denote the family of indices from S for which x ∈ P i .
Revised: S ( x ) denote the family of indices from S for which x ∈ P i .

Original: Let a family of projection operators { Π i } i ∈ S , such that for any x the projections { Π i ( x ) } i ∈ S ( x ) all map to the same point.
Revised: Let a family of projection operators { Π i } i ∈ S , such that for any x the projections { Π i ( x ) } i ∈ S ( x ) all map to the same point.

Original: We call a projective compression operator with respect to { Π i } i ∈ S a mapping C : R n → R n such that
Revised: We call a projective compression operator with respect to { Π i } i ∈ S a mapping C : R n → R n such that

Original: [Equation]
Revised: [Equation]

Original: for any i ∈ S ( x ) .

Original: For example, in the case of the Top-k compression operator, we can define a projection for each subset
Revised: For example, in the case of the Top-k compression operator, we can define a projection for each subset

Original: A of coordinates of cardinality k .
Revised: A of coordinates of cardinality k .

Original: We say that given a vector x , the set A ∈ S ( x ) iff the largest k coordinates of x in absolute value (with ties broken lexicographically) are supported in A .
Revised: We say that given a vector x , the set A ∈ S ( x ) iff the largest k coordinates of x in absolute value (with ties broken lexicographically) are supported in A .

Original: Lemma 3 (Continuously differentiable functions induce few singularities after compression) .
Revised: Lemma 3 (Continuously differentiable functions induce few singularities after compression) .

Original: Let
Revised: Let

Original: L : R n → R be a continuously differentiable function, and let C be a projective compression operator.
Revised: L : R n → R be a continuously differentiable function, and let C be a projective compression operator.

Original: Then the function g ( x ) :
Revised: Then the function g ( x ) :

Original: = L ( C ( x )) is continuously differentiable everywhere except at a set of points of measure 0 .
Revised: = L ( C ( x )) is continuously differentiable everywhere except at a set of points of measure 0 .

Original: Furthermore, so is the robustified function L CrAM ( x ) :=
Revised: Furthermore, so is the robustified function L CrAM ( x ) :=

Original: max ∥ δ ∥≤ ρ L ( C ( x + δ )) .
Revised: max ∥ δ ∥≤ ρ L ( C ( x + δ )) .

Original: Proof.
Revised: Proof.

Original: First we note that the boundary of any convex set has measure zero by standard arguments in convex analysis (Lang, 1986).
Revised: First we note that the boundary of any convex set has measure zero by standard arguments in convex analysis (Lang, 1986).

Original: Since a countable union of sets of measure zero has measure zero, it follows that the union of the boundaries of P
Revised: Since a countable union of sets of measure zero has measure zero, it follows that the union of the boundaries of P

Original: i ’s has measure zero.
Revised: i ’s has measure zero.

Original: Now since L is continuously differentiable, within any set P i , we have that g ( x ) =
Revised: Now since L is continuously differentiable, within any set P i , we have that g ( x ) =

Original: L (Π i ( x )) , and hence it remains continuously differentiable.
Revised: L (Π i ( x )) , and hence it remains continuously differentiable.

Original: Therefore the only region for which we can not argue about continuous differentiability is the complement of the union of interiors of P
Revised: Therefore the only region for which we can not argue about continuous differentiability is the complement of the union of interiors of P

Original: i ’s, ( ∪ i int P i ) c ⊆
Revised: i ’s, ( ∪ i int P i ) c ⊆

Original: ∪ i ∂P i which is a set of measure zero.
Revised: ∪ i ∂P i which is a set of measure zero.

Original: Since g is well-behaved almost everywhere, all that remains to argue is that this is the same case with L CrAM .
Revised: Since g is well-behaved almost everywhere, all that remains to argue is that this is the same case with L CrAM .

Original: For any fixed direction ∆ θ , we define the mapping
Revised: For any fixed direction ∆ θ , we define the mapping

Original: [Equation]
Revised: [Equation]

Original: and its robustification
Revised: and its robustification

Original: [Equation]
Revised: [Equation]

Original: Hence the directional derivative w.r.t.
Revised: Hence the directional derivative w.r.t.

Original: ∆ θ of L CrAM ( θ ) is discontinuous only when (cid:102) M ( θ ) is discontinuous.
Revised: ∆ θ of L CrAM ( θ ) is discontinuous only when (cid:102) M ( θ ) is discontinuous.

Original: Finally, we note that this almost never happens, as M is continuous almost everywhere, and thus so must be (cid:102) M .
Revised: Finally, we note that this almost never happens, as M is continuous almost everywhere, and thus so must be (cid:102) M .

Original: Thus, all directional derivatives are continuous except at a set of measure 0 , which concludes the proof.
Revised: Thus, all directional derivatives are continuous except at a set of measure 0 , which concludes the proof.

Original: Finally, just like in the previous case, maximizing L ( C ( θ + δ )) over small perturbations is generally intractable.
Revised: Finally, just like in the previous case, maximizing L ( C ( θ + δ )) over small perturbations is generally intractable.

Original: So we instead consider obtaining a good enough maximizer via a standard iterative method which has shown good performance in practice.
Revised: So we instead consider obtaining a good enough maximizer via a standard iterative method which has shown good performance in practice.

Original: More precisely we consider the projected gradient ascent method, which provides strong theoretical guarantees, even when the projection is performed onto non-convex domains (Peste et al., 2021).
Revised: More precisely we consider the projected gradient ascent method, which provides strong theoretical guarantees, even when the projection is performed onto non-convex domains (Peste et al., 2021).

Original: In the case where the compression operator represents magnitude pruning, this corresponds to the iterative hard thresholding (IHT) method, frequently employed in the sparse recovery literature.
Revised: In the case where the compression operator represents magnitude pruning, this corresponds to the iterative hard thresholding (IHT) method, frequently employed in the sparse recovery literature.

Original: To reach a good iterate within this specific domain we instead perform a single step of (projected) gradient ascent, which matches the IHT iteration:
Revised: To reach a good iterate within this specific domain we instead perform a single step of (projected) gradient ascent, which matches the IHT iteration:

Original: [Equation]
Revised: [Equation]

Original: We have therefore obtained a re-derivation of the CrAM update in Equation 2.
Revised: We have therefore obtained a re-derivation of the CrAM update in Equation 2.