SequentialAttention++ for Block Sparsification: Differentiable Pruning Meets Combinatorial Optimization

Authors: taisuke  yasuda
mohammadhossein  bateni
vahab  mirrokni
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:22:54

Abstract: Neural network pruning is a key technique towards engineering large yet scalable, interpretable, and generalizable models. Prior work on the subject has developed largely along two orthogonal directions: (1) differentiable pruning for efficiently and accurately scoring the importance of parameters, and (2) combinatorial optimization for efficiently searching over the space of sparse models. We unite the two approaches, both theoretically and empirically, to produce a coherent framework for structured neural network pruning in which differentiable pruning guides combinatorial optimization algorithms to select the most important sparse set of parameters. Theoretically, we show how many existing differentiable pruning techniques can be understood as nonconvex regularization for group sparse optimization, and prove that for a wide class of nonconvex regularizers, the global optimum is unique, group-sparse, and provably yields an approximate solution to a sparse convex optimization problem. The resulting algorithm that we propose, SequentialAttention++, advances the state of the art in large-scale neural network block-wise pruning tasks on the ImageNet and Criteo datasets.

Main paper: 
 introduction: 
 Pruning methods for neural networks [LeCun et al., 1989 ] replace dense weight matrices by sparse approximations, which offer improved generalization and inference efficiency in terms of storage, energy consumption, and other computational resources. In various common formulations, the problem of computing the best sparse approximation to a dense weight matrix is intractable as it generalizes the sparse linear regression problem, which is known to be NP-hard even to approximate [Natarajan, 1995 , Foster et al., 2015 , Gupte and Vaikuntanathan, 2021 , Price et al., 2022] . Despite this fact, a wide variety of techniques have proven to be quite successful in practice. This includes magnitude pruning, ℓ 1 regularization, greedy coordinate descent, sampling, among others.While earlier works have focused on unstructured (i.e., entrywise) sparsity, which has been an active and fruitful area, researchers have rapidly recognized the importance of structured sparsity, which enforces that the sparse approximation respects certain patterns, such as block structure. These structural constraints often lead to further efficiency gains due to improved hardware utilization [Anwar et al., 2017 , Pool and Yu, 2021 , Liu et al., 2022] . Our work thus focuses on developing new and improved techniques for structured sparsification of weight matrices, and in particular on block sparsification [Ma et al., 2023] , which allow for a balance between performance gains from hardware utilization and reduced computation due to sparsity [Gale et al., 2023] .We argue that existing approaches to neural network pruning have developed along two orthogonal directions: algorithms for importance scoring and algorithms for combinatorial optimization. We roughly think of importance scoring algorithms as those that aim to select a small number of important entries (or blocks) of weight matrices, while we think of combinatorial optimization algorithms as wrapper methods that use the importance scoring algorithms as oracles to iteratively construct the desired (block) sparse weight matrix.Among importance scoring algorithms, early popular choices have included magnitude pruning [Thimm and Fiesler, 1995, Han et al., 2015] , where the magnitude of each trainable parameter serves as a proxy for its importance, as well as methods based on gradients [Karnin, 1990 , Sanh et al., 2020] , Hessians [LeCun et al., 1989 , Hassibi et al., 1993 , Singh and Alistarh, 2020 , Frantar and Alistarh, 2023] , and other statistics of the weights. Other works have incorporated ℓ 1 regularization [Wen et al., 2016 , Yang et al., 2019] to encourage sparsity. More recently, a class of techniques broadly termed differentiable pruning inspired by techniques for differentiable neural architecture search [Liu et al., 2019] have increased in popularity, where importance scores and/or soft masks are trained together with the network weights in a differentiable manner [Xiao et al., 2019 , Voita et al., 2019 , Kang and Han, 2020 , Ramakrishnan et al., 2020 , Savarese et al., 2020 , Zhang et al., 2022] . Variations of this idea use the network weights themselves to represent the "importance scores", and simply use a transformation of the original network weights [Schwarz et al., 2021 , Vanderschueren and Vleeschouwer, 2023 , Cho et al., 2023] .As for the combinatorial optimization aspects of pruning, the use of iterative or greedy procedures has long been explored and is known to improve sparsification quality over "one-shot" uses of importance scoring algorithms [LeCun et al., 1989 , Hassibi et al., 1993 , Ström, 1997 , Frankle and Carbin, 2019] . The work of Halabi et al. [2022] gives a theoretical justification of this observation via connections to weakly submodular optimization. Combinatorial optimization algorithms beyond greedy approaches, especially local search methods that improve sparsity patterns via local swaps such as iterative hard thresholding (IHT), have long been known in the submodular optimization literature, and have recently been shown to be extremely effective when combined with magnitude pruning [Evci et al., 2020 , Peste et al., 2021 , Kuznedelev et al., 2023b , Benbaki et al., 2023] . The work of Peste et al. [2021] also provides strong theoretical guarantees for their approach, ACDC. Similar ideas have also been termed as "neuroregeneration" in work of Liu et al. [2021] .Given these two highly fruitful approaches to the problem of pruning neural networks, it is natural to ask how recent advances in importance scoring algorithms and combinatorial optimization algorithms can work in concert. We investigate this question from both theoretical and empirical perspectives.We first present a theoretical investigation of differentiable pruning techniques for block sparsification when the objective function ℒ : R 𝑛 → R is strictly convex and differentiable. This already captures several interesting problems where block sparsification of weight matrices is desired, such as multinomial logistic regression and multiple response linear regression. We take the 𝑛 variables of our objective function to be partitioned into disjoint groups {𝑇 𝑖 } 𝑡 𝑖=1 where 𝑇 𝑖 ⊆ [𝑛] and possibly have varying size. For instance, in the context of block sparsification, ℒ could correspond to the multinomial logistic regression objective function with 𝐾 classes and 𝑑 features, and the 𝑛 = 𝐾𝑑 variables could be partitioned into 𝑡 blocks 𝑇 1 , 𝑇 2 , . . . , 𝑇 𝑡 . Furthermore, we will also consider an ℓ 2 regularization term on the parameters 𝛽, that is, we study variants of the problem min 𝛽∈R 𝑛 ℒ(𝛽) + 𝜆‖𝛽‖ 2 2 . Note that explicit ℓ 2 regularization is a standard component of machine learning architectures, and also appears implicitly whenever a loss function is optimized with gradient descent [Shalev-Shwartz, 2012] , with the regularization parameter 𝜆 being controlled by learning rate parameters and early stopping [Suggala et al., 2018] .Our contributions are twofold: (1) we show that a wide variety of differentiable pruning techniques can all be understood as an implementation of nonconvex regularization that generalizes the group LASSO, and (2) we show that a wide class of nonconvex regularizers give a unique 1-sparse global minimum that coincides with the unique 1-sparse global minimum of a corresponding group LASSO problem. These two results together establish that many differentiable pruning techniques work simply by identifying the same 1-sparse solution as the group LASSO. In turn, it is known that the 1-sparse solution found by the group LASSO is the variable block with the largest squared gradient [Axiotis and Yasuda, 2023] , which is equivalent to the Orthogonal Matching Pursuit [Pati et al., 1993 , Shalev-Shwartz et al., 2010 , Liberty and Sviridenko, 2017 , Elenberg et al., 2018] when applied sequentially (see Appendix B). Thus together, these results make progress towards understanding the inner workings of modern differentiable pruning methods.For our first contribution, we observe that if we minimize the loss ℒ with each of the variable groups 𝛽| 𝑇𝑖 for 𝑖 ∈ [𝑡] replaced by a "masked" version 𝑞(w 𝑖 )𝛽| 𝑇𝑖 , and with regularization on w and 𝛽, then this problem is equivalent to another optimization problem that simply optimizes ℒ with a different, and often sparsity-inducing, regularizer. A basic version of this observation already appears in works of Hoff [2017] , Axiotis and Yasuda [2023] , where it is shown that if the masks 𝑞 are just the identity, then we recover the usual group LASSO problem, that is,min w∈R 𝑡 ,𝛽∈R 𝑛 ℒ({w 𝑖 𝛽| 𝑇𝑖 } 𝑡 𝑖=1 ) + 𝜆 2 (︀ ‖w‖ 2 2 + ‖𝛽‖ 2 2 )︀ = min 𝛽∈R 𝑛 ℒ(𝛽) + 𝜆 𝑡 ∑︁ 𝑖=1 ‖𝛽| 𝑇𝑖 ‖ 2where {w 𝑖 𝛽| 𝑇𝑖 } 𝑡 𝑖=1 denotes the concatenation of the "masked" groups w 𝑖 𝛽| 𝑇𝑖 for 𝑖 ∈ [𝑡]. We generalize this observation and show that this framework also applies to other ideas popular in the differentiable pruning literature, such as applying ℓ 1 regularization on the masks w to induce sparsity [Yang et al., 2019] or applying softmax-type masks such as exp(w 𝑖 ) [Yasuda et al., 2023] . We note that prior to our work, there was little theoretical understanding on the value of applying such techniques in the context of differentiable pruning.We also apply similar ideas to differentiable pruning techniques that use the network weights themselves as importance scores [Schwarz et al., 2021 , Cho et al., 2023] . Here, the basic observation is that if one optimizes a loss function ℒ with variables 𝛽 replaced by the (signed) entrywise square 𝛽 ⊙ 𝛽, then this results in a "rich get richer" dynamic where large weights evolve to be larger while smaller weights are driven to zero, resulting in sparse solutions. This idea also has connections to exponentiated gradient descent which also results in sparse solutions [Vaskevicius et al., 2019, Amid and Warmuth, 2020a,b] . However, prior work only handles entrywise sparsity and does not address the question of structured pruning. We show that these ideas can also be understood in the framework of sparsity-inducing regularizers, even in the group setting. Here, we show that "masking" each of the variable groups 𝛽| 𝑇𝑖 by its ℓ 2 norm ‖𝛽| 𝑇𝑖 ‖ 2 gives a natural group generalization of this technique, and that this gives an optimization problem that is again equivalent to the group LASSO.Our second set of contributions is to analyze the solutions of a wide class of nonconvex regularizers. We now consider the following regularized problem, where 𝑞 : R + → R + is a strictly increasing and subadditive function with 𝑞(0) = 0, and 𝜆 > 0 is a regularization parameter:EQUATION)For instance, some popular choices of 𝑞 include the absolute value 𝑞(𝑥) = |𝑥|, 𝑝-th powers 𝑞(𝑥) = |𝑥| 𝑝 for 𝑝 < 1, or logarithmic regularizers such as 𝑞(𝑥) = log(1 + 𝑥). In general, the class of such 𝑞 (strictly) contains the set of all concave functions 𝑞 that vanish at the origin. Note that the form of (1) slightly differs from the usual form of nonconvex regularizers, as it applies 𝑞 -1 on the sum ∑︀ 𝑡 𝑖=1 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ) rather than taking the regularizer to just be ∑︀ 𝑡 𝑖=1 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ). This does not substantially change the nature of the optimization problem as it is the Lagrangian dual for the same constraint. The main result of this section is Theorem 1.1, which relates the group 𝑞-regularized min 𝛽∈R 𝑛 ℒ(𝛽) + 𝜆 𝑡 ∑︁ 𝑖=1 ‖𝛽| 𝑇𝑖 ‖ 2 .(2)Theorem 1.1 (Unique sparse global minima). Let 𝑞 : R + → R + be strictly increasing, subadditive (i.e., 𝑞(𝑎 + 𝑏) ≤ 𝑞(𝑎) + 𝑞(𝑏) for 𝑎, 𝑏 ∈ R + ), and satisfy 𝑞(0) = 0. If (2) has a unique minimizer 𝛽 * with group sparsity at most 1, then 𝛽 * is also the unique minimizer for (1).We make several remarks about Theorem 1.1. First, we justify why the assumption of the theorem is not vacuous: that is, we explain why the group LASSO objective (2) has sparse solutions. In recent work of Axiotis and Yasuda [2023] the following are shown if ℒ is strongly convex and differentiable:• If 𝜆 ≥ 𝜏 for 𝜏 = max 𝑡 𝑖=1 ‖∇ℒ(0)| 𝑇𝑖 ‖ 2 , then(2) has a unique global minimizer at 𝛽 = 0. • If 𝜆 < 𝜏 is sufficiently close to 𝜏 , then (1) has a unique 1-sparse global minimizer.Thus, when 𝜆 is large enough, Theorem 1.1 establishes that (1) has a unique sparse global minimum. Furthermore, Axiotis and Yasuda [2023] also show that the above global minimizer of the group LASSO problem (2) with group sparsity 1 is supported on a group 𝑇 𝑖 that maximizes ‖∇ℒ(0)| 𝑇𝑖 ‖ 2 , that is, it selects the group of variables that locally provides the largest improvement in the objective function cost. Repeatedly alternating between selecting such a feature and re-optimizing over the support is an algorithm known as the group orthogonal matching pursuit (group OMP), and has provable guarantees for group sparse convex optimization when ℒ satisfies the restricted strong convexity and restricted smoothness properties [Axiotis and Yasuda, 2023] . It is also shown that a related local search algorithm known as group orthogonal matching pursuit with replacement (group OMPR) also applies in this context, which has improved guarantees.Finally, we emphasize that it is generally difficult to establish structural results for nonconvex optimization problems, even for simple convex problems with nonconvex regularizers. Thus, we believe that our results may be of independent interest in the literature of nonconvex optimization.We now apply our theoretical insights of combining differentiable pruning and combinatorial optimization to develop a novel algorithm for block neural network pruning, which we call Sequen-tialAttention++. SequentialAttention++ is primarily a fusion of two prior techniques: Sequential Attention, a feature selection technique based on differentiable pruning developed in work of Yasuda et al. [2023] , and ACDC, which is a highly effective stochastic adaptation of the classic iterative hard thresholding (IHT) algorithm [Peste et al., 2021] from the combinatorial optimization literature.Sequential Attention [Yasuda et al., 2023] is an algorithm for feature selection on neural networks, that introduces a softmax mask that is trained together with the neural network weights. Each of the 𝑛 input features is scaled by a differentiable mask 𝐴 𝑖 = exp(𝐿 𝑖 )/ ∑︀ 𝑛 𝑗=1 exp(𝐿 𝑗 ) for a vector 𝐿 ∈ R 𝑛 of logits. Note that our theoretical results on differentiable pruning, and in particular Lemma 2.1, suggests that this roughly corresponds to performing a log-sum regularization on the corresponding weights for these features. We first extend this to the block sparsification setting by instead scaling each block of weights to prune by a similar softmax mask (see Figure 1 ). Note that in this new setting, Lemma 2.1 shows that this corresponds to a group log-sum regularization on each of the blocks.We then use this differentiable pruning technique as part of a local search procedure inspired by ACDC [Peste et al., 2021] . In the originally proposed ACDC algorithm, the neural network is trained in multiple phases, where the phases alternate between a "dense" training phase and a "sparse" training phase. During the dense phases, the weights are trained in the standard way, whereas in the sparse phases, only a sparse set of weights corresponding to the top 𝑘 weights at the beginning of the phase (i.e., chosen by magnitude pruning) are used. The idea here is that if a suboptimal sparse support is selected during the sparse phase, then this support can be modified during the dense phase. We note that one of the weaknesses of this algorithm is the use of the weight magnitudes as a proxy for the importance of the weights, whereas improved parameter importance estimation is possible by introducing differentiable pruning techniques. Thus in our SequentialAttention++ algorithm, we modify the ACDC algorithm by training a softmax mask together with the neural network weights during the dense phase as in Figure 1 , and then using the softmax mask to select a sparse support during the sparse phases. Our theoretical results establish provable guarantees for a slightly modified version of this algorithm, by showing that log-sum regularization can be integrated with a similar local search algorithm that alternates between dropping small weights from the support, selecting weights via regularization, and optimizing on the new support (see Theorem B.3 and Appendix B). 
