Preserving Deep Representations in One-Shot Pruning: A Hessian-Free Second-Order Optimization Framework
Authors: Ryan Lucas
Rahul Mazumder
Venue: ICLR 2025 Poster (2025)
Keywords: Neural Network Pruning
Structured Pruning
Optimization
Hessian-free Optimization
Primary area: optimization
TL;DR: An improved one-shot pruning framework using Hessian-free optimization

Abstract: We present SNOWS, a one-shot post-training pruning framework aimed at reducing the cost of vision network inference without retraining. Current leading one-shot pruning methods minimize layer-wise least squares reconstruction error which does not take into account deeper network representations. We propose to optimize a more global reconstruction objective. This objective accounts for nonlinear activations deep in the network to obtain a better proxy for the network loss. This nonlinear objective leads to a more challenging optimization problem---we demonstrate it can be solved efficiently using a specialized second-order optimization framework. A key innovation of our framework is the use of Hessian-free optimization to compute exact Newton descent steps without needing to compute or store the full Hessian matrix. A distinct advantage of SNOWS is that it can be readily applied on top of any sparse mask derived from prior methods, readjusting their weights to preserve deep feature representations. SNOWS obtains state-of-the-art results on various one-shot pruning benchmarks including residual networks and Vision Transformers (ViT/B-16 and ViT/L-16, 86m and 304m parameters respectively). Our open-source implementation is available at https://github.com/mazumder-lab/SNOWS.

Published as a conference paper at ICLR 2025
PRESERVING DEEP REPRESENTATIONS IN ONE-SHOT PRUNING: A FREEE SECOND-ORDER OPTI- MIZATION FRAMEWORK
Ryan Lucas" Rahul Mazumder'
ABSTRACT
We present SNOWS, a one-shot post-training pruning framework aimed at re- ducing the cost of vision network inference without retraining. Current leading one-shot pruning methods minimize layer-wise least squares reconstruction error which does not take into account deeper network representations. We propose optimize a more global reconstruction objective. This objective accounts for non- linear activations deep in the network to obtain a better proxy for the network loss. This nonlinear objective leads to a more challenging optimization problem-wa demonstrate can be solved efficiently using specialized second-order optimiza- tion framework. A key innovation of our framework is the use of Hessian-free optimization to compute exact Newton descent steps without needing com- pute or store the full Hessian matrix. A distinct advantage of SNOWS is that it can be readily applied on top of any sparse mask derived from prior methods, readjusting their weights to preserve deep feature representations. SNOWS ob- tains state-of-the-art results on various one-shot pruning benchmarks including residual networks and Vision Transformers (ViT/B-16 16 and ViT/L--16, 86m and 304m parameters respectively). Our open-source implementation is available at https:: /ttt///gttttu com/ www.om/maaammdeeee -aaeeerraaab///aaa SNOWS.
INTRODUCTION
Modern deep learning- bbaeed vision models, particularly convolutional neural networks (CNNs) (Le- cun et al., 1998) and more recent vision transformers (ViTs) (Dosovitskiy et al., 2021). 1). have achieved remarkable performance in various computer vision tasks, including image classification, object de- tection, and visual tracking (Chai et al., 2021). However, their success comes at the cost of substan- tial computational and memory resources, and these models are only trending towards increasingly large scales (Dehghani et al., 2023). Neural network pruning, class of techniques aimed at re- ducing the number of parameters by selectively removing weights or connections, has emerged as solution to reduce the inference demands of large-scale vision networks (Goel et al., 2020; He & Xiao, 2023; Kuznedelev et al., 2023), alongside methods such as quantization (Cheng et al., 2024).
Many pruning approaches require expensive retraining phase following weight removal to recover network performance (Liu et al., 2019; Frankle & Carbin, 2019). Especially as networks scale, practitioners looking to prune and cheaply deploy model often lack the resources fully retrain it and may not even have access to the model's 's dataset or optimizer. While classical one-shot pruning approaches (where weights are removed in single step rather than iteratively) typically included retraining Hagiwara (1994), recent work has explored one-shot pruning without this computationally expensive step (Frantar & Alistarh, 2023; 2022; Singh & Alistarh, 2020).
One-shot pruning approaches typically use second-order approximations of the neural network loss function their pruning criteria objective function, an approach dating back to (LeCun et al., 1989; Hassibi & Stork, 1992). Modern second-order approaches fall into two main categories: global methods, which consider the entire network at once, and local methods, which prune one layer at a time. Global methods decide which weights to prune and how to update the remaining
MITT Operations Research Center (ryan fryannu@mii u@mi tommit.ee eduu)
1 MIT Sloan School of Management, MIT Operations Research Center, and MIT Center for Statistics (rahul maz@mi t. edu)

---

Published as a conference paper at ICLR 2025
weights, based second-order approximation of the downstream task loss function. However, since computing and inverting the full Hessian matrix is expensive, these methods often replace the Hessian with approximation such as the Fisher information matrix further approximations thereof (Singh & Alistarh, 2020; Frantar et al., ; Benbaki et al., 2023; Wang et al., 2019). Like retraining approaches, global pruning methods can be difficult to scale to the largest vision networks due to their need to take full network gradients (e.g. in the Fisher computation). Local methods on the other hand, which have become popular for large networks, replace the task loss with local layer-wise least squares reconstruction loss (Dong et al., 2017a; Frantar & Alistarh, 2022: Meng et al., 2024a), for which the Hessian can be computed in a simple, straightforward fashion. In these local layer-wise reconstruction methods, the only optimization variables are the weights at a single layer, hence layer-wise Hessians can be isolated and the network pruned one layer at time. These local methods are computationally efficient, but may not provide the most accurate proxy of the loss of the network and neglect the influence of deeper feature representations.
In this paper, we present SNOWS neural network pruning framework that bridges the gap between these global and local one-shot pruning approaches. Unlike global or retraining-base methods, SNOWS only requires calculating gradients for individual layers. At the same time, the SNOWS formulation goes beyond standard layer-wise reconstruction and optimizes a more global reconstruction objective (cf. Equation 3). The motivation is to preserve not just local set of fea- tures, but to consider the effect of pruning on the more global set of learned representations. While this objective leads to more challenging nonlinear optimization problem that cannot be solved by standard layer-wise methods, we show that it can solved efficiently using specialized second- order optimization framework. Unlike prior global approaches that make approximations to the Hessian, SNOWS optimizes an exact second-order approximation. A key driver of the scalability of our method is the use of Hessian-free optimization (Martens, 2010), which allows us to efficiently compute and optimize the second-order approximation without explicitly forming storing the full Hessian matrix. This enables SNOWS to scale to the largest CNNs and ViTs. For instance, using method the ViT-L/16 model (== 304m parameters) can be pruned on single A100 GPU.
Semi-structured pruning approaches have received more attention since they lead to more tangible acceleration on modern computing hardware. We primarily target N:M sparsity (Mishra et al., 2021), a semi-structured pruning approach where N out of every M (contiguous, aligned) weights are retained. This sparsity pattern leads to greater acceleration than unstructured sparsity by aligning with hardware indexing patterns, such as those in NVIDIA's Sparse Tensor Cores (Pool & Yu, 2021), and leads lesser accuracy loss than structured formats (Zhang et al., 2023). Unlike previous methods that require training N:M sparse networks from scratch (Zhou et al., 2021; Zhang et al., 2022) to achieve performance competitive with dense networks, SNOWS can prune N:M sparsity in one-shot using just few thousand samples.
Contributions. We make the following contributions to the one-shot pruning problem:
We introduce novel post-training pruning framework that operates at the layer-wise level to optimize a more global reconstruction objective to account for nonlinear activations deep in the network. We demonstrate empirically that this approach provides a better proxy for the network loss compared to layer-wise least squares reconstruction methods.
We develop efficient second-order optimization framework to solve the challenging nonlinear optimization problem posed by our reconstruction objective. We use Hessian-frea optimization integrated with a customized Conjugate Gradient (CG) method to efficiently solve the second- order approximation. Our CG method exploits sparsity in the linear system defined by the second- order approximation to compute exact Newton descent steps in memory- way, scaling to problems with layer-wise Hessians as large as 4.2m x 4.2m.
SNOWS delivers offte-of-thh-ar performance in one-shot pruning for computer vision, achieving superior results on various network architectures and datasets. Our method can prune ResNet50 to 1: :4 sparsity with a 1.5% drop in test accuracy on CIFAR-100 and a 9.4% drop on ImageNet- Ik using just 3,000 training examples, where prior state-of-the-art one-shot N : M pruning methods have a drop of 7.6% and 29.42% respectively. Similarly, for unstructured sparsity, our method can prune MobileNet to 70% sparsity with a 5.6% drop in accuracy, whereas the best available
'Stochsttic Newton Optimal Weight Surgeon (SNOWS)
2

---

Published as a conference paper at ICLR 2025
competing method has 12.55% drop. In addition to improving performance on CNNs, our method scales to modern ViT networks with up to 304m parameters.
2 RELATED WORK
One-shot pruning without retraining. Recent work has demonstrated the possibility of pruning networks in single step without requiring expensive retraining Frantar & Alistarh (2023; 2022); Singh & Alistarh (2020). In this setting. one typically assumes access to a limited sample of cali- bration data (e.g. few thousand examples) and the goal is to compress the network with minimal aceuracy loss relative to the dense model. This modern approach to one-shot pruning builds on clas- sical work that developed second-order approximations of the task loss function. Early approaches (Lecun et al., 1998; Hassibi & Stork, 1992) considered the entire network's weights but were com- putationally intractable for large models. More recent methods like WoodFisher Singh & Alistarh (2020) and Combinatorial Brain Surgeon (Singh & Alistarh, 2020), as well as CHITA (Benbaki al., 2023), improve scalability by replacing the Hessian with the Fisher information matrix or using block-diagonal approximations. Other works use Kronecker product approximations of the Fisher matrix (Wang et al., 2019).
These global approaches optimize a more accurate proxy for the network loss function, but often struggle to scale the largest vision networks, hence layer-wise approaches which focus on recon- struction have become a popular alternative. Layer-wise OBS (Dong et al., 2017a) formulates the problem of reconstructing the linear activations of the dense network as a linear regression problem which has a standard least squares Hessian, and show that the overall network reconstruction error bounded by a quantity depending on the error given each of the layer-wise problems. The OBC framework (Frantar & Alistarh, 2022) introduces an efficient method for greedy backward elimi- nation on this layer-wise problem using fast low rank updates of the Hessian inverse. There have been limited attempts to incorporate nonlinearity into layer-wise approaches, but several papers have reported the need to consider broader notions of network connectivity as pruning criteria and loss functions (Hoang al., 2023; He & Zhou, 2024). The Net-Trim framework solves a layer-wist regression problem with ReLU activation (Aghasi et al., 2017), accounting for the non-linearity by reformulating the ReLU using linear inequalities and solving the resulting convex program. In the Net-Trim formulation, the weights are optimized in a way that is "aware" of the subsequent non- linearity by incorporating the effects of the ReLU activation into the pruning loss function. Our objective is similar in spirit, but we make no assumptions about the exact form of the non- non-lincarity, only requiring that the underlying functions are wiice-differentiable. Moreover, our formulation goes beyond the activations at single layer.
Structured and semi-structured pruning in computer vision. Unstructured pruning, which in- volves pruning individual weights in an unrestricted format, leads to high compression ratios but often cannot realize tangible speed-ups except at very high sparsity levels or when run on special- ized hardware (Cheng et al., 2024). Hence there has been movement towards structured pruning methods (He & Xiao, 2023), which lead to actual speed-ups by imposing regularity on the sparsity pattern. In CNNs, structured pruning cific channels (Meng et al., 2024c). Recently. NVIDIA (Mishra et al., 2021) released N:M Sparse Tensor Cores, the current version of which includes a 2:4 sparsity pattern that leads to double the throughput on dense matrix multiplications. Hence there has been several works applying N:M pruning in CNNs (Pool & Yu, 2021; Zhou et al., 2021). There has been more limited work on com- pressing ViT networks, in part due to their scale. Most of the recent work focuses on unstructured sparsity (He & Zhou, 2024). The CAP framework (Kuznedelev al., 2023) obtains strong results in one-shot unstructured pruning in DeiT networks (Touvron et al., 2021), though they use 4096 gradient evaluations (R= 520k samples) in their one-shot pruning experiments, and use fine-tuning to recover network accuracy. Previous work on structured and semi-structured VIT pruning has taken various approaches. Several works (Zhu et al., 2021; Yu & Wu, 2021) focused on pruning redundant channel dimensions (rows) in the VIT. (Yu et al., 2022a) proposed pruning both the width (number of neurons) and depth (number of layers) of ViTs, which centers on architectural changes rather than specifically weight pruning. Notably, almost all existing work ViT pruning requires fine-tuning to recover dense model performance, and hence the literature on one-shot structured pruning in VITs is very limited.
involve removing entire filters (Zhou et al., 2016) or spe-
3