theory: 
 In Section 2, we present our theoretical results on differentiable pruning and local search algorithms for DNN sparsification. Missing proofs can be found in Appendix A.In this section, we show how a wide variety of differentiable pruning techniques studied in the literature can be viewed as nonconvex regularizers. As described earlier in Section 1.2.2, we later show that nonconvex regularization can in fact be connected to provable guarantees for sparse convex optimization by implementing the orthogonal matching pursuit algorithm and its variants. Thus, together, we give the first steps towards a full theoretical analysis of many popular differentiable pruning techniques in the literature.The softmax is a popular differentiable sparsity-inducing technique, where a vector is transformed by exponentiating each entry and normalizing the result. The softmax forms the backbone of many modern ML techniques ranging from multinomial logistic regression to differentiable architecture search [Liu et al., 2019] to attention mechanisms and transformers [Vaswani et al., 2017] , and thus a theoretical understanding of the softmax is critical mission for modern machine learning theory.We take a step towards this by considering unnormalized softmax, which corresponds to a simple entrywise exponentiation. The unnormalized softmax is a popular alternative to the usual softmax as it still captures its sparsity-inducing properties [Amid and Warmuth, 2020a,b] , while its simplicity allows for more efficient implementations. We show that, in fact, unnormalized softmax can be viewed as a type of log-sum regularization, which is a popular relaxation of the ‖•‖ 0 norm that has been often considered in the machine learning and signal processing literatures [Rao and Kreutz-Delgado, 1999 , Wipf and Nagarajan, 2009 , Qiao et al., 2020 , Tugnait, 2022 , Zhou et al., 2023] .Lemma 2.1 (Unnormalized softmax as log-sum regularization).min w∈R 𝑡 ,𝛽∈R 𝑛 ℒ({exp(w 𝑖 )𝛽| 𝑇𝑖 } 𝑡 𝑖=1 ) + 𝜆 (︀ ‖w‖ 2 2 + ‖𝛽‖ 2 2 )︀ = min u∈R 𝑛 ℒ(u) + 𝜆 𝑡 ∑︁ 𝑖=1 𝑞(‖u| 𝑇𝑖 ‖ 2 )where 𝑞(𝑎) = 𝑊 (2𝑎 2 ) 2 /4 + 𝑊 (2𝑎 2 )/2 and 𝑊 is the Lambert 𝑊 function, i.e., the inverse of 𝑓 (𝑊 ) = 𝑊 𝑒 𝑊 .Next, we consider the idea of applying a sparsity-inducing regularization on a mask (see, e.g., the work of Yang et al. [2019] ). We show that by regularizing the mask instead of the parameters themselves, the resulting optimization leads to a "more nonconvex" regularizer.Lemma 2.2 (ℓ 1 -regularized masks as ℓ 𝑞 regularization).min w∈R 𝑡 ,𝛽∈R 𝑛 ℒ({w 𝑖 𝛽| 𝑇𝑖 } 𝑡 𝑖=1 ) + 𝜆 (︀ ‖w‖ 1 + ‖𝛽‖ 2 2 )︀ = min u∈R 𝑛 ℒ(u) + 3 2 2 1/3 𝜆 𝑡 ∑︁ 𝑖=1 ‖u| 𝑇𝑖 ‖ 2/3 2 2.1.3 PowerpropagationFinally, we study differentiable pruning techniques that use the network weights themselves as importance scores. The most straightforward implementation of this idea is to square each of the weights, as explored in works such as powerpropagation for neural networks [Schwarz et al., 2021] , but more complex versions have also been considered [Cho et al., 2023] . We show how these techniques can be generalized to handle the group setting, and show how they can also be interpreted as an implementation of group sparsity-inducing regularization. Lemma 2.3 (Group powerpropagation as Group LASSO).min w∈R 𝑡 ,𝛽∈R 𝑛 ℒ({‖𝛽| 𝑇𝑖 ‖ 2 𝛽| 𝑇𝑖 } 𝑡 𝑖=1 ) + 𝜆‖𝛽‖ 2 2 = min u∈R 𝑛 ℒ(u) + 𝜆 𝑡 ∑︁ 𝑖=1 ‖u| 𝑇𝑖 ‖ 2We will prove the following theorem in this section, which establishes natural conditions for which nonconvex regularization of a convex function produces a unique group-sparse global minimum.As discussed in Section 1.2.2, this theorem is the main crucial result for proving that local search algorithms give provable guarantees for sparse convex optimization. Theorem 1.1 (Unique sparse global minima). Let 𝑞 : R + → R + be strictly increasing, subadditive (i.e., 𝑞(𝑎 + 𝑏) ≤ 𝑞(𝑎) + 𝑞(𝑏) for 𝑎, 𝑏 ∈ R + ), and satisfy 𝑞(0) = 0. If (2) has a unique minimizer 𝛽 * with group sparsity at most 1, then 𝛽 * is also the unique minimizer for (1).We have the following lemma that shows that if 𝑞 is strictly increasing and subadditive, then the group 𝑞-regularization is always larger than group LASSO regularization. Thus, the group LASSO objective is always a lower bound on the 𝑞-regularized objective. Lemma 2.4. Let 𝑞 : R + → R + be strictly increasing and subadditive. Then,𝑡 ∑︁ 𝑖=1 ‖𝛽| 𝑇𝑖 ‖ 2 ≤ 𝑞 -1 (︃ 𝑡 ∑︁ 𝑖=1 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ) )︃Proof. Since 𝑞 is invertible, applying the subadditivity condition on 𝑞( ∑︀ 𝑡 𝑖=1 ‖𝛽| 𝑇𝑖 ‖ 2 ) and then applying 𝑞 -1 on both sides of the inequality yields the result. Furthermore, note that for solutions 𝛽 that have group sparsity at most 1, the group 𝑞-regularization has the same value as the group LASSO regularization. That is, the lower bound value can be achieved when the group sparsity is at most 1. Lemma 2.5. Let 𝑞 : R + → R + be strictly increasing and satisfy 𝑞(0) = 0. Then, for any 𝛽 ∈ R 𝑛 with group sparsity 1,𝑡 ∑︁ 𝑖=1 ‖𝛽| 𝑇𝑖 ‖ 2 = 𝑞 -1 (︃ 𝑡 ∑︁ 𝑖=1 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ))︃ .Proof. If 𝛽 has group sparsity at most 1, say supported on 𝑇 𝑗 for some 𝑗 ∈ [𝑡], then we have𝑞 -1 (︃ 𝑡 ∑︁ 𝑖=1 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ) )︃ = 𝑞 -1 (︀ 𝑞(‖𝛽| 𝑇𝑗 ‖ 2 ) )︀ = ‖𝛽| 𝑇𝑗 ‖ 2 .Together, Lemmas 2.4 and 2.5 imply that if the group LASSO objective has a unique sparse minimum, then this is a lower bound on the optimal value that can be achieved by the 𝑞-regularized objective. This proves Theorem 1.1. The formal argument can be found in Appendix A. 
the sequentialattention++ algorithm: 
 Weight magnitude is a simple and reliable importance score used to prune candidates (in our case, blocks) in a sparse optimization problem. In many cases, however, the magnitudes do not correlate very well with the true importances of the candidates. This has been observed e.g. in Axiotis and Sviridenko [2021, 2022] , who showed that the magnitude pruning criterion used in the IHT algorithm is provably suboptimal even for simple sparse regression tasks, and proposed an adaptive weight decay to deal with this issue. One reason for the suboptimality of magnitude pruning is that the weights are not encouraged to be sparse during model training, leading to redundancy. Methods such as Powerpropagation [Schwarz et al., 2021] and Sequential Attention [Yasuda et al., 2023] have been proposed to address this issue by explicitly encoding a non-convexity that encourages weights to be concentrated on a sparse subset (this can be viewed as weight re-parameterization or concave regularization, as shown in Section 2).To test the hypothesis that softmax attention weights are higher-quality importance scores, we consider one-shot block pruning based on the softmax attention scores used in Sequential Attention (see Figure 1 on how to apply it to blocks), and we compare it with block magnitude (Frobenius norm) pruning. The results in Figure 2a suggest that softmax attention scores are generally more reliable as block importance scores, especially for larger block sizes. This leads us to adopt the softmax parameterization in our algorithm.As observed e.g. in Peste et al. [2021] , one-shot pruning approaches are significantly suboptimal compared to iterative pruning approaches such as ACDC. We use a similar alternating compressed and decompressed phases approach as ACDC, but we apply it on the softmax attention weights instead of the block magnitudes. This establishes SequentialAttention++ as a combination between Sequential Attention and ACDC. The basic algorithm can be seen in Algorithm 2.Algorithm 1 Feed-forward layer with the basic version of SequentialAttention++ to select top 𝑘 parameters from a kernel W.function FF(X ∈ R 𝑏×𝑛 : input batch, 𝑡 : training step) Trainable params: One drawback of sparse/dense (compression/decompression) phases is that the dense-to-sparse transition is abrupt. Since the lowest-magnitude weights are instantly pruned, this neglects correlations between these pruned parameters. If we were to re-train the model after pruning one parameter at a time, the picture could be drastically different, since low-magnitude weights could grow (this could happen e.g. due to parameter redundancy). In fact, this effect was highlighted by Kuznedelev et al. [2023a] , who devised a backward selection method based on correlations as captured by the Hessian.Kernel W ∈ R 𝑛×𝑚 , Logits L ∈ R 𝑛×𝑚 A = 𝑛𝑚 • 𝑒 L / ∑︀ 𝑒 L Ŵ = W ⊙ A ⊙ Mask(A,Inspired by this approach, we incorporate a backward selection phase between the DENSE and SPARSE phases, which we call the SPARSIFICATION phase. In this phase, we gradually prune the least important features based on the attention weights. This gradual process allows the model to re-adjust the attention weights after some parameters are pruned. The importance of this phase is validated by ablation experiments in Appendix D.1. We use an exponential pruning schedule, to prune more aggressively in the beginning of the phase, and more carefully at the end (as we approach the desired number of candidates 𝑘). A comparison of the sparsity schedules of ACDC and SequentialAttention++ can be found in Figure 2b . We use the sparsity schedule sparsity(𝑡) = 𝑠 • 1-𝑒 where 𝑠 is the target sparsity. This interpolates between sparsity 0 and 𝑠, and constitutes a single SPARSIFICATION phase. We choose the constant 𝑐 = 4 (for an ablation analysis, see Appendix D.2). 
experiments: 
 We evaluate our algorithms on sparsification tasks where a dense DNN is approximated by blocksparse counterparts, at various block sizes 𝐵 and sparsities 𝑝, where a sparsity 𝑝 indicates that the DNN layer will only have a 1 -𝑝 fraction of nonzero entries, and a block size of 𝐵 indicates that the nonzero entries are arranged in 𝐵 × 𝐵 blocks. Note that for a fixed sparsity, larger block sizes generally translate to improved efficiency due to improved hardware utilization, but also degrades quality. Block size of 1 corresponds to unstructured pruning. Our experiments are performed on the ImageNet and Criteo datasets. More details on the setup can be found in Section C.1.We compare our SequentialAttention++ algorithm to three other representative prior algorithms for DNN pruning. The first is basic magnitude pruning, which is a popular and effective algorithm where the weights are sparsified by keeping the weights with the largest magnitude after training [Frankle and Carbin, 2019] . We use it in the block setting by keeping the largest blocks in Frobenius norm.The second algorithm is a block generalization of Powerpropagation [Schwarz et al., 2021] , which combines magnitude pruning with a differentiable pruning technique where sparsity is encouraged by squaring the weights. While the original Powerpropagation algorithm did not handle the block sparsification setting, we show that multiplying each block by the Frobenius norm leads to a provable generalization (see Lemma 2.3). Finally, we consider ACDC [Peste et al., 2021] , which is an adaptation of iterative hard thresholding (IHT) [Blumensath and Davies, 2009] to the setting of neural network sparsification, and has produced the state-of-the-art pruning results for ImageNet [Kuznedelev et al., 2023b] . For all algorithms and datasets, we include a fine-tuning phase at the end of training, using the pruned model, and evaluate the final pruned model on the test set.Our results on ImageNet are summarized in Table 1 . The sparsities range over 58-95% and the block sizes over 8, 16, 32, 64. We compare ACDC and SequentialAttention++. Our ACDC implementation closely follows the implementation in Peste et al. [2021] 3 . We use the phase schedule suggested by Kuznedelev et al. [2023b] (10% dense, 7 equal SPARSE-DENSE phases where the last dense phase is extended by 5%, 15% sparse). For SequentialAttention++, we additionally replace each sparse-dense phase by a SPARSIFICATION-SPARSE-DENSE phase, as described in Section 3.1, and we replace the last of the 7 phases (including its extension) by a SPARSIFICATION phase. We use a batch size of 2048 and a maximum learning rate of 0.8.We observe that SequentialAttention++ generally outperforms ACDC on the block sparsification task, across all different block sizes and sparsities that we tested. It should be mentioned that this comes at the cost of introducing additional trainable parameters to the model (one parameter per block). This overhead could be concerning in some applications if block size is too small (e.g., 1), in which case the model's parameters are being doubled. However, the overhead is negligible for larger (e.g., ≥ 8) block sizes.Our results on the Criteo dataset are presented in Table 2 . The sparsities range over 𝑝 ∈ {90%, 95%, 97%, 98%, 99%} and block sizes over 𝐵 ∈ {5, 10, 20}. In this experiment, we used a schedule of 10 sparse-dense phases, in addition to a 20% initial dense phase and a final 20% sparse phase. Note that for this experiment, we used masking instead of pruning for ACDC, meaning that unselected blocks are not pruned but multiplied with an all-zero mask. We observe that SequentialAt-tention++ is the best performing algorithm. In fact, we notice that the gap widens with large block sizes and high sparsity, suggesting that SequentialAttention++ is a highly accurate block sparsification algorithm for large block sizes and extreme sparsities. 
conclusion: 
 In this work, we unified, generalized, and improved prior approaches to neural network pruning via a framework which combines differentiable pruning with combinatorial optimization algorithms, in particular local search techniques. Theoretically, we gave a unified analysis of a wide class of existing techniques via a connection to nonconvex regularization, and proved novel properties about sparse convex optimization with nonconvex regularization. In particular, we established natural conditions under which nonconvex regularization yields a unique group-sparse global minimum that is supported on the group that maximizes the ℓ 2 norm of the gradient, thus yielding provable guarantees for group sparse convex optimization. Empirically, we proposed a novel algorithm, SequentialAttention++, which outperforms prior methods on standard benchmark datasets for neural network sparsification.We conclude with a few open directions which we believe to be interesting for future work. The first is on characterizing the nature of critical points and local minima of nonconvex-regularized convex problems. This would be a more practically useful variation on our result, which only establishes provable guarantees for the global minimizer. For our second question, we ask whether one can theoretically establish that nonconvex regularization yields better optimization guarantees than the LASSO. In our work, we have only shown that the quality of solutions found by nonconvex regularization can match the LASSO for a wide variety of nonconvex regularizers, but we do not theoretically establish that this formulation is better. It would be interesting to show, e.g., that nonconvex regularizers allow for faster convergence to the sparse global minimizer. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: [TODO]Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: [TODO]Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA]Justification: [TODO]Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: [TODO] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Proof of Lemma 2.1. Note first that for a fixed 𝑎 > 0, the function𝑤 ↦ → 𝑤 2 + 𝑎 2 / exp(2𝑤) is minimized at 𝑤 satisfying 2𝑤 -2𝑎 2 exp(-2𝑤) = 0, that is, 𝑤 = 𝑊 (2𝑎 2 )/2. Then, for each group 𝑖 ∈ [𝑡], we can set u| 𝑇𝑖 = exp(w 𝑖 )𝛽| 𝑇𝑖 so w 2 𝑖 + ‖𝛽| 𝑇𝑖 ‖ 2 2 = w 2 𝑖 + ‖u| 𝑇𝑖 ‖ 2 2 exp(2w 𝑖 ) ≥ 𝑤 2 + 𝑤where 𝑤 = 𝑊 (2‖u| 𝑇𝑖 ‖ 2 2 )/2. Summing over the groups 𝑖 ∈ [𝑡] gives the desired result.Proof of Lemma 2.2. Note first that for a fixed 𝑎 > 0, the function 𝑤 ↦ → 𝑤 + 𝑎 2 /𝑤 2 is minimized at 𝑤 satisfying 1 -2𝑎 2 𝑤 -3 = 0, that is, 𝑤 = 2 1/3 𝑎 2/3 . Then, for each group 𝑖 ∈ [𝑡], we can set u| 𝑇𝑖 = w 𝑖 𝛽| 𝑇𝑖 so|w 𝑖 | + ‖𝛽| 𝑇𝑖 ‖ 2 2 = |w 𝑖 | + ‖u| 𝑇𝑖 ‖ 2 2 w 2 𝑖 ≥ 3 2 𝑤 where 𝑤 = 2 1/3 ‖u| 𝑇𝑖 ‖ 2/32 . Summing over the groups 𝑖 ∈ [𝑡] gives the desired result.Proof of Lemma 2.3. Set u| 𝑇𝑖 = ‖𝛽| 𝑇𝑖 ‖ 2 𝛽| 𝑇𝑖 . Then,‖u| 𝑇𝑖 ‖ 2 = ‖‖𝛽| 𝑇𝑖 ‖ 2 𝛽| 𝑇𝑖 ‖ 2 = ‖𝛽| 𝑇𝑖 ‖ 2 2so summing over the groups gives the claimed result.Proof of Theorem 1.1. Suppose that the optimal group LASSO solution 𝛽 * of objective (2) has group sparsity at most 1. Then for any other solution 𝛽 ′ , we have thatℒ(𝛽 ′ ) + 𝜆𝑞 -1 (︃ 𝑡 ∑︁ 𝑖=1 𝑞(‖𝛽 ′ | 𝑇𝑖 ‖ 2 ) )︃ ≥ ℒ(𝛽 ′ ) + 𝜆 𝑡 ∑︁ 𝑖=1 ‖𝛽 ′ | 𝑇𝑖 ‖ 2 by Lemma 2.4 > ℒ(𝛽 * ) + 𝜆 𝑡 ∑︁ 𝑖=1 ‖𝛽 * | 𝑇𝑖 ‖ 2 by optimality = ℒ(𝛽 * ) + 𝜆𝑞 -1 (︃ 𝑡 ∑︁ 𝑖=1 𝑞(‖𝛽 * | 𝑇𝑖 ‖ 2 ))︃ by Lemma 2.5.Thus, 𝛽 * must be the unique minimizer of (1).We show that our results from Section 2.2 together with recent work of Axiotis and Yasuda [2023] give provable guarantees for a local search algorithm based on orthogonal matching pursuit with replacement using nonconvex regularization.We first introduce some definitions needed to state our result. ℒ(𝛽 + Δ) -ℒ(𝛽) -⟨∇ℒ(𝛽), Δ⟩ ≤ 𝐿 𝑠 2 ‖Δ‖ 2 2 .We will now obtain provable guarantees for Algorithm 3 in Theorem B.3.Initialize 𝑆 arbitrarily such that |𝑆| = 𝑘 ′ for 𝑖 = 1, . . . , 𝑅 do Letβ = arg min 𝛽∈R 𝑛 ℒ(𝛽) + 𝜆 • 𝑞 -1 (︃ ∑︁ 𝑖 / ∈𝑆 𝑞(‖𝛽| 𝑇𝑖 ‖ 2 ) )︃for 𝜆 sufficiently large Let 𝑖 / ∈ 𝑆 be the group maximizing β| 𝑇𝑖 and 𝑗 ∈ 𝑆 be the group minimizing ‖𝛽‖ 2 | 𝑇𝑗 𝑆 ← 𝑆 ∪ {𝑖} ∖ {𝑗} end for Theorem B.3 (OMPR via nonconvex regularization). Let 𝑞 : R + → R + be strictly increasing, subadditive, and 0 at the origin. After 𝑅 iterations of Algorithm 3 with 𝑘 ′ ≥ 𝑘(︂ 𝐿 2 2 𝜇 2 𝑘+𝑘 ′ + 1 )︂ , for 𝑅 ≥ 𝑘 • 𝐿 2 𝜇 𝑘+𝑘 ′ log ℒ(𝛽 (0) ) -ℒ(𝛽 * ) 𝜀 ,then β has group sparsity ‖𝛽 ∞ ‖ group ≤ 𝑘 ′ and satisfies ℒ(𝛽 ∞ ) ≤ ℒ(𝛽 * ) + 𝜀 , where 𝜇 𝑘+𝑘 ′ is a lower bound on the restricted strong convexity constant of 𝑙 at group sparsity 𝑘 + 𝑘 ′ and 𝐿 2 is an upper bound on the restricted smoothness constant of 𝑙 at group sparsity 2 (see Definition B.2).Proof. By Theorem 1.1, if the optimization problem in Algorithm 3 with 𝑞 replaced by the absolute value function has a unique minimizer with group sparsity at most 1, then β is a unique global minimizer with group sparsity at most 1, and coincides with this Group LASSO solution. Lemma 3.2 of Axiotis and Yasuda [2023] then establishes that this solution is supported on the group that maximizes the ℓ 2 norm of the gradient, which in turn implies Theorem B.3 via guarantees for the group orthogonal matching pursuit with replacement algorithm (Corollary A.10 of Axiotis and Yasuda [2023] ). ImageNet [Deng et al., 2009] . ImageNet is the most widely used vision dataset and is considered as the de facto benchmark in the neural network pruning literature, culminating in the state of the art results in Kuznedelev et al. [2023b] . We use ResNet50 and a standard training setup (90 epochs, SGD with cosine learning rate and momentum, weight decay). We reshape the 4-dimensional (𝐻 × 𝑊 × 𝐶 in × 𝐶 out ) kernel tensors used in convolutional layers to 2D matrices of shape 𝐻𝑊 𝐶 in × 𝐶 out , which define the 2D block candidates for pruning. We prune all layers uniformly, except for layers with < 100 blocks, which we do not prune at all, to avoid degeneracy at high sparsities.Criteo [Diemert et al., 2017] . Criteo is a standard public dataset for the clickthrough rate (CTR) prediction task, which consists of 33M training examples with 13 numerical and 26 categorical features. The model we sparsify is a standard fully connected DNN with three 400-width layers and an additional embedding layer to transform each input feature into an embedding vector of size 10 (for a total embedding width of 390). We note that a simple MLP is often a fairly competitive model for this task [Naumov et al., 2019] . We prune the first dense layer after the embedding layer. We use Adam optimizer with a learning rate that decays exponentially from 2 • 10 -2 to 3 • 10 -4 . We train to minimize the cross-entropy loss for 25 epochs with a batch size of 32768.We give our block sparsification results on the Criteo dataset in In addition to the basic algorithm described in Section 3, our implementation of SequentialAtten-tion++ incorporates several other ingredients for improved empirical performance. First, we confirm the observation of Peste et al. [2021] that resetting the optimizer between each phase of SequentialAt-tention++ is crucial for good performance. We note that this is also suggested by our theoretical results (Theorem B.3), which suggests that each of the dense and sparse phases should be thought of as a separate optimization problem that is solved independently. Similarly to Kuznedelev et al. [2023b] , we also observe that weight decay significantly boosts performance, even when applied to the attention logits.Second, we observe that pruning each layer of the network separately performs better than a global pruning algorithm which attempts to prune all layers at once. We suggest that this may be the case due to "bottlenecking" behavior, where a global pruning algorithm may choose to almost completely eliminate a layer which may destroy the connectivity of the neural network. While this is not the case when pruning individual parameters, pruning large blocks can easily eliminate a layer. We use uniform sparsity across layers, but choose not to sparsify layers containing less than 100 blocks. This is because layers have greatly varying sizes, and want to avoid a sharp quality drop from overpruning smaller layers, which was observed in experiments. Finally, we clip attention weights to the range [𝑛 • density, 𝑛/density] to avoid them becoming too small or too large.We provide additional plots for our experiments in Figures 4 and 5 . In Figure 4 , we plot tradeoffs between the validation accuracy and weight matrix sparsity for SequentialAttention++ and ACDC Peste et al. [2021] . In Figure 5 , we plot tradeoffs between the validation loss and AUC against weight matrix sparsity for SequentialAttention++ and our three baseline algorithms of Magnitude Pruning, Powerpropagation Schwarz et al. [2021] , and ACDC Peste et al. [2021] . The results in Table 3 show that, on average over different block sizes and densities, removing the SPARSIFICATION phase decreases validation accuracy by 0.009, or 0.9 percentage points. We conclude that the SPARSIFICATION phase is an important feature of SequentialAttention++.D.2 Choice of the SPARSIFICATION exponent.In this section, we try different values of the constant used in the exponent of the schedule of the SPARSIFICATION operation. We remind that during a SPARSIFICATION phase, the sparsity varies as sparsity(𝑡) = 𝑠 • 1-𝑒 Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We sanity-checked our ACDC implementation by verifying that the accuracy of 90% unstructured global pruning matches that of the ACDC paper (75.01 vs 75.03). 