---

Published as a conference paper at ICLR 2025
3 PROBLEM DEFINITION AND BACKGROUND
Notation. In a neural network with L layers, the network function is represented as f(x", W), where W (W1, w W is the collection of weight matrices across layers, and X Raxdo is the input, with number of data points n and input dimension dg- For & 2 1, we use X E R=xde-1 1 to refer to the inputs layer E and Y R to refer to the outputs of layer E; therefore, the input to layer e + is the output of layer E, or X = =Y', The weight matrix for the l-th layer is denoted W E Rdd-1xdd The output of layer is computed as Y f W where the function f" denotes the operations at layer E, which typically consist of an linear transfor- mation followed by a point-wiss nonlinear activation, i... f W activation function applied element- wise, though l can also include much more general operations such as residual connections, pooling operations, etc. For a wwice-differentiable loss function de- pending on the network parameters W or a subset, we denote the gradient as VL(W) 0L/OW and the Hessian H = 8 C/OW².
o (X where a an
Background on Global Model. Global methods, originating with the Optimal Brain Surgeon (OBS) framework (Lecun et al., 1998; Hassibi neural network compression:
Stork, 1992), minimize the following loss function for
min L(W) x E(X, W)) Wew
(1)
where Y are the ground truth labels, f(x,, W) represents the final output of the compressed net- work, E-) is the task loss, e.g., cross-entropy for classification, and W is predefined set of sparse weight matrices. The canonical example of W is the set of S-sparse matrices, W = {W] [Wllo S S}. though W can also represent much more general sets such as structured-spare, matrices. The very natural idea here is to compress the network to preserve performance on the downstream task. To solve the problem, OBS uses second-order approximation to Eqn (1). Intuitively, second- order method is well-motivated by the importance of interactions in pruning. Time has revealed however that there are practical issues limiting the applicability of the OBS framework. Most impor- tantly, the OBS solution to Eqn (1) requires computing the Hessian of the entire network's 's weights. For modern deep neural networks, computing the full Hessian matrix is simply not tractable.
Background on Layer-wise Model. A more practical alternative to global methods is the popular layer-wise proxy, replacing Eqn (1) with the following regression problem:
min W'EW
X' - X W
The goal in Eqn (2) is to preserve the linear activations of the model at a particular layer while enforcing sparsity in We, This formulation, which was introduced in (Dong et al., 2017a), has practical appeal since the only decision variables are the weights at single layer W so the Hessian has reduced dimensionality (de-1 x d compared to that of the full network. Moreover, since the operations on W' are linear, the objective has a standard least squares form, and its Hessian is given by H 2 E meaning it can be computed in parallel across the row dimension of X
THE SNOWS METHOD
Our Improved Layer-wise Model. While more tractable than global approaches, the layer- wise formulation in Eqn (2) focuses solely on preserving the linear activations at a particular layer without considering the impact on subsequent layers' representations. In this paper we instead consider the following layer-wise model:
min L(W") W'EW
Σ
yltk- plitt W9
(3)
where (.) is defined as the composition of the next k operations starting from layer f: pletk (x', w - pt+((, w o peth-1(. (., we o f'(,, w²)
We omit the details of the second-order approximation for brevity, but provide more in A.1.1
4

---

Published as a conference paper at ICLR 2025
and Y = - W are the dense target outputs after the k-th operation starting from layer f. Hence, the loss function considers the effect of W on all layers and operations up to E+ K. A natural motivation for Eqn (3) is that pruning weights at a particular layer affects not only the outputs of that layer but also activations in subsequent layers. Namely, the solution to the layer-wise problem in Eqn (2) does not necessarily minimize the discrepancy in the outputs of deeper layers:
argmin We x x'
# argmin Wr yeth f (x', |
for all A: >> O, where the constituent functions can include non-linear activations, pooling operations, and residual connections, among others. We provide more detail on the motivation for the SNOWS loss function in subsubsection A.1.2.
Special Cases of Eqn (3). The loss function is a flexible proxy for the neural network loss function, where the parameter K can be varied depending on available computational resources. When K O, the loss function reduces to the layer-wiss wise framework. When K = 1, the formulation is analogous to Net-Trim for ReLU networks Aghasi et al. (2017), but also generalizes it other activation functions. Setting K- = Kmxx which denotes the number of operations after layer & up to the output layer, extends the optimization to all operations following layer f. In Figure 1, we observe how the value of K influences network performance when pruning the full ResNer20 and ResNet50 networks to 1:4 sparsity using SNOWS. In Table in subsubsection A.2.4, we provide an ablation on the run-time and peak memory usage when varying K.
ResNet20 (CIFAR-10)
ResNet50 (CIFAR-100)
90
Aocurss (88)
80
70
65
K=0
K=10
60
K=30
K=1
K=220
K=3
K=40
55
K-60
60
K=5
K=40
50
K=10
13
1.0
Pruning treation (Block)
Feration (Block)
Figure 1: Effect of varying K in the loss function in Eqn (3) on out-of-sample accuracy pruning ResNet20 on CIFAR-I0 and ResNet50 on CIFAR-100 to 1:4 sparsity (74% and 66% respectively). Increasing K improves the accuracy of the pruned network, at the cost of higher computation time.
Cheap layer-wise gradients without computing the full computational graph. The gradient of Eqn (3) function with respect to the pruned weights W can be expressed as:
VL(W> Σ 2 (x
fott
aftitth (x* (')
(4)
oW
An important advantage of this formulation is that allows for reduced memory requirements com- pared to computing full network gradients. Specifically, the gradient has dimensionality dim W*), which significantly smaller than the total number of parameters in the full network. Moreover, the computational graph required to compute this gradient is shorter because it includes only the operations from layer up to É K, rather than the entire network, which minimizes the number of activations that need to be constructed and stored on the computational graph.
Each function féiE+k can represent lin- ear transformations and several non-linear operations, so while this is still a layer- wise problem depending only on W it no longer a linear least squares prob- lem. Hence, the Hessian H is no longer equal to 2 in general. To form the Hessian we need to compute all second-order partial derivates
Dense Sub-Network
x
r
y
y
and does not separate
We
fe+1
fe+k
Sparse Sub-Network
ViPC(W9
Figure 2: Computional graph to minimize Eqn (3).

---

Published as a conference paper at ICLR 2025
of flif* for each k O...., K which is also not tractable for large networks. Nonetheless, now present a tractable framework to optimize Eqn (3).
An Efficient Second-Order Optimization Framework to Optimize Eqn (3) Our K- K-step ob- jective induces a discrete, nonlinear optimization problem which cannot necessarily be solved in a similar fashion to previous layer-wise methods (Dong al., 2017a; Aghasi et al., 2017; Frantar & Alistarh, 2022). Thus, we adopt a similar framework to more global methods by minimizing a local quadratic approximation of our nonlinear objective. However, given the difficulty of forming the Hessian directly, we use Hessian-free optimization (Martens, 2010) to optimize the local quadratic approximation, combined with CG method that exploits sparsity for efficient updates. We now describe the key components of our framework.
Hessian-free Optimization. Given the difficulty of forming the Hessian directly, we opt use Hessian-free free optimization (Martens, 2010) to optimize the local quadratic approximation. Hessian- free optimization is second-order optimization method that avoids explicit computation of the Hessian matrix by exploiting the fact that while the Hessian matrix H is expensive to compute, the so-called Hessian product H 8 can be computed as:
VL(W' +c) -VCWW)
6 lim c00
(5)
which has cost of additional gradient evaluation when computed via finite differences. Note also that dim H 5) = dim(W* meaning that the memory requirements of H ô are the same for W unlike the Hessian which is usually far large fit on GPU memory. Moreover, since we operate on smaller localized sections of the network, from layer to operation + K, our formulation is less susceptible numerical error that can harm Hessian-free optimization full network training approaches (Martens, 2010). This allows us to compute exact Hessian products as opposed to Fisher or Gauss-Newton approximations used in prior work on matrix-free training and pruning (Martens & Sutskever, 2012; Singh & Alistarh, 2020; Frantar et al., 2021).
Decoupling the Discrete and Continuous Problems. In designing our algorithm for scalability to the largest networks, we develop approach that is agnostic to the mask. Our approach can be integrated with other techniques for mask selection, or can stand alone with simple methods such as magnitude pruming. Our alternative approach for optimizing the local quadratic approximation can be devised by recognizing Eqn (3) as a bi-level optimization problem:
min ZEZ W. E xd,
min
L(WE)
y pest (x',
W2,WZ O 2 2
where Z E Z is a binary mask imposing the same sparsity pattern as W and W denotes the weights when the sparsity pattern is fixed. We omit the superscript l from Z for simplicity, though the mask is typically chosen and applied layer-wise also. In our case, fixing Z allows us to perform second-order Taylor expansion around the masked solution for the weights:
1 8222202
L W + +5z) = L(W VL(WE) 8z
2
In this approximation, we consider perturbing the weights W for given mask Z, which induces sparsity in the solution. We can now design strategies to minimize Eqn (7) which simplifies the problem dramatically compared solving coupled discrete and continuous problem.
Sparsity in the Second-Order Expansion. The Taylor expansion in Eqn (7) has several key prop- erties related to the sparsity induced by the mask Z. First, consider the gradient VL(W)) aw?"
ac
It follows that (VL); 0 wherever Z, = 0. Similarly, the Hessian Hz exhibits sparsity, with (H2),, = 0 wherever Z; = 0 or Zj = 0. Therefore, the sparsity pattern of Hz is the square of the sparsity pattern of Z. Specifically, if Z has a sparsity level s E [0, 1], then the Hessian Hz will have a sparsity level s². Moreover, any row or column (Hz) in the Hessian corresponding to Z, = 0 will be entirely zero. We exploit this sparsity in Hz to reduce the dimensionality of the system required to minimize Eqn (7). is important to note that unlike in the global approximation, we cannot assume VL(W>) 2 = 0 in Eqn (7). In the global case this is justified since they start by computing the gradient and Hessian of the fully dense network and each time make optimal weight

---

Published as a conference paper at ICLR 2025
adjustments. Whereas in our approximation, this would amount making the assumption that the current weights are optimal for any mask Z, even after other weights have been deactivated. Thus the gradient term remains and minimizing Eqn (7) with respect to the weight direction óz gives:
min 82
62 V 2 + 2
(8)
Setting the gradient of the above quadratic equal to zero yields the following reduced linear system:
Hz 1Iz-zz= = -V==V((W
The matrix H E Rm is the Hessian restricted to the active weights, where m represents the number of weights with Z; F O. Correspondingly, the gradient VL(W)) R includes only the components associated with these active weights. As a result, the weight update 6z E I can be calculated exclusively for the active weights dz H HZ' which yields the celebrated Newton update (Battiti, 1992), but on a greatly reduced linear system.
Directly forming and inverting the Hessian Hz, even if sparse, is still not tenable. Here we exploit the fact that the left-hand side Hessian product in Eqn (9) can be computed via Eqn (5), and hence the corresponding linear system can be solved by querying the Hessian product within the CG method (Nocedal & Wright, 2006). To ensure that the Hessian is positive definite, particularly when the objective function is not strongly convex, add Levenberg- Levenberg-Maqquard regu- larization to the Newton update. This is a small dampening regularization to the Hessian diagonal, resulting in the following system:
CG Method.
+ 82 -VC
(10)
This leads to Algorithm 1 to solve for the Newton update, which relies solely on products with the Hessian. Since each call to Equ (5) requires computing the gradient, the number of iterations in CG is a critical factor in determining the run-time of the overall algorithm. However, although CG can require up to m iterations to fully converge, it often makes significant progress in far fewer steps. In subsubsection A.2.3, we provide an ablation study where we vary the maximum number of steps. CG can typically be safely terminated early after a reasonable number of iterations without significantly degrading the quality of the solution. Given the method for computing the Newton di- rection, the iterative scheme for performing Newton updates then involves minimizing the quadratic approximation of the loss function at each step to adjust the current solution:
dz RJ argmin oz
8 VL
(W))
+ 2 H022222
(via CG)
W W/ + adz
(Newton Update)
where a is chosen to satisfy the Armijo sufficient decrease condition (Byrd et al., 2016). We provide the full condition in subsubsection 1.1.3. In practice, due hardware memory constraints, run a batched version of Newton descent called sub-sampled Newton descent. Instead of computing full gradients or Hessian products at once, we sample mini-batches of data to compute stochastic estimates of the gradient and Hessian products. Stochastic Newton descent is a common technique in the literature on second-order optimization methods and has been shown to exhibit strong conver- gence guarantees (Roosta-Kee Khorasani & Mahoney, 2016). The only alteration we make compared to traditional stochastic Newton algorithms is that we sample Hessian products rather than the Hessian directly. The full algorithm is given in Algorithm 2. We remark that, in practice, the stochastic Newton algorithm typically converges in very few mini-batches. In subsubsection A.2.1 and sub- subsection A.2.2, provide ablations where we compare SNOWS to Stochastic Gradient Descent (SGD) and the Fisher approximation to optimize Eqn (3), finding in both cases that SNOWS obtains better minima and converges faster.
7

---

Published as a conference paper at ICLR 2025
Algorithm 1 CG Algorithm to Compute åz
A complete network pruning algorithm. So far our descriptions have focused on the method for
1: Input: Fixed mask Z, sparse weights W
the layer-wise pruning problem. We now present
2: Compute gradient f- 2
a recap of the full layer-wise pruning process (AI- gorithm 2) as well the algorithm to prune the full network (Algorithm 3). Our approach is flex- ible and can accommodate any choice of masks .....
3: Initialize dz (- 0, r <- 9. P 4- r
4: while ||| 2 tol do
5: Compute Bp)) = Hzp + Ap via Eqn (5)
{z'..... Z Algorithm 3 uses a cascading mechanism for pruning layer by layer, where the next layer's input is adjusted after pruning the cur- rent layer. This follows the same core idea as Net- Trim (Aghasi et al., 2017). Unlike parallel ap- proach, where each layer is pruned independently. this method passes the output of one pruned layer as the input to the next. Specifically, after prun- ing layer l the input to the next layer is updated
6:
"
-'r p B(p)
7:
Update 8z x 8z + up
8:
Nnew t Bp)
9:
7
+ r
10:
p Tnew + Yp
11:
r 5- Fnew
12: end while
13: return öz
as X f'(X,, W This ensures that each layer's s optimization accounts for changes introduced via pruning in the previous layer.
Algorithm 2 SNOWS: Layer-wise Optimization Algorithm 3 SNOWS: Full Network Algorithm
1: Input: Layer input Xt, Dense weights W Sparse mask z'
Input: Network input X dense weights {W* ..... _ sparse masks ,z,,..,,Z! Z²),
2: Initialize:
layers L, horizon parameter K
W' w' z
for l = L,...,oo do
3: for b 1, B do (I> mini-batches)
Find W as:
4: Compute ( via Algorithm 1:
z'
argmin Σ yetk p W + 23/13 W
K(E)
(()) argmin
(
1 8
H (6) 2 2:0
8
via Algorithm 2 with ((e) min (, Kmax). Cascade inputs to the next layer:
5: Update weights on the active set:
Wz W + (00)
x-1=f W
6: end for
end for
7: Output: Pruned weights W
Output: Pruned weights { (WI. ..... W43
EXPERIMENTS ON CNNS AND VISION TRANSFORMERS
CNN Pruning. CNNs are a fundamental architecture for computer vision tasks. In a CNN, the weight tensor W has dimensions (dom x dm x k5 x Rw.). Here, deut is the number of output channels (filters), din is the number of input channels, and kn and kn are the kernel height and width, respectively. We focus on a semi-structured sparsity pattern that has become more popular in recent years, called N:M sparsity. This sparsity pattern requires that, for every (non-overlapping) block of M weights, exactly N weights are non-zero. For CNNs, the N:M sparsity format requires that the weights be sparse along the input dimension dm (Pool & Yu, 2021). We follow this convention throughout our experiments. We provide an example of an N:M sparse convolutional kernel in Figure 11 in subsubsection A.3.1 and additional details on how we handled the K-step problem.
Experimental setup. We evaluate SNOWS across several standard CNN benchmarks. We use the 100 CIFAR-100, and ImageNet-1k kk datasets (Krizhevsky & Hinton, 2009; Deng et al., 2009), where we evaluate performance pruning the ResNet20 and ResNet50 architectures (He et al., 2015) to N:M sparsity. Additionally, for experiments on unstructured sparsity, we use the MobileNetV1 architecture (Howard et al., 2017) to compare to previous studies on unstructured pruning. Due to the size of the original ImageNet dataset and the fact that one-shot methods require just few thousand samples, we use the Minilmagenet-I 1k dataset, a sample from the ImageNet- Ik dataset. We use 10,000 samples from this dataset for evaluation, with 3,000 calibration samples being used
8

---

Published as a conference paper at ICLR 2025
for the N:M sparsity experiments. For unstructured sparsity, use 1,000 calibration samples to be consistent with the methods we compare with (Benbaki et al., 2023).
N:M sparsity. We first evaluate the impact of rumning SNOWS on top of existing mask selection techniques, OBC (Frantar & Alistarh, 2022), SparseGPT (Frantar & Alistarh, 2023), and MP, that support N:M sparsity. Adding SNOWS on top of other pruning techniques consistently improves performance across different methods and datasets, as shown in Figure 1. This is evident in more challenging datasets like CIFAR-I00 and ImageNet-Ik; where other methods alone perform poorly, but the addition of SNOWS recovers close to the performance of the dense model. For 2:4 sparsity, the choice of mask selection technique has less of an impact on the results. Once SNOWS is applied, the methods achieve similar performance. However, at higher sparsity (I:4), more intensive mask selection techniques, particularly OBC, combined with SNOWS give the best results. We provide latency performance estimates for the N:M patterns shown in Figure 1 in subsubsection A.3.5. We also report hyperparameters for SNOWS, including the value of K, the dampening factor X, and the CG tolerance in subsubsection A.3.6.
Results.
Unstructured Sparsity. We also compare SNOWS against several popular one-shot pruning meth- ods for unstructured sparsity, including Magnitude Pruning (MP) Gupta et al. (2024), WoodFisher (Singh & Alistarh, 2020), Combinatorial Brain Surgeon (Yu et al., 2022b), and CHITA (Benbaki et al., 2023), using the reported results from the CHITA paper for the competing methods. We remark that our base ResNet20 and MobileNetV1 models start from slightly different starting accu- racies, thus we instead show the accuracy change compared to the baseline model in Figure 3 and provide the raw numbers in subsubsection A.3.3. Other one-shot methods optimize both the mask and weights, whereas SNOWS optimizes from a mask obtained through magnitude pruning. Despite this, SNOWS consistently outperforms existing one-shot approaches at all sparsity levels. In sub- subsection A.3.4, we provide additional retraining experiments for SNOWS and the best competing method in Figure 3, FALCON Meng et al. (2024b), showing SNOWS can further improve with some additional fine-tuning. We also provide additional run-time comparisons subsubsection A.3.7.
Restiet 20 (FFAA-16))
CIFAR-10
CIFAR-100 ImageNet-Ik 1k (ResNet50)
accuesT Chunne PH)
Method
(ResNet20)
(ResNet50)
-20
1:4
2:4
1:4
2:4
1:4
2:4
-40
-60
Dense Model
92.58
77.40
75.71
-50
OBC
76.77
90.98
69.79
77.03
46.29
73.74
0.5
+SNOWS
89.09 91.86
76.05
77.11
66.33
73.87
Mob leketV]] [imagetes:1100
SparseGPT
11.06
86.05 8.27
74.10
1.61
69.46
-20
+SNOWS
86.10 91.94
73.66
77.10
49.32
73.59
MP
10.99
63.45 9.45 70.76
63.45
+SNOWS
87.82 91.88 75.87 77.24 62.75
73.97
-20
0/6
6.7
Table 1: Top-1 Test Accuracy integrating SNOWS with other popular mask selection algorithms for N:M pruning.
Sparaty
CBS
FALCON
WF
CHITA
SNOWS
Figure 3: Comparing SNOWS to one-shot pruning methods for unstructured sparsity.
VIT pruning. ViTs (Dosovitskiy et al., 2021) have become popular in computer vision by applying the Transformer architecture, originally designed for natural language processing, to image recogni- tion tasks. In ViTs, an image is divided into a sequence of patches and embedded into a vector. These embeddings are then processed through blocks that consist of Multi-Head Self-Attention (MHSA) layers and FreddFoowward Neural Networks (FFNs). (Chen et al., 2022) introduce a framework to make the intermediate attention outputs sparse, while leaving the weights for the query, key, and value as dense matrices. They justify this saying it is particularly hard to prune these weights prior to the attention operations, since it is not the raw weight values that are important, but their interac- tions within the attention mechanism. In contrast, we jointly optimize over these weights and across all heads, and hence we inherently account for the interactions between the pruned weights and how they influence the attention output. The detailed formulation is provided in subsubsection A.3.3. We also discuss how to prune the output projection module of the MHSA block and MLP layers within the FFN modules in the VIT, which are more standard formulations than for MHSA.
9

---

Published as a conference paper at ICLR 2025
Experimental setup. We evaluate SNOWS on ViT-B/16 and ViT-L/16 models, pruning them to 2:4 sparsity on MinilmageNet-Ik Ik. Since methods for mask selection in N:M sparsity do not extend easily to vision transformers, employ MP to obtain the 2:4 masks.
Results. As shown in Table 2, SNOWS consistently outperforms MP, with higher Top-1 and Top-5 accuracies. When pruning the QKV, output projection, and MLP layers, SNOWS achieves Top-1 accuracy of 76.57% on ViT/B-16, significantly better than 70.89% for MP. In Figure 4, we plot the attention maps generated by the last attention layer of the dense model, the model pruned with MP and the model obtained when SNOWS is applied on top of MP. SNOWS better preserves the atten- tion patterns the dense model since it optimizes to preserve the activations via the reconstruction loss. In Figure 5, we show this leads to superior transfer learning performance, where fine- tinittnnng the classifier layer of a ViT pruned with SNOWS outperforms MP when transferring to smaller datasets.
Original
Dense Model
MP
MP*SNOWS
(a)
(b)
(c)
(d)
Figure 4: Visualizing (a) the original test image and attention maps from the last layer of (b) the dense VIT/B-16 model, (c) the model obtained by applying 2:4 MP mask, and (d) the model after applying SNOWS on top of MP. SNOWS optimizes to reconstruct learned activations, preserving features learned by the dense network even in the deepest layers.
Pruned Layers
Top-1
CIFAR-10
QKV (SNOWS)
79.45
Accoracy (S)
00
VIT/B-16
QKV (MP)
78.97
Denac
Dense: 80.42%
QKY, Out, (SNOWS)
HP
SNOWS
QKV, Out, MLP (MP)
70.89
VIT/L-16
QKY (SNOWS)
81.01
CIFAR-100
Dense: 84.17%
QKV (MP)
75.55
Acunccy (8)
MLP (SNOWS)
69.71
(MP)
Denoe
HP
SNOWS
Table 2: Pruning VIT/B- 16 and VIT-L-16 2:4 spar- sity on ImageNet-Ik in one shot. SNOWS uses the 2:4 mask obtained from MP. Note that no retraining involved. We use 5, 000 calibration samples for VIT-B/16 and n = 20, 000 for VIT-L/I6.
28
15
28
Epoch
Figure 5: Test accuracy fine-tuning the classifier layer of the 2:4 sparse VIT/B-16 on CIFAR- 10 and CIFAR-100.
6 LIMITATIONS AND CONCLUSIONS
This work proposes a new layer-wise pruning framework, bridging between the local layer-wise and global pruning loss functions introduced in prior work. The objective of our modified loss function induces a harder non-linear optimization problem compared to the standard layer-wise formulation. We propose an efficient Hessian-free method for optimizing this objective, and demonstrate its effi ciency on large-scale networks. There are several limitations to the proposed work: firstly, SNOWS does not do mask selection. This makes it flexible since it can integrate with existing techniques for mask selection, but using discrete optimization as is done in (Benbaki et al., 2023; Meng et al., 2024b;c) on the SNOWS loss function may lead to stronger results. Moreover, the method has little dependence the computational modules involved in the reconstruction loss. As such, the frame- work could readily be extended to Text-to-Speed to-Speech or Large Language Models. However, this would require specializing the formulation of the K-step problem
did for CNN and ViT modules.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGEMENTS
This research is supported in part by grants from the Office of Naval Research (N0001421 12841 and (000142212665)) We acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPCC resources that have contributed to the research results reported within this paper.
REFERENCES
Alireza Aghasi, Afshin Abdi, Nam Nguyen, and Justin Romberg. Net-trim: Convex pruning of deep neural networks with performance guarantee, 2017. URL httpss//axxii..rgg////////////// 1611, 62.
Roberto Battiti. First- and second-order methods for learning: Between steepest descent and new- ton's method. Neural Computation, 4(2):141-166, 1992. doi: 111162//nco.1992.422
Riade Benbaki, Wenyu Chen, Xiang Meng, Hussein Hazimeh, Natalia Ponomareva, Zhe Zhao, and Rahul Mazumder. Fast as chita: Neural network pruning with combinatorial optimization, 2023.
URL https:// /aapiv//aaivvvv/
org/ abs/2302 14623.
R. H. Byrd. L. Hansen, Jorge Nocedal, and Singer. A stochastic quasi- method for large-scale optimization. SIAM Journal Optimization, 26(2):1008-1031, 2016. doi: 10.1137/ 140954362. URL https://///// /tw.i org/ 1117/110954362.
Junyi Chai, Hao Zeng, Anming Li, and Eric W.T. Ngai. Deep learning in computer vi- sion: critical review of emerging techniques and application scenarios. Machine Learn- ing with Applications, 6:100134, 2021. ISSN 2666-8270. doi: http:///oo..rg//0.016/// 2021.100134. URL https: www. sciencedi rect ccmp//cciennce/atticee//pll $266682702100066
Zhaodong Chen, Yuying Quan, Zheng Qu, Liu Liu, Yufei Ding, and Yuan Xie. Dynamic n:m fine-grained structured sparse attention mechanism, 2022. 2203. 00091.
Hongrong Cheng, Miao Zhang, and Javen Qinfeng Shi. A survey on deep neural network pruning- laxonomy, comparison, analysis, and recommendations, 2024. URL aattps://arxine abs/2300 06767.
Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability, 2022. URL https: arxiv.. erg/abs/2100 00065.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp. 7480-7512. PMLR, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. pp. 248 255, 2009. doi: 10.1 12000688..20995266888..
Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017a. URL https: hattps://procee httpp///e////ccceenccccss. neurips cc/paper fille/paperr/2111 htt.:////SScccceeeeeeeeeeeeaeeee /tt.:///ee///ee/eooooacacaaaaaaa
Papeer pdf.
Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. Guyon, U. Von Luxburg, Bengio, H. Wallach, R. Fergus, S. Vish- wanathan, and R. Gamett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017b. URL https: httpppps///ccccennnns..eee neurips cc cc/paper. pdf.

---

Published as a conference paper at ICLR 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021. URL wwwps:ttt//oom
abs 2011 2.110.11922. 92...
Utku Evci, Fabian Pedregosa, Aidan Gomez, and Erich Elsen. The difficulty of training sparse neural networks, 2020. URL https: http//aaxvv....aaaa////
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks, 2019. URL https: /aaxivivvvv../////// org/ awtaaassssss 8038003.000 03635.
Elias Frantar and Dan Alistarh. Optimal brain compression: A framework for accu- rate post-training quantization and pruning. In S. Koyejo, S. Mohamed, A. Agar- wal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Pro- cessing Systems, volume 35, pp. 4475-4488. Curran Associates, Inc., 2022. URL
/trpperrf.lessppper/22222... neurips wttpssppcpppppppeeeeeeeeeee
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models be accurately pruned in one-shot, 2023. URL hwwph:tpps...m
http:/aabb2201.00770
Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of second-order information, 2021. URL https: /ttp//axxv...ggabbee/ssee 2107 03356.
Abhinay Goel, Caleb Tung, Yung-Hsiang Lu, and George K Thiruvathukal. A survey of methods for low-power deep learning and computer vision. In 2020 IEEE 6th World Forum on Internet of Things (WF-IoT), pp. 1-6. IEEE, 2020.
Manas Gupta, Efe Camci, Vishandi Rudy Keneta, Abhishek Vaidyanathan, Ritwik Kanodia, Ashish James, Chuan-Sheng Foo, Min Wu, and Jie Lin. Is complexity required for neural network prun- ing? case study global magnitude pruning. In 2024 IEEE Conference Artificial Intelli- gence (CAI), pp. 747-754, 2024. doi: 10. WWW.C/AAII1869222..00011000
Masafumi Hagiwara. A simple and effective method for removal of hidden units and weights. Neurocomputing, 6(2):207-218, 1994. ISSN 0925-2312. doi: https:/doii.og 10. 00.1016/0925--312299922 22110992255-122299990555558888 URL ht / /www. sci sciencedirect com/ scinnccence/ Backpropagation, Part IV.
Moritz Hardt, Benjamin Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent, 2016. URL https: http://////wxoooo/////////// 1atps./1.........
Babak Hassibi and David Stork. Second order derivatives for network pruning: Op- timal brain surgeon. In S. Hanson, J. Cowan, and C. Giles (eds.), Advances in Neural Information Processing Systems, volume htteeee_ieeeeappeee/99////////e neurips cc.psscccceeo wwpppper.pppppppeeeepeeeeeee..ee files pttppppppp////lelee/e/ 304cdec6B46ab33662993333333
1992. URL
Paper aapper.od pdf.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog-
nition, 2015. URL https: /rg/ /twt//b///////e /112.11111111 10512.131.333.
Yang He and Lingao Xiao. Structured pruning for deep convolutional neural networks: A survey. IEEE transactions on pattern analysis and machine intelligence, 2023.
Yang He and Joey Tianyi Zhou. Data-independent module-aware pruning for hierarchical vision transformers, 2024. URL /
22404 13648.
Dan Hendrycks and Kevin Gimpel. Gaussian eтror linear units (gelus), 2023. URL https://
arxiv. wr..org//aa abs 16066600 08415.
Duc N.M Hoang. Shiwei Liu, Radu Marculescu, and Zhangyang Wang. REVISITING PRUNING
AT INITIALIZATION THROUGH THE LENS OF RAMANUJAN GRAPH. In The Eleventh International Conference on Learning Representations, 2023. URL https: / /oppnreview. net/forut d=uVeDssOff
12

---

Published as a conference paper at ICLR 2025
Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andrectto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications, 2017. URL htttst///axxii... http://///abb///////////// 04861.
Sheng- Chun Kao, Amir Yazdanbakhsh, Suvinay Subramanian, Shivani Agrawal, Utku Evci, and Tushar Krishna. Training recipe for n.m structured sparsity with decaying pruning mask, 2022.
URL https:tpp/////...t.m /ttr///axxx /taaxaivvvooo .rgg aaaaab/s/2200 07617.
Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009. URL https: /wwww es toronto. wdw.:/nttt.eddeeeeeoeoeeeeeeeeee wttr///iiiiaiagg-eeeeeeeee wwwaarnnnn-eetuuee-s-rr 2009--RR- pdf.
Denis Kuznedelev, Eldar Kurtic, Elias Frantar, and Dan Alistarh. Cap: Correlation pruning for acchly-accurate sparse vision models, 2023. URL https: arxia. org. absp/210 09223.
Y. Lecun, L. Bottou, Y. Bengio, and P. Haffner. Gradient- based learning applied to document recog- nition. Proceedings of the IEEE, ()(11):227 2324, 1998. doi: 10.1 79.77999....699
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. In D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 2. Morgan-Kaufmant 1989. URL
newrips cc/ Paapp..ppp...
Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning, 2019. https http::/aaa.vvo.oeeee. wrtp:/s/////aaaa/.o.o /1100/11100000 05270.
James Martens. Deep learning via hessian-free optimization. In International Conference on Machine Learning, 2010. URL https: wewanticccoolaarre.............. wtt.ssscoooaaaooooorrrsosssstsee
11154521.
James Martens and Ilya Sutskever. Training deep and recurrent networks with hessian-free op- timization. In Neural Networks, 2012. URL https: //pp//app...//// http://ppp.....a.i.....c semanticscholar oor/
CorpusID:9247220
Xiang Meng, Kayhan Behdin, Haoyue Wang, and Rahul Mazumder. Alps: Improved optimization for highly sparse one-shot pruning for large language models. arXiv preprint arXiv:2406.07831, 2024a.
Xiang Meng, Wenyu Chen, Riade Benbaki, and Rahul Mazumder. Falcon: Flop-aware combina- torial optimization for neural network pruning, 2024b. URL https::/arxiv..rg/// /bbs/ 2403. 07094.
Xiang Meng, Shibal Ibrahim, Kayhan Behdin, Hussein Hazimeh, Natalia Ponomareva, and Rahul Mazumder. Osscar: One-shot structured pruning in vision and language models with combinato- rial optimization, 2024c. URL https: //tp:/i//////////////a//: org/ wwg/abs////eooooooooo...oo.ae.sr 2403 12983.
Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu. and Paulius Micikevicius. Accelerating sparse deep neural networks, 2021. URL https: arxiv. o abs////////////// 221./211440.88
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of paraectrtniinninnnonnn in generalization of neural networks, 2018. URL https://arxe org/ 1b00
Jorge Nocedal and Stephen J Wright. Conjugate gradient methods. Numerical optimization, pp- 101-134, 2006.
Jeff Pool and Chong Yu. Channel permutations for n:m sparsity. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.). Advances in Neural In- formation Processing Systems, volume 34, pp. 13316-13327. Curran Associates, Inc., 2021. URL
www.iingg..euuuiae hwt.erriisss.cppeee.eeeeeeeeeeee Gile/6e844 fttt:e/eeeocbbb33325222222222223
13

---

Published as a conference paper at ICLR 2025
Alex Renda, Jonathan Frankle, and Michael Carbin. Comparing rewinding and fine-tuning in neural 02389.
network pruning, 2020. URL http:::://a.....///.//oaoo/ www.a.itw..aaaaaaaaoooooo
Farbod Roosta- Khorrasii and Michael W. Mahoney. Sub-sampled newton methods i: Globally
convergent algorithms, 2016. URL https: bttpps:/aaxiv..o////aaa.
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural
network compression, 2020. URL htt.://se/aaaau...oo///////////p
wttp://///.///ooo////////// .13.004.4444900 14440.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and
Hervé Jégou. Training data-efficient image transformers and distillation through attention, 2021.
URL https:t//////..... aawiv.vviv. org/ abs /ttp/ 12877.
Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang. Eigendamage: Structured pruning in the kronecker-factored eigenbasis. In International conference on machine learning, pp. 6566- 6575. PMLR, 2019.
Fang Yu, Kun Huang, Meng Wang, Yuan Cheng, Wei Chu, and Li Cui. Width and depth pruning for vision transformers. In AAAI Conference on Artificial Intelligence, 2022a. URL https: // //api. www.nticscoooaatt..........maaae org/ Corpus 250294994
Hao Yu and Jianxin Wu. A unified pruning framework for vision transformers, 2021. URL https:
org/ www//oggg//b////eee/eee/eee abs/21s111111111 15127.
Xin Yu, Thiago Serra, Srikumar Ramalingam, and Shandian Zhe. The combinatorial brain surgeon: Pruning weights that cancel one another in neural networks, 2022b. URL https: arxiv. org/abs/ 2203 04466.
Yuxin Zhang, Mingbao Lin, ZhiHang Lin, Yiting Luo, Ke Li, Fei Chao, Yongjian Wu, and Rongrong Ji. Learning best combination for efficient n:m sparsity. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 941-953. Curran Associates, Inc., cc wttps/pr.cppaepeeeeeeeeeeeeeeeee httpp//peeee.lllepppe//////////
2022. URL file E1200999e8885550000011666eeeeee0
pdd..
Yuxin Zhang, Mingbao Lin, Yunshan Zhong, Mengzhao Chen, Fei Chao, and Rongrong Spatial re-parameterization for n:m sparsity, 2023. /
arwpiv wrg. 6306. 05612.
Aojun Zhou, Yukun Ma, Jumnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hong- sheng Li. Learning n:m fine-grained structured sparse neural networks from scratch, 2021. URL
https: http///aavvvv//////////ee//..ave org/ abss 2202.0022220 04010.
Hao Zhou, Jose M. Alvarez, and Fatih Porikli. Less is more: Towards compact cnns. 2016. doi: 10. 1011007997883333194449999-.0.0
Mingjian Zhu, Yehui Tang, and Kai Han. Vision transformer pruning, 2021. URL https:// a arxiv.org/abs/ 08500.
A APPENDIX
A.1 PRIMARY METHODS AND ANALYSIS
A.1.1 BACKGROUND ON OBS AND SECOND-ORDER APPROXIMATIONS.
For pruning the network achieve the sparsity target, the proposed OBS solution is start from the dense weights W and minimize a second-ore -rderr Taylor expansion of the loss function with respect perturbation of the weights:
min L(W + + - -
(II)
14

---

Published as a conference paper at ICLR 2025
subject to the constraint that 6, W = 0, ie. single weight W is to be removed, with VL(W) 0L/0W being the gradient and H @L/OW22 being the Hessian matrix. The OBS procedure minimizes this approximation iteratively with respect the weight perturbation 6, each time obeying the constraint that single weight is removed. Assuming the dense network is fully trained, they argue VL(W) = O, which leaves just the second-order term involving ò in the approximation. Minimizing the resulting constrained formulation with respect to Å using Lagrange multipliers yields exact saliency scores for each weight Wy. given as Py
W 1 H
eg as
well as a method for optimally readjusting the remaining weights ó [H H H where Cy
the standard basis with all zeros except the q-th element which is one. Here, Py is solution to the discrete problem (a ranking among all weights) while 5 is solution to the continous problem (an optimal weight readjustment for weights other than q, when & is removed).
A.1.2 MOTIVATION FOR THE SNOWS LOSS FUNCTION
The original motivation of the layer-wise wise reconstruction objective given in Dong et al. (2017b) was minimize the impact of pruning at each layer the final network output. They show that the error ماع - 1/1 of the entire pruned network the final layer L is bounded as:
E
W
VdE
VOEL
where SE is the layer-wise error between the outputs layer L. i... SE
In other words, the difference between the final output of the dense and pruned network is bounded by product of the layer-wise errors. This means if the layer-wise erross can be kept small, the overall error should be controlled. However, there is an interesting subtlety to how the errors propagate. Say we are pruning layer f. It would seem as though the contribution by layer to EL is confined to ò However, notice that
f'= W
So pruning at the current layer changes the input the next layer and moreover all subsequent layers since X + = f ( We+k-1), for k 1,..., L l.
So after pruning layer f, we can observe the expression for the emror in the subsequent layers
8F = 1 x x W X
At the time of pruning layer f, With Weth since these weights have not yet been pruned, but notice that already X &++ * X since our pruning decision at layer has affected all layers after it. The error term for the future layers becomes
1 n || x W
SE
Hence the error at the future layers is determined by the magnitude of the difference c+h induced by pruning layer e. Our formulation, by minimizing Σ l| Yetk1
F
K [ x can be interpreted as directly minimizing the subsequent terms F during the pruning of layer l as opposed to only SE
A.1.3 ARMIJO SUFFICIENT DECREASE CONDITION.
In the Newton descent procedure, we use the Armijo decrease condition to determine an appropriate step size that ensures sufficient decrease in the loss function:
15

---

Published as a conference paper at ICLR 2025
argmax a s.t. C <(W2 + tdzz) SC < agl
a 03352VV VC (2)
The parameter 3 is small constant. We use 10-5
A.1.4 A TOY EXAMPLE SHOWING THE ADVANTAGE OF NEWTON'S METHOD OVER SGD UNDER VARYING CURVATURE.
Consider the following toy problem:
1 C(w)) (Arwp A2u3)
(12)
2
where =(w1, Wy]" is the weight vector and λ1 the loss (curvature) along the wi and W2 axes, respectively. Now consider a situation where Å1, A2 very different. For example, (high curvature along W2). The condition number of the problem is Ai >max A = > = 100.
are constants representing sensitivity of
choose A1 (low curvature along w1) and X2 100
This creates ill-conditioned problem with significant difference curvature between the two directions. Now compare the convergence behavior for Newton and SGD. For SGD the update is:
w (++1) (N) 7/ (V
where 7) is the learning rate and VL(w)- = [m
The updates can written as:
w (++1) w(()
n>1 w m/1111 (b) (1 m7111 ) (R)
(++1)
17/2112 = =(1 1-112/ 10 1112)1122 (k)
112
W12
To ensure convergence, the learning rate needs to satisfy 0 < 1) < Amax e.g. see (Cohen et al., 2022) for a detailed convergence analysis. With Amax A 100, natural choice is = 0.01. This
2
implies:
(k+1) (1 )w? (b) (1 0.01 x 1) (k) = 0.99w, (R)
This means wy decreases by 1% per iteration, leading to slow convergence along this direction due to the low curvature. Namely, rewriting the update formula:
(0.99)>w (0)
w1
The optimal solution is clearly 0. Assuming wi (0) 1, we can find the number of iterations k required for wi to be less than a small threshold e.g. 16
In(10-6) In(9.9
(M) u1 Se
Rs 1,375 iterations.
In contrast, along Wz!
(k+1) = - 1222 = (1 - 0.01 x 100)wa (R) = Ove = 0. 22
Here, W2 converges to the optimum in one iteration. Thus, importantly overall convergence is dom- inated by the slowest direction (low curvature along w1). Now consider Newton's method on the same problem:
w (k+1) (A) H VV
16

---

Published as a conference paper at ICLR 2025
A 0
0
where the Hessian is =
and its inverse is: --1
X2
X2
The updates become:
(k+1) (8)
Ai A (k) = 0,
11
122 112 A2 X2 (k) = O.
(k+1)
(k)
Thus both wt and w2 reach the minimum in one iteration, regardless of the condition number k. Newton's 's method overcomes the issue of varying curvature by adjusting its updates according the inverse of the Hessian. Rescaling the gradient components means takes appropriate sized steps in each direction, i.e. larger steps in low-curvature directions and smaller steps in high-curvature ones. This results in convergence to the minimum single iteration for both directions, regardless of the condition number K. The trajectories shown in Figure 6.
Gradient Descent Trajectory
Newton's Mathod Trajectory
1.0
1.0
(low curvature)
(low curvature)
(high curvature)
0.8
Wy (high curvature)
wwightieba
0.6
0.6
0.6
0.-
0.2
0.0
200
400
6:0
500
2000
1200
1000
0
Baration
Beration
Figure 6: Comparison of convergence rates between SGD and Newton's 's method on the toy opti- mization problem in Eqn (12) with highly varying curvature (Λı = 1, A2 = 100). While Newton's ss method achieves convergence single iteration, SGD requires approximately 1,375 iterations due to the high condition number 100.
A.2 ABLATION STUDIES
A.2.1 COMPARING SGD VERSUS NEWTON'S METHOD TO OPTIMIZE EQN (3).
In Figure 7, we compare the time required to optimize the first layer of ResNer20 on CIFAR-10 and ResNet50 ImageNet-Ik, both using 2:4 mask Z obtained via magnitude pruning. The Stochastic Newton algorithm, when applied to optimize Wz, makes significantly faster progress compared to SGD, especially as problem size increases. In the section that follows provide detailed loss trajectories for both methods on this same problem. Our results indicate that Newton's method moves much farther from the starting solution with each step, whereas SGD progresses much more slowly and generally does not move far from the starting solution. We also provide more detailed discussion possible reasons why SGD converges slowly in this setting.
17

---

Published as a conference paper at ICLR 2025
560 (1=0.011
Normalizd Loss (lq Scale)
560 (P=0.003)
(1-00,001)
100
150
Figure 7: Comparing Newton's method and SGD with minimizing Eqn (3) with K = 10, for the first layer in the model. The solution starts with 2:4 sparse mask obtained via magnitude pruning. Newton's 's method converges very quickly and requires less time to solve the K K-step reconstruction problem. SGD is trained for 2000 steps whereas Newton converges in less than 10.
The empirical success of SGD in training neural networks from scratch has been very well- demonstrated (Hardt al., 2016). Several papers have remarked that this success seems owe in part to the effect of overparameterization (Neyshabur et al., 2018), where neural networks are trained with more parameters than are necessary for a given task. At the same time, empirical work on sparse fine-tuning has shown that fine-tuning sparse network with SGD often does not lead to good performance (Renda et al., 2020). Other work has found that training a sparse network from scratch with knowledge of good initialization is likewise difficult (Evci et al., 2020). Our em- pirical findings in Figure 7 corroborate this fact, suggesting SGD performs poorly and converges slowly the optimization problem defined in Eqn (6). In Figure 8 show the weight trajectories for the same problem, finding that SGD does not move far from the starting solution, even after thou- sands of iterations. We hypothesize that SGD may perform poorly in the sparse setting due uneven curvature (loss sensitivity) in the coordinates, which is also present in the overparameterized setting sparse networks.
but may become exaggerated by the lack of redundancy show a small toy example comparing SGD and Newton's 's method for a swo-dimensional quadratic optimization problem where the weight directions have varying eurvature. The intuition that SGD must choose a learning rate larger than the largest eigenvalue of the Hessian since if it is larger, then loss will diverge on that direction. In overparameterized problems, it may be possible to get away with having long-tail of low curvature directions that take more time to converge. However, in sparse models redundancy is significantly reduced, and this means the optimizer must be more sensitive to each direction. In such cases, choosing learning rate that works well for all directions becomes increasingly difficult, as there are fewer parameters to absorb the mismatch between the learning rate and the curvature.
the next section, we
18

---

Published as a conference paper at ICLR 2025
ResNetZ0 ICIFAR 10)
RcsNet50 (ImageNet-110
Thelectry
SGD Lau Contour with Fejedory
Rowler Cantar With Thgedtory
ND Ranine
Srwin Suiitny
A150
Nevin Tiiettoy
Ad813
4.20
+0.3
+8.00 +0.00
-0.33
04.40
cafs
0.25
8:90
A0/o
0415
9400
(a) ResNet20: Loss surface for SGD (left) and Newton (b) ResNet50: Loss surface for SGD (left) and Newton (right).
(right).
Neeights
Rebstive Distance Weigtts
Gi Siatactic Neatre
SCD
(c) ResNet20: Relative distance between full weight (d) ResNet50: Relative distance between full weight matrix and starting weights.
matrix and starting weights.
Figure 8: Optimization of the first layer for both ResNet20 and ResNet50 using SGD and stochastic Newton from Figure 7. Top row: Loss trajectories for 2000 iterations of SGD versus respectively 6 (CIFAR-10) and 13 (ImageNet-1k) iterations of stochastic Newton for the two largest starting weights. Bottom row: Relative distance between the full weight matrix and the starting weight
matrix for both methods. The relative distance is defined as dist(W,) |W,-- Wol the weights at iteration & and Wo are the initial weights.
where W, are W.ll
A.2.2 USING TRUE HESSIAN VERSUS FISHER APPROXIMATION TO OPTIMIZE EQN (3)
We also observe significant performance improvement when using the true Hessian over the Fisher approximation. Specifically, we replace the Hessian in Eqn 7 with the Fisher approxima- tion F= = / Σ V W, (. wz (, T and employ the updates 8z -F--FF-LWWW The main advantage of employing the true Hessian via the Hessian- free approach is the accelerated speed of convergence. As illustrated in Figure 9, we plot the objective value when using the Fisher approximation versus the exact Hessian-free steps for all layers in ResNet20 and the first three layers in ResNet50. The results indicate that using Newton's 's method instead of the Fisher approximation leads to much faster convergence, not only in terms of iterations but also in overall runtime. For the larger ResNet50 networks, running the Fisher approximation to convergence is actually prohibitive, and can take almost half an hour per layer for even the first layers of ResNet50.
(a) ResNet20
(b) ResNet50
Figure 9: Comparison of Newton's method and Fisher approximation for all layers in ResNer20 and the first three layers of ResNet50.
A.2.3 EFFECT OF THE NUMBER OF CG ITERATIONS
The plot in Figure 10 shows the impact of varying the maximum number of CG iterations in Al- gorithm on the reconstruction loss across different layers and models. As expected, decreasing
19

---

Published as a conference paper at ICLR 2025
the maximum CG iterations generally results in slower progress per iteration, lower caps (e.g., CG 5) reduce the loss at a slower rate compared to higher iteration limits (e.g., CG 500). However, despite this slower per-iteration progress, CG with fewer iterations often converges faster overall, as the algorithm compensates by taking a more iterative approach, covering more batches quickly while solving the second order approximation for each batch more approximately. This means that, although each step is smaller, CG with lower iteration caps refines the solution more frequently, leading to comparable or even faster convergence relative higher iteration caps.
Varying maximum number of iterations
Ra120 (1R Ressa (OPAN- 180) Reedo (lowgeter)
Layer
Layor2
Layer
20-1
4x
Ex
6x
10
3'"
202
300
400
40
150
140
107
6x
4x-
10-1
0.0
0.5
2.0
255
2.0
Layer3
10"-
10"
10a
55 12"
1x:12
4x.12"
60-0--
25:
3x 10
Iteration
C6
CG 10
50
CG 100
500
Figure 10: Normalized reconstruction loss starting from the masked dense weights. The plot shows effect of varying the maximum CG iterations in Algorithm 1. While fewer iterations result in slower progress per step, they often lead to faster overall convergence by refining the solution more frequently.
A.2.4 EFFECT OF THE PARAMETER K ON RUN-TIME
We provide in Table 3 a comparison (1) run-time in minutes for pruning the entire network and (2) peak GPU memory usage, the parameter K varies. In practice K should be chosen as large as possible while not exceeding available GPU memory or run-time expectations. In practice moderate values of K can lead to good performance while having very small run-times.
ResNet20 (CIFAR-10)
ResNet50 (CIFAR-100)
ResNet50 (ImageNet-Ike
K
Time
Peak Memory (GB)
Time
Peak Memory (GB)
Time
Peak Memory (GB)
0
0.7
2.2
11.3
8.7
10.5
8.6
1
0.8
2.5
10.7
10.8
12.9
15.1
3
1.6
2.8
11.1
11.4
25.2
17.3
5
2.3
3.3
11.4
12.1
39.7
25.8
10
4.8
3.3
13.3
15.5
94.8
38.3
20
8.1
4.3
16.1
21.9
184.1
54.0
30
11.2
17.0
24.9
289.5
66.9
40
11.8
5.1
18.8
28.1
0OM
O0M
60
max
max
20.3
31.6
OOM
OOM
Table Effect of varying on the run-time and peak memory usage of the full SNOWS algorithm. The run-time shown is in minutes.
A.3 IMPLEMENTATION AND EXPERIMENTAL DETAILS
A.3.1 N:M SPARSITY IN CNNS
N:M sparsity enforces a fine-grained sparsity pattern along specific dimensions of the weight tensor. In the context of convolutional neural networks, apply N:M sparsity along the input dimension
20

---

Published as a conference paper at ICLR 2025
di while keeping the other dimensions fixed. This means that, for each output channel n and each spatial location in the kernel (i.j). we group the weights W n.G.i.J along the imput channels into groups G of M consecutive weights such that [lW.,.... ll = N. Within each group, exactly N weights are non-zero, and the remaining - weights are zeroed out. An example for 2:4 sparsity is shown in Figure 11.
dou x knp x lew
d
Figure 11: An example of 2:4 sparse convolutional kernel.
A.3.2 K-STEP LAYER-WISE RECONSTRUCTION PROBLEM IN CNNS.
Each convolutional layer typically involves several linear and non-linear operations. The exact op- erations can differ across architectures, but we present example for the ResNet architecture (He et al., 2015). The functions in Eqn (3), which captures the operations at layer f, can be mathemat- ically defined f((X,, W Conv(X', W*) refers to the convolution operation. The Conv operation can be written as a linear operation so for K = 0 this problem reduces to the layer-wise reconstruction problem in Equation 3. For larger K, say K 1, we define
(X W ReLU (BN (f'("', W xx)
where BN refers to batch normalization, ReLU(M) = max(M, 0) and the addition of X the residual connection. For K 1, directly jump to reconstruct the output of the ReLU since BN a small linear rescaling applied to the convolution, so we do not directly reconstruct its output and likewise for residual connections, though these inherently absorbed in the reconstruction the output of the ReL.U.
A.3.3 RAw NUMBERS FOR UNSTRUCTURED COMPARISONS
For Figure 3, the benchmarking numbers from (Benbaki et al., 2023) and (Meng et al., 2024b). However, their dense versions ResNet20 and MobileNetV1 have slightly different start- ing accuracy than ours. This the reason report relative accuracy the baseline in Figure 3.
21

---

Published as a conference paper at ICLR 2025
ResNef20 on CIFARIO
Sparsity
Baseline Accuracy: 91.36%
Baseline: 92.56 %
MP
WF
CBS
CHITA
FAL.ON
SNOWS
0.5
88.44 (-2.92)
90.23 (-1.13)
90.58 (-0.78)
90.60 (-0.76)
90.87 (-0.49)
92.28 (-0.28)
0.6
85.24 (-6.12)
87.96 (-3.40)
88.88 (-2.48)
89.22 (-2.14)
89.67 (-1.69)
91.90 (-0.66)
0.7
78.79 (-12.57)
81.05 (-10.31)
81.84 (-9.52)
847.2
84.42 (-6.94)
90.87 (-1.69)
54.01 (-37.35)
62.63 (-28.73)
51.28 (-40.08)
57.90 (-33.46)
65.17 (-26.19)
87.82 (-4.74)
0.9
11.79 (-79.57)
11.49 (-79.87)
13.68 (-77.68)
15.60 (-75.76)
19.14 (-72.22)
64.80 (-27.76)
MobileNetV1 on ImageNet-1K
Sparsity
Baseline Accuracy: 71.95%
Baseline: 70.89 %
MP
WF
CBS
CHITA
FALCON
SNOWS
0.5
62.61 (-9.34)
68.91 (-3.04)
70.21 (-1.74)
70.42 (-1.53)
70.35 (-1.60)
70.10 (-0.79)
0.6
41.94 (-30.01)
60.90 (-11.05)
66.37 (-5.58)
67.30 (-4.65)
67.18 (-4.77)
68.70 (-2.19)
0.7
6.78 (-65.17)
29.36 (-42.59)
55.11 (-16.84)
59.40 (-12.55)
58.40 (-13.55)
65.28 (-5.61)
0.8
0.11 (-71.84)
0.24 (-71.71)
16.38 (-55.57)
29.78 (-42.17)
25.82 (-46.13)
54.83 (-16.06)
Table 4: Raw accuracy (%) with absolute accuracy change (in brackets) compared to the baseline. The baseline accuracies are shown above the methods.
A.3.4 RETRAINING EXPERIMENTS
In Table 5, we provide retraining experiments for SNOWS and the best competing method from unstructured sparsity experiments, FALCON Meng al. (2024b). The retraining experiments use step decay learning rate schedule to fine-tune. For CIFAR-10, 10, we retrained ResNet20 for 90 epochs, using SGD with momentum 0.9 and weight decay 5e-5. The learning rate started 3e-3 and was reduced by factor of 0.1 every 30 epochs. For ImageNet- 1K, we retrained MobileNetVI for 5 epochs using the same optimizer settings, with the learning rate decaying every epochs. Throughout retraining, only non-zero weights to be updated. As shown in the results table, this retraining strategy was particularly effective at high sparsity levels (0.8-0.9), where SNOWS significantly outperformed FALCON in retaining model accuracy.
maintained the pruning mask preserve the network sparsity, allowing
FALCON
SNOWS
Dataset
Sparsity
Original
+RT
Original
+RT
0.5
90.87 (-0.49)
91.07 (-0.29)
92.28 (-0.28)
92.32 (-0.24)
0.6
89.67 (-1.69)
90.58 (-0.78)
91.90
92.13 (-0.43)
ResNet20
0.7
84.42 (-6.94)
89.64
90.87 (-1.69)
91.85 (-0..1)
(CIFAR-10)
0.8
65.17 (-26.19)
87.59 (-3.77)
87.82 (-4.74)
91.06 (-1.50)
0.9
19.14 (-72.22)
81.60 (-9.76)
64.80 (-27.76)
87.69 (-4.87)
0.5
70.35 (-1.60)
70.66 (-1.29)
70.10 (-0.79)
70.66 (-0.23)
MobileNetV1
0.6
67.18 (-4.77)
68.89 (-3.06)
68.70 (-2.19)
70.10 (-0.79)
(ImageNet-IK)
0.7
58.40 (-13.55)
64.74 (-7.21)
65.28 (-5.61)
68.76 (-2.13)
0.8
25.82 (-46.13)
53.84 (-18.11)
54.83 (-16.06)
64.70 (-6.19)
Table 5: Raw accuracy (%) and absolute accuracy change (in brackets) for FALCON and SNOWS before and after retraining different sparsity levels. The baseline accuracies are as given in Table 4.
A.3.5 INFERENCE TIME SPEED-UPS FOR SEMI-STRUCTURED SPARSITY
Calculating exact inference time speed-ups for N:M sparsity patterns presents significant challenges due to limited native PyTorch support and the need for specialized CUDA kernels to realize speed- ups. Moreover, currently only the 2:4 pattern is supported by NVIDIA's sparse tensor cores. To
22

---

Published as a conference paper at ICLR 2025
provide meaningful performance metrics, instead report end-to-end FLOPs reduction across ious architectures and sparsity patterns in Table 6, accounting for all network operations and layers.
Table 6: FLOPs and Accuracy Reduction Across Models and Semi-Structuerd Sparsity Patterns
Model
Dataset
Sparsity Pattern
FLOPs Reduction
Accuracy Change
ViT-B-16
ImageNet-Ik
2:4
1.97x
-3.85%
1:4
2.99x
-9.38%
ResNe150
ImageNet-Ik
2:4
1.79x
-1.74%
1:4
3.13x
-1.35%
ResNet50
CIFAR-100
2:4
1.83x
-0.16%
1:4
3.80x
-3,49%
ResNet20
CIFAR-10
2:4
1.97
-0.64%
A.3.6
HYPERPARAMETER CONFIGURATIONS FOR SNOWS USED IN EXPERIMENTS
K-Step Value
Architecture
Dataset
Dampening À
CG Tol
N:M Sparsity
K Unstructured Sparsity
ResNet20
CIFAR-10
K
le-4
le-3
ResNet50
CIFAR-100
K 100
le-4
le-3
ResNet50
ImageNet-IK
K = 30
le-4
le-3
MobileNet
ImageNet-Ik
K
le-4
le-3
VIT/B-16
ImageNet-IK
le-4
5e-4
K 4 (QKV)
VITVL-16
ImageNct-1k
1e-4
5e-4
K (MLP)
Table 7: K Ktstep values, dampening factors, and CG tolerances used for N:M sparsity and unstruc- tured sparsity experiments across different architectures.
A.3.7 SNOWS ALGORITHM RUNTIME
For moderate values of K, our approach performs better and can run much faster than prior ap- proaches, and scales much better terms of architecture size. In Table 6, we compare the runtime of our algorithm to CHITA (Benbaki et al. (2023)). We remark that other competitors WoodFisher and CBS require hours to prune even MobileNet with a small sample size, see section 4.1.1 of (Benbaki et al. (2023)). We show below the run-time for pruning three networks of different sizes: ResNet20 (250k parameters), MobileNet (4m parameters), and ResNet50 (25m parameters). The main advantage of our method in terms of efficiency is scaling to large networks since avoid computing. inverting, or storing the full Hessian or its approximation. and moreover not even stor- ing inverting the smaller layer-wise Hessian. Notably, our method outperforms CHITA even with moderate values K (though larger values used in our final experiments can obtain the best results e.g. in Figure 3).
23

---

Published as a conference paper at ICLR 2025
SNOWS
CHITA
Model-Sparsity
Samples
Acc (%)
Runtime (min)
Ace (%)
Runtime (min)
ResNet50 (25m parameters)
500
75.7
9.66
73.6
231.20
MobileNet (4m parameters)
1000
61.1
5.56
59.4
16.82
ResNer20 (250k parameters)
1000
90.0
0.56
84.1
2.23
Table 8: Comparison of performance versus run-time comparing SNOWS and CHITA pruning dif- ferent architectures to 70% unstructured sparsity. ResNet-20 is run CIFAR-10, MobileNet ImageNet-1K, and ResNe150 on CIFAR-100. We use 1000 samples for ResNer20 and MobileNet, but for ResNet50 we use 500 samples since this was the only sample size below which CHITA runs in <5 hours.
A.3.8 LAYER- WISE RECONSTRUCTION PROBLEMS FOR VIT MODULES.
WEWER-WISE
MHSA Pruning. In a Transformer model with L layers, the input to layer & denoted as X E Raxd where is the sequence length and d is the embedding dimension. The Multi-Head Self- Attention (MHSA) mechanism at layer & uses multiple attention heads to capture different aspects of the input sequence. Specifically, the input X is linearly projected produce the concatenated queries (Q'), keys (K*), and values (Vt) across all heads:
Q' = =''W6. ='=X'Wa V'=
where W W Ko W E RdxHdd are the learnable weight matrices for all heads, with H being the number of heads and du the dimensionality each head. The self-attention mechanism then computes the relevance of each token in the sequence by performing the scaled dot product of queries and keys, followed by a softmax operation to produce normalized attention scores:
Y softmax (o'K²²)
V
Here, Y represents the output of the MHSA layer at the current layer É. In the general multi-step layer-wise reconstruction problem for K 2 O, we aim to reconstruct the outputs over the subsequent K layers while optimizing jointly the query, key and value matrices. This is expressed
K
y px (W W W6))
min W'WW
where f+.++ (W q: W which can include the FFN outputs, GeLU activations, Layer Normalization outputs, and subse- quent attention operations. For specific case: when K = 0, we focus on reconstructing only the output of the attention block in layer é, which can be written as:
K: W represents the transformation of the input X through K lay-
y softmax min W.EW
X R'W
Vd
Through joint optimization we inherently account for the interactions between the pruned weights. Moreover, since we solve the problem across all attention heads, we can learn redundancies that exist across multiple heads, thereby preserving the overall expressiveness the MHSA block after pruning.
Output Projection Pruning. In MHSA, after the attention operation, the concatenated outputs of the attention heads are linearly projected back to the original embedding dimension d via an output projection matrix W Ho xd, This matrix reduces the dimensionality of the concatenated outputs z' concat Z "...., ZH), where each Z is the output from one of the H heads. The output projection is computed as:
y - z'w!
24

---

Published as a conference paper at ICLR 2025
In the case where K the output projection pruning reduces to standard linear least squares reconstruction problem.
min Wgww
Y Z' W
For larger K, the output projection reconstruction problem involves reconstructing the outputs of the MLP, Attention blocks that proceed it.
FFN Pruning. In vision transformers, feedforward networks (FFNs) are responsible for a substan- tial portion of the overall model parameters, often accounting for more than half of the total param- eter count and FLOPS (Kao al., 2022). The FFN block applies an initial linear transformation, followed by GeLU activation function (Hendrycks & Gimpel, 2023), followed by an additional linear transformation. To optimize and prune the FFN, we break the MLP into its constituent linear parts, addressing each separately. Specifically, we handle the first linear layer (MLPO) and the sec- ond linear layer (MLP3) as distinct optimization tasks. As special cases, for = 0, both MLPO and MLP3 reduce to standard linear reconstruction problem:
min
X - X
W'EW
For K = 1, MLPO becomes a non-linear problem involving the GeLU activation function:
Gel.U(X'W b²) GeI.U(X W bb)) 2 W'ew
min
2
For larger K, the MI.P reconstruction problem involves reconstruction subsequent attention opera- tions and MIP blocks. In all cases, MI.PO and MLP3 are optimized independently.
25


Reviews:
[{"review_id": "am04b2fm9w", "reviewer": "Reviewer_TGxm", "summary": "The paper argues it is important to consider the effect of pruning of the whole network instead of just layerwise when performing one-shot post-training pruning. They use a Hessian-free second-order method to optimize their formulation. The results show significant improvement in image-classification tasks on resnets and vits.", "questions": "1. There are a few cases where other pruning methods fail (very low accuracy, Fig. 3, Tab. 1), could you comment on the potential reasons?\n2. Does retraining improves the performance further after pruning using the proposed method?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The formulation is neat, and the use of tools from optimization literature (Hessian-free conjugate gradients) is elegant.\n2. The results show significant improvements over OBC and pruning using the layerwise objective.\n3. Writing is clear and the main claims are supported by the experiments.", "weaknesses": "1. The emphasis on using true Hessian information is not supported in the experiments. This is not a crucial point but in most scenarios in deep learning, Fisher matrix-based approximation of Hessian [a] (or even Ada-grad [b] or Adam) works well. It would be worth providing that baseline to better support the proposed method.\n\n- [a] Martens, James. \"New insights and perspectives on the natural gradient method.\" Journal of Machine Learning Research 21.146 (2020): 1-76. \n- [b] Duchi, John, Elad Hazan, and Yoram Singer. \"Adaptive subgradient methods for online learning and stochastic optimization.\" Journal of machine learning research 12.7 (2011).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Usx7ukOQBJ", "reviewer": "Reviewer_XYMY", "summary": "The following paper presents SNOWS, a one-shot post-training pruning technique that performs pruning at a layer-wise level while globally optimizing joint layers within a deep neural network to include nonlinear activation functions within. While it resulted in a more challenging optimization problem due to nonlinearities, it is possible to use the Hessian-free approach integrated with a customized Conjugate Gradient (CG) method that efficiently solves second-order approximation during gradient computation while simultaneously allowing reduced memory requirements. Experiments on various Computer Vision tasks under both CNN-based architecture and Vision transformer highlighted the effectiveness of SNOWS.", "questions": "See the weakness section.\n\n===EDIT: Rating is raised from 3 to 5=====", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Besides achieving performance improvement compared to other one-shot pruning methods, the proposed method also improves the classification performance in image classification tasks when built on top of existing sparse mask selection algorithms for N:M pruning, denoting its practicability for accelerated training.\n2. Highly appreciate the fact that the source code is available for inspection.", "weaknesses": "1. There are no conclusions and limitations section in the paper, and the writing in the methods section can be more concise in Section 4, which mainly focuses on explaining the SNOW method. This could leave more space to explain the experiment section in the main paper properly along with the results. Overall, the organization of this paper needs a lot of improvement and it feels like the paper is not ready for publication in the current form. \n2. Typos:\n- \"Like retraining approaches\" -> \"Like retraining, approaches\" (Section 1)\n- \"network parammeters $\\mathbf{W} = (\\mathbf{W}^1, ..., \\mathbf{W}^1)$ should have been $\\mathbf{W} = (\\mathbf{W}^1, ..., \\mathbf{W}^L)$ (Section 3)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VfJyDYbfbo", "reviewer": "Reviewer_96mw", "summary": "This paper proposes a novel pruning method called SNOWS (Stochastic Newton Optimal Weight Surgeon), which is an adaptation of pruning methods like the Optimal Brain Surgeon method, that uses the second-order information of the weights to do the pruning.\nHowever, SNOWS does not calculate the full hessian but merely the hessian products, and those too are not all calculated (approximated) at once and the optimization is performed using stochastic newton descent.", "questions": "Q1- Since this method is able to achieve semi-structured sparsity, how does its latency compare to unstructured pruning methods? A comprehensive latency study would be very helpful.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The following are the strengths of the paper:\n1. The idea contributed by this paper seems to be novel and is very interesting. \n\n2. The theoretical aspect of the paper seems sound.\n\n3. The proposed method seems universally applicable to convolution and ViT-based models. \n\n4. The writing of the paper is quite clear.", "weaknesses": "The paper's introduction makes some glaring errors in defining concepts. For example, it describes one-shot pruning methods as pruning methods that do not require retraining after pruning. However, this is not accurate. Pruning can be done iteratively or one-shot; however, some finetuning might be required after one-shot pruning, as shown by [1], [2], and many other methods covered by [3]. \n\nWith the above definition for one-shot in mind, [4] achieves unstructured pruning with more than 90% sparsity while not dropping any clean performance for multiple architectures. The proposed SNOWS method should also be evaluated against [4] for unstructured pruning performance. \n\n\nLastly, the major motivation of the proposed work is gain in speed when performing pruning (as the entire hessian need not be calculated), however these no exists latency comparison to other pruning methods. \n\nImproving on these weaknesses might require a major rewrite, and thus the recommendation.\n\n**References**\n\n[1] Sanh, Victor, Thomas Wolf, and Alexander Rush. \"Movement pruning: Adaptive sparsity by fine-tuning.\" Advances in neural information processing systems 33 (2020): 20378-20389.\n\n[2] Sehwag, Vikash, et al. \"Hydra: Pruning adversarially robust neural networks.\" Advances in Neural Information Processing Systems 33 (2020): 19655-19666.\n\n[3] Hoefler, Torsten, et al. \"Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks.\" Journal of Machine Learning Research 22.241 (2021): 1-124.\n\n[4] Hoffmann, Jasper, et al. \"Towards improving robustness of compressed cnns.\" ICML Workshop on Uncertainty and Robustness in Deep Learning (UDL). Vol. 2. 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8QOMohaaYs", "reviewer": "Reviewer_EtDe", "summary": "(+) The authors propose optimizing a more global reconstruction objective referred to as SNOW, which accounts for non-linear activations deep in the network to obtain a better proxy for the network loss.\n\n(+) This nonlinear objective leads to a more challenging optimization problem - the authors demonstrate that it can be solved efficiently using Hessian-free optimization.\n\n(+) SNOWS demonstrates its effectiveness with state-of-the-art results on various one-shot pruning benchmarks, including  CNNs and Vision Transformers (ViTs).", "questions": "Please see the above weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(+) SNOW can be seamlessly applied to any sparse mask from prior methods, refining weights to leverage the nonlinearities in the deep feature representations of CNNs and ViTs.\n\n(+) The updated rule (Newton) and the illustrative examples in the Supplementary section are well articulated.", "weaknesses": "(-) The motivations for the improved layer-wise model (Eq. 3) need clarification. Why are the subsequent operations $f^{l:l+K}$ necessary-are they intended to minimize residuals? Please clarify the motivation behind SNOW by explaining the concept of Fig. 2 and including additional captions.\n\n(-) While SNOW’s performance with increasing K iterations is well documented, computational comparisons with baselines are lacking. For consistent comparisons, all experimental results should explicitly include the number of K-iterations in captions of tables and figures. Moreover, I recommend the authors add an ablation study on the time complexity of K-iterations to the experimental result tables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wWDeCjTgEN", "reviewer": "Reviewer_YAxu", "summary": "This paper proposes SNOWS, a one-shot pruning framework aimed at reducing vision network inference costs without retraining. Unlike traditional approaches that prioritize layer-wise error minimization, SNOWS optimizes a global objective that captures deeper, nonlinear network representations. To efficiently solve the resulting complex optimization problem, SNOWS leverages a Hessian-free second-order method that computes Newton descent steps without needing the full Hessian matrix. SNOWS also enhances any pre-existing sparse mask by refining weights to better leverage nonlinearities in deep features. The framework demonstrates state-of-the-art results across multiple one-shot pruning benchmarks, including residual networks and Vision Transformers. Results are shown for CIFAR10/100 and Mini-Imagenet.", "questions": "Suggest that the authors define the MP metric/method when it is first used and provide a citation if it refers to a specific published method.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "**===After rebuttal, I raised score from 5 to 6=====**\n+ A method with a new pruning strategy at the level of global reconstruction shows promising results on common datasets CIFAR and Mini-Imagenet.\n+ The proposed framework is reportedly efficient where it does not require storing the full Hessian matrix, allowing it to perform running with large neural networks.\n+ Providing more details of pruning with ViT models gives some new insights.", "weaknesses": "**Concerns**\n+ The baselines seem to be out-of-date, it would be helpful to specify additional recent baselines from 2024. More baselines should be considered for example [1,2].\n\n+ Considering specific aspects of the method that could be condensed or moved to an appendix, and particular experimental analyses or comparisons should be expanded:  Accuracy and latency speed-up with different network architectures: Res50, MobileNet, ViT.\n\n+ Section 5, \"Experimental setup\" is written too long, it contains also the results and a description of results that is not \"setup\". It should be separated into a new paragraph or subsection. More breaks between paragraphs should be added for better readability.\n\n+ Considering particular efficiency metrics the authors could report (e.g., runtime, memory usage) and specific methods (in Table 3) they should compare against to demonstrate their efficiency claims.\n\n**Other suggestions**\n+ Mentioning that experiments on ImageNet-1k (lines 102, 444)  might be misleading since the fact that the authors only use the Mini-ImageNet with 10,000 samples compared to ImageNet-1k (1.2M samples). Replacing ImageNet-1k with MiniImagenet-1k is more appropriate.\n\n+ There are two \"Experimental setup\" parts which are the same level as the network part (CNN running, Vision transformers, ...), which might cause confusion, suggesting making it a subsection for each type of network.\n\n+ The so-called \"Figure 3\" (Line 472-484) should be a \"Table 3\". The text referred to that table should be revised accordingly.\n\n+ For better readability, I suggest rather \"minimize (7)\" --> \"minimize Eq. 7\", Taylor expansion in (7) --> \"Taylor expansion in Eq. 7\", each call to (5) --> \"each call to Eq. 5\", etc.\n\n+ Equations are written in a separate line with a center and should be enumerated, for example for line 214, line 382, and 384.\n\n+ Although it is lookable, Figure 4 should make the text inside the figure larger for better readability.\n\n[1] FALCON: FLOP-Aware Combinatorial Optimization for Neural Network Pruning, AISTATS 2024\n\n[2] LayerMerge: Neural Network Depth Compression through Layer Pruning and Merging, ICML 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)