Decision: This work addresses structured neural network pruning by unifying differential pruning and combinatorial optimizations. A nonconvex regularization for group sparse optimization is shown to underpin many differentiable pruning methods, and its global optimum is analyzed.  A practical algorithm SequentialAttention++ is then proposed, whose effectiveness is demonstrated on block-wise pruning tasks on the ImageNet and Criteo.

This paper is well written.  It provides new insights on optimization and neural network pruning, and the empirical algorithm is effective.  I am personally a bit concerned that most theoretical results assume the group sparsity is at most 1.  However, overall, the merits of the paper make it a solid contribution to the conference.

Reviews: review: paper_summary: paper_summary: This work studies the task of neural network pruning. The authors unify the two main directions of the literature: differential pruning and combinatorial optimizations. Specifically, they point out that most differentiable pruning techniques can be considered as non-convex regularization group sparse optimization problems. Based on this theoretical analysis, they propose SequentialAttention++. The empirical evaluation shows their method reaching state-of-the-art results.
strength_weakness: strength_weakness: The writing is slightly hard to follow. Specifically, the exact motivation of SequentialAttention++ seems a bit unclear. Is it just a combination of Sequential Attention and ACDC because they are currently state-of-the-art methods? Or is there a stronger reason and theoretical support behind this?
questions: questions: Do you have more reasons behind the combination of Sequential Attention and ACDC? Also, can you explain why in most of the experiment results, larger block size seems to perform better? Wouldn't a smaller block size mean more flexibility in pruning?
limitations: limitations: Yes, the authors discuss the limitation of the work.
paper_summary: paper_summary: This paper analyzes sparse training with group sparsity and proposes a general theoretical framework to analyze both hard thresholding (discrete) and scoring (continuous) pruning methods. The proposed theoretical framework of block sparsification encompasses multiple existing sparsification methods via non convex regularizers and shows the existence of a unique global minimizer for such problems by establishing an equivalence to the group LASSO problem. Based on the theoretical insights, the proposed algorithm Sequntial Attention++ iteratively prunes parameters combining the approach of AC/DC with block sparsification and score based pruning.
strength_weakness: strength_weakness: 1. The authors discuss that the magnitude is not necessarily the best importance scoring metric for sparsification. However, experimentally it has been observed to be the best performing criteria across multiple datasets for hard thresholding methods with unstructured pruning (see [1]). I would like to know if the authors observe a different trend for (i) structured pruning and (ii) score based sparsification instead of hard thresholding i.e. does magnitude based scoring perform worse in case of differentiable pruning. 
2. In the proposed algorithm, the authors use a softmax based scoring method. Does using an unnormalized softmax lead to similar results?
3. If I understand correctly, the proposed results also hold for scoring methods which use the network weights as a score like in STR  [2] or like powerpropagation. However, does the additional overparametrization in case of using a separate scoring parameter improve the overall performance of the network?
4. For ImageNet it would be nice to see comparisons with structured DST [2] as it achieves a state of the art performance for structured sparsity.
5. Have the authors tried the same algorithm without an AC/DC approach, by simply training with a scoring method and pruning after every few epochs. Is the dense phase necessary?

[1] Hoefler, Torsten, et al. "Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks." Journal of Machine Learning Research (2021).

[2] Kusupati, Aditya, et al. "Soft threshold weight reparameterization for learnable sparsity." International Conference on Machine Learning (2020).

[3] Lasby, Mike, et al. "Dynamic Sparse Training with Structured Sparsity." The Twelfth International Conference on Learning Representations.
questions: questions: See above.
limitations: limitations: The paper sheds light on the connection between sparsification methods via nonconvex regularizers to the group LASSO problem, thus establishing a theoretical framework to analyze continuous sparsification methods, which is also empirically validated.
paper_summary: paper_summary: This paper focuses on bringing a theoretical understanding to differentiable pruning for neural networks and reveals connections to group lasso.
strength_weakness: strength_weakness: Please see the “Questions” section.
questions: questions: I think this paper provides some useful connections between optimization theory and pruning. I haven’t checked the math carefully, but things seem reasonable overall. It is interesting that the differential pruning techniques can be explained using group lasso. The empirical results are nice to have, but I view the theoretical connections made in the paper to be the main contribution.

1) Does the training curve in Figure 3 represent the entirety of the training? It looks like for the orange curve, the peak around step=48000 is higher than the accuracy at the end of the training. 

2) Also, what happens if the models are trained longer? Will the blue curve continue to do better than the orange curve?
limitations: limitations: Yes.

Comments: title: Response to rebuttal
comment: Thank you for the answers.
Thank you for the detailed response and for answering every question, substantiated with experiments. I appreciate the effort.
Thank you for answering my questions!