Compression-aware Training of Neural Networks using Frank-Wolfe

Many existing Neural Network pruning approaches either rely on retraining to compensate for pruning-caused performance degradation or they induce strong biases to converge to a specific sparse solution throughout training. A third paradigm obtains a wide range of compression ratios from a single dense training run while also avoiding retraining. Recent work of Pokutta et al. (2020) and Miao et al. (2022) suggests that the Stochastic Frank-Wolfe (SFW) algorithm is particularly suited for training state-of-the-art models that are robust to compression. We propose leveraging $k$-support norm ball constraints and demonstrate significant improvements over the results of Miao et al. (2022) in the case of unstructured pruning. We also extend these ideas to the structured pruning domain and propose novel approaches to both ensure robustness to the pruning of convolutional filters as well as to low-rank tensor decompositions of convolutional layers. In the latter case, our approach performs on-par with nuclear-norm regularization baselines while requiring only half of the computational resources. Our findings also indicate that the robustness of SFW-trained models largely depends on the gradient rescaling of the learning rate and we establish a theoretical foundation for that practice.


Introduction
The astonishing success of Deep Neural Networks relies heavily on over-parameterized architectures [72] containing up to several billions of parameters. Consequently, modern networks require large amounts of storage and increasingly long, computationally intensive training and inference times, entailing tremendous financial and environmental costs [59]. To address this, a large body of work focuses on compressing networks, resulting in sparse models that require only a fraction of memory or floating-point operations while being as performant as their dense counterparts. Recent techniques include the pruning of individual parameters [9,22,26,27,41,46] or group entities such as convolutional filters and entire neurons [2,44,48,68], the utilization of low-bit representations of networks (quantization) [34,63] as well as classical matrix-or tensor-decompositions [60,67,74] in order to reduce the number of parameters.
While there is evidence of pruning being beneficial for the ability of a model to generalize [9,30], a higher degree of sparsification will typically lead to a degradation in the predictive power of the network. To reduce this impact, two main paradigms have emerged. Pruning after training, most prominently exemplified by Iterative Magnitude Pruning (IMP) [26], forms a class of algorithms charaterized by a three-stage pipeline of regular (sparsity-agnostic) training followed by prune-retrain cycles that are either performed once (One-Shot) or iteratively. The need for retraining to recover pruning-induced losses is often considered to be an inherent disadvantage and computationally impractical [18,46,47,65]. In that vein, pruning during training or regularization approaches avoid retraining by inducing strong inductive biases to converge to a sparse model at the end of training [13,37,47,75]. The ultimate pruning then results in a negligible performance degradation, rendering the retraining procedure superfluous. However, such procedures incorporate the goal sparsity into training, requiring to completely train a model per sparsity level, while IMP needs just one pretrained model to generate the entire accuracy-vs.-sparsity frontier, albeit at the price of retraining.
A third paradigm, which is the focus of this work, naturally emerges when no retraining is allowed, but training several times to generate the accuracy-vs.-sparsity tradeoff frontier is prohibitive. Ideally, the optimization procedure is "compression-aware" [3] or "pruning-aware" [50], allowing to train once and then being able to compress One-Shot to various degrees while keeping most of the performance without retraining (termed pruning stability). Pruning-aware training procedures are expected to yield state-of-the-art dense models which are robust to pruning without its (regularization) hyperparameters being selected for a particular level of compression. An optimization approach that is particularly suited in that regard is the first-order and projection-free Stochastic Frank-Wolfe (SFW) algorithm [7,21,50,54,61]. While being valued throughout various domains of Machine Learning for its highly structured, sparsity-enhancing update directions [12,39,71], the algorithm has only recently been considered for promoting sparsity in Neural Network architectures.
Continuing the work of Pokutta et al. [54] and Miao et al. [50], the goal of this work is to thoroughly analyze the capabilities of SFW in the pruning-aware setting. We drastically improve upon the recent results of Miao et al. [50] by showing that the k-support norm [4] offers a significantly more natural and effective constraint on the parameters. Leveraging the ideas of Rao et al. [55] and McDonald et al. [49] in the context of Deep Learning, we show that SFW is also able to achieve excellent results in the structured pruning domain, namely convolutional filter pruning and low-rank matrix decomposition. To the best of our knowledge, our work is the first to apply SFW for structured pruning tasks. In analyzing the techniques introduced by Pokutta et al. [54], we find that the gradient rescaling of the learning rate is of utmost importance for obtaining high performing and pruning stable results. We lay the theoretical foundation for this practice by proving the convergence of SFW with gradient rescaling in the non-convex stochastic case, extending results of Reddi et al. [57].
Contributions. The major contributions of our work can be summarized as follows: 1. We significantly improve upon the results of Miao et al. [50] regarding the usage of SFW in the pruning-aware context by leveraging k-support norm constraints, mitigating the drawbacks of the existing approach and leading to more robustness to pruning. We also demonstrate that our approach is more stable with respect to the hyperparameter selection.
2. We extend the SFW-prune framework to the structured compression domain. First, we show how to expand the rationale from the unstructured pruning case to that of convolutional filter pruning by constraining the group-k-support norm of convolutional layers, actively decaying the norm of seemingly unimportant filters to ensure stability when pruning. Further, we propose to enforce robustness to low-rank tensor decomposition by interpreting convolutional tensors as matrices and constraining their spectral-k-support norm, resulting in low-rank updates throughout training. Our approach significantly outperforms other baseline approaches in terms of efficiency by requiring only the first few singular values to compute the update direction.
3. Our empirical results suggest that the robustness of SFW can largely be attributed to the usage of the gradient rescaling of the learning rate, which increases the batch gradient norm and effective learning rate throughout training, even though the train loss constantly decreases. To justify the usage of gradient rescaling theoretically, we prove the convergence of SFW with batch gradient dependent step size in the non-convex setting.
Compression-aware training is a promising research direction for training models to state-of-the-art performance while encouraging stability to pruning. We believe that our work is an important building block in the design of structured training algorithms. One strength is the fact that the proposed methods cover a wide range of compression domains, i.e., unstructured pruning, structured pruning, as well as matrix decomposition. Our results show the suitability of the SFW algorithm and highlight the importance of the learning rate rescaling, which we justify theoretically in the hope of enabling further research.
Related Work. The Frank-Wolfe (FW) [21] or conditional gradient [43] algorithm has been studied extensively in the convex setting, enjoying popularity throughout various domains of Machine Learning for being able to efficiently deal with complex structural requirements [e.g. 20,33,39,52,71]. With convexity being a relatively strong assumption, Lacoste-Julien [38] extended the convergence theory of FW to the non-convex setting, while Hazan and Luo [28] and Reddi et al. [57] provide convergence rates for the stochastic variant of the algorithm (SFW). Several different accelerated variants have been proposed, including variance reduction methods [28,58,69], adaptive gradients [16] and momentum [14,51].
With the theoretical foundation for the non-convex stochastic setting being laid, SFW has received a surge of interest in the context of training Neural Networks: Ravi et al. [56] advocate the usage of parameter constraints in Deep Learning, Xie et al. [66] train shallow networks using SFW, Berrada et al. [7] design a variant specifically for Neural Networks, Tsiligkaridis and Roberts [61] employ SFW for adversarial training and Pokutta et al. [54] show that SFW can reach state-of-the-art performance in benchmark image classification tasks. While there is a rich literature on classical FW being applied to sparsity problems in Machine Learning, only few have considered exploiting the structure-enhancing properties of (stochastic) FW in the field of Deep Learning. Grigas et al. [24] remove neurons from three layer convolutional architectures. Pokutta et al. [54] propose to constrain the parameters to lie within a k-sparse polytope, resulting in a large fraction of the parameters having small magnitude. Miao et al. [50] leverage this idea in the context of unstructured magnitude pruning with a focus on pruning-aware training, being a compression setting of training 'once-for-all' sparsities [10]. For a detailed account of different sparsification approaches we refer to the excellent survey of Hoefler et al. [30].
Outline. We begin by introducing the necessary notation, the problem setting as well as the SFW algorithm. Section 2 further contains precise explanations of the feasible regions we intend to use in the unstructured and structured pruning case as well as for matrix decompositions. Section 3 is devoted to experimentally comparing the employed constructions to existing approaches. Section 4 contains an analysis of two learning rate rescaling mechanisms. Finally, we conclude and discuss the findings of this paper in Section 5.

Preliminaries and Methodology
For x ∈ R n , we denote the i-th coordinate of x by [x] i . The diagonal matrix with x on its diagonal is denoted by diag(x) ∈ R n,n . For p ∈ [1, ∞], the L p -ball of radius τ is denoted by B p (τ ). x 0 denotes the number of non-zero components of x ∈ R n . For any compact convex set C ⊆ R n , let us further denote the L 2 -diameter of C by D(C) = max x,y∈C x − y 2 . As usual, we denote the gradient of a function L at θ by ∇L(θ) and the batch gradient estimator by ∇L(θ). For the sake of convenience, we abuse notation and apply univariate functions to vectors in an elementwise fashion, e.g., |x| denotes the vector |x| := (|x 1 | . . . , |x n |).

Stochastic Frank-Wolfe
We aim at optimizing the parameters θ of a Neural Network while enforcing structure-inducing constraints by considering the constrained finite-sum optimization problem where the loss functions i and L are differentiable in θ and C is a compact, convex set. When using SGD, imposing hard constraints requires a potentially costly projection back to C to ensure feasiblity of the iterates. Our focus however lies on the Stochastic Frank-Wolfe (SFW) algorithm [7,21,54], which is projection-free and perfectly suited for yielding solutions with structural properties. In each iteration t, SFW calls a linear minimization oracle (LMO) on the (stochastic) batch gradient which is then used as the direction to update the parameters using the convex combination where η t ∈ [0, 1] is a suitable learning rate. If the initial parameters θ 0 are ensured to lie in the convex set C, then the update rule ensures feasibility of the parameters in each iteration. Solving Equation (2) is often much cheaper than performing a projection step [15,33], in many cases even admitting a closed-form solution. Choosing a feasible region where the v t are highly structured can be beneficial to generalization [54,56] and is very relevant in the context of sparsity, an example being the cross-polytope C = B 1 (τ ) which results in update vectors having only a single non-zero entry [54]. If v t is sparse, then by Equation (3) all remaining parameters are discounted strongly, encouraging convergence towards sparse solutions [50,54].
If C is given by the convex hull of (possibly infinitely many) vertices, a so-called atomic domain, then the solution to Equation (2) is attained at one of these vertices [33]. If C is full-dimensional and symmetric around the origin, then the corresponding gauge function defines a proper norm [25], i.e., C is a ball with respect to some atomic norm, which might be non-trivial to compute. Throughout the following sections, we will often omit the actual norm at hand but rather rely on the specification of C as an atomic domain. Note however that the norm is essential in ensuring the feasibility of θ 0 , since the typical initializations [e.g. 23,29] are not guaranteed to lie in C.

Structured feasible regions
Generalizing the L 1 -ball B 1 , Pokutta et al. [54] introduced the k-sparse polytope of radius τ , C = P k (τ ), which is defined as the convex hull of all vectors v ∈ {0, ±τ } n with at most k non-zero entries. This set can be characterized as the ball of radius τ with respect to the k-sparse norm · P k = max { · ∞ , · 1 /k}. The solution to Equation (2) can be given as the k-sparse vector v satisfying is the set of k indices corresponding to the k largest elements of |∇ t |. The update of Equation (3) then only updates the k parameters corresponding to the directions of highest decrease in loss, while the remaining n − k parameters are decayed towards zero. Pokutta et al. [54] showed that training under these constraints leads to a larger fraction of the parameters being relatively small. Miao et al. [50] investigated this idea further in the context of magnitude-based (unstructured) pruning and showed that, for certain properly chosen values of τ and k, SFW training under k-sparse polytope constraints can lead to networks that are robust to pruning for a wide variety of sparsities.
This approach however suffers from two drawbacks: first of all, all activated parameters will receive an update of same magnitude, namely τ . This hinders convergence, especially when k is large. Secondly, Pokutta et al. [54] and Miao et al. [50] specify the desired L 2 -diameter D of C = P k (τ ) to control the regularization strength and then in turn choose the radius τ such that D(P k (τ )) = D. Defined this way, τ depends on k as τ = D/(2 √ k). This is counter-intuitive, since k controls both the amount of activated parameters as per design of the LMO as well as the magnitude of the parameter updates, resulting in unnecessarily coupled parameters.

Unstructured Pruning using the k-support norm
To remedy these drawbacks, we propose using the feasible region C k (τ ) given by the convex hull over all k-sparse vectors with bounded euclidean norm: The LMO solution to Equation (2) of C k (τ ) is k-sparse and takes the distribution of magnitude between the k largest entries of ∇ t into account, as stated by the following lemma.
be the vector obtained by setting the n − k magnitude-wise smallest entries of ∇ t to zero.
A proof can be found in the appendix. Similarly to the k-sparse polytope, SFW applied to C k (τ ) updates the k parameters whose (stochastic) gradient entries correspond to those of fastest loss decrease, however accounting for the distribution of magnitude among them instead of using the same magnitude for all parameters. Due to the convex update rule, the remaining parameters are decayed, eventually resulting in few of them not being close to zero and thus making the trained network robust to magnitude pruning. As opposed to the k-sparse polytope, the diameter D(C k (τ )) does not depend on k and hence decouples the parameters k and τ as desired. C k (τ ) is known in the regression and statistics literature as the k-support norm ball [4]. McDonald et al. [49] further generalized this notion by introducing the (k, p)-support norm, which is the gauge function corresponding to the set C (k,p) (τ ) obtained by replacing the L 2 -norm with the L p -norm in Equation (4) for arbitrary p ∈ [1, ∞]. Clearly, this notion incorporates both the k-support norm (which is the case of p = 2), but also the k-sparse norm. For p = ∞, the vertex set of C (k,∞) (τ ) is precisely that of P k (τ ).
In Section 3.1 we extensively study the usability of the k-support norm approach to learn pruning stable solutions in the context of unstructured magnitude pruning. Our findings show that leveraging the k-support constraints significantly outperforms the approach based on the k-sparse polytope. Further, we find that the former is more stable to hyperparameter tuning.

Structured Filter Pruning using the Group-k-support norm
The rationale for introducing the k-support norm is not directly transferable to the setting of pruning structural components such as convolutional filters. However, following Rao et al. [55], the k-support norm can be extended to consider groups of parameters as follows: given a disjoint partition of the parameters into groups G ∈ G, define the group-k-support norm ball of radius τ as where v 0,G is the smallest number of groups that are needed to cover the support of v. Choosing G as the set of (disjoint) convolutional filters, we obtain a feasible region that extends the rationale of the k-support norm to the structured pruning case, since we can state the solution to Equation (2) as follows. For a group G ∈ G let x G be the L 2 -norm of x when only considering elements of G. Given the batch gradient of the l-th convolutional layer ∇ l t , let G 1 , . . . , G k be the k groups (i.e. filters) with the largest L 2 -norm ∇ l t Gi and let H : Then the solution to Equation (2) is given by otherwise. A proof can be found in Rao et al. [55]. Section 3.2 will be devoted to studying this group extension of the k-support norm in the structured pruning case, namely that of filter pruning. Clearly, this approach is easily extendable to the pruning of other groups, such as neurons. We compare SFW to the baseline of SGD with a group penalty on the filters and show the superiority of the former with respect to generating pruning-robust solutions.

Low-Rank Matrix Decomposition using the Spectral-k-support norm
Instead of removing individual parameters or groups thereof, networks can also be compressed after training by decomposing parameter matrices into a product of smaller matrices, allowing one to replace a layer by two consecutive ones that require a drastically smaller amount of FLOPs at inference [17]. The key ingredient is the truncated singular value decomposition (SVD), where setting the smallest singular values to zero leads to an optimal low-rank approximation by virtue of the Eckart-Young-Mirsky theorem. More precisely, given a rank r parameter matrix W ∈ R n,m with singular values σ(W) = (σ 1 (W), . . . , σ r (W)) and SVD W = where the magnitude of 'pruned' singular values quantifies the error in approximation. A detailed account of this approach can be found in the appendix. A natural approach to ensure robustness to matrix decomposition is hence based on penalizing the nuclear norm W * := σ(W) 1 [3,60], which requires the costly computation of the full SVD in each iteration.
When constraining the parameters to have bounded nuclear norm instead of penalizing it, the LMO solution to Equation (2) utilized by SFW can be computed efficiently by requiring only the first singular value-vector-pair [33]. Extending this notion to consider the k largest singular pairs, McDonald et al. [49] propose the Spectral-k-support norm, whose ball of radius τ is defined as where σ(W) 2 is the 2-Schattennorm of the singular values σ(W) [33]. Using the insights from the previous sections, we propose to constrain the Spectral-k-support norm of the parameter matrices. The following lemma allows us to cheaply compute the LMO solution.
that only the k largest singular values are kept. Then W t is a solution to Equation (2).
A proof can be found in the appendix. Note that similar to the k-support norm ball taking the magnitude of k largest gradient entries into account, the scaling by Σ k σ(Σ k ) −1 2 takes the magnitude of k largest singular values into account. In Section 3.3 we study the capabilities of SFW when constraining the Spectral-k-support norm of convolutional tensors, which account for the majority of FLOPs at inference [26]. While there exist higher-order generalizations of the SVD to decompose tensors directly [see e.g. 35,40], we follow the approach of interpreting the tensor W ∈ R n×c×d×d with c in-channels, n convolutional filters and spatial size d as an (n × cd 2 )-matrix [3,32]. In comparison to other baseline approaches, we analyze to what extend the SFW algorithm is able to generate matrix-decomposition robust networks while needing to compute significantly less singular values in each iteration.

Experimental Setup
Let us describe the general methodological approach used throughout the experiments, deferring the details to the respective sections and the appendix. All experiments are conducted using the PyTorch framework [53], where we relied on Weights & Biases [8] for the analysis of results. To enable reproducibility, our implementation is publicly accessible at github.com/ZIB-IOL/compression-aware-SFW.
The focus of our experiments will lie on image classification tasks: we train deep convolutional architectures such as Residual Networks [29] and Wide Residual Networks [70] on the benchmark datasets CIFAR-10 and CIFAR-100 [36]. The exact training setups as well as grid searches used can be found in the appendix. As a general remark, we follow the experimental guidelines of Blalock et al. [9] towards standardized comparisons between sparsification methods. All results are averaged over two seeds with min-max bands indicated for plots and standard deviation for tables.
In the pruning-aware setting we are interested in finding single hyperparameter configurations that perform well under a wide variety of compression rates, i.e., without tuning hyperparameters for each sparsity. However, strongly regularizing configurations might be underperforming for moderate sparsities or even when keeping the model dense, while achieving decent results in the extremely high sparsity regime. This is undesired since we seek for hyperparameter configuration that both reach state-of-the-art accuracy for dense training and are robust to compression. To avoid such behaviour, we filter out all dense models that suffer from a greater than five percent drop in train accuracy compared to a default training setup (details can be found in the appendix).

Compression-aware Optimization using the Frank-Wolfe Algorithm
To evaluate SFW in the compression-aware training context, we compare the proposed feasible regions to suitable baselines in the domains of unstructured magnitude pruning, (structured) convolutional filter pruning, and low-rank matrix decomposition. The main part contains results for the ResNet-18 architecture trained on CIFAR-10, other training configurations are deferred to the appendix. Figure 1a compares the proposed k-support norm ball approach to the k-sparse polytope training of Miao et al. [50] as well as SGD with weight decay. Each line corresponds to the single hyperparameter configuration which results in the highest on-average test accuracy after pruning. k-support norm training with SFW leads to models that are robust to pruning both in the medium to high sparsity regime, experiencing only little degradation when removing 98% of the parameters. The k-sparse approach performs well in the medium to high sparsity regime, but quickly collapses for higher compression rates. A clear advantage of k-support norm ball constraints is that SFW is able to obtain this performance in the high compression regime while not suffering from underperformance before pruning. This distinguishes our approach from SGD with weight decay, which is able to obtain better results than the approach of Miao et al. [50] if weight decay is properly tuned, however at the expense of performing worse in the medium sparsity regime. Further, Figure 4 in the appendix shows the   successful decoupling of the radius and k: the k-support norm ball is less sensitive to hyperparameter changes and obtains better results throughout a wide range of hyperparameter configurations than its k-sparse counterpart.

Structured Filter Pruning
For the pruning of convolutional filters, we follow the PFEC approach of Li et al. [44]: at the end of training we sort the filters of each convolutional layer by their L 1 -norm and remove the smallest ones until the desired level of compression is met. Figure 1b compares the proposed group-k-support norm constraint approach against the natural baseline of SGD with a group penalty [2], details as well as exact grids used can be found in the appendix. All lines correspond to the tradeoff curves when selecting the configuration that performs best on-average among all sparsities at stake. Although not depicted here, these curves closely match the ones obtained when allowing different hyperparameter configurations for different sparsities. For SFW, non-convolutional parameters are constrained to lie in appropriate L 2 -balls.
SFW converges to solutions that are more robust to pruning than the penalty based SGD-approach. That becomes especially visible when pruning between 65% to 75% of the convolutional filters in each layer: SFW is able to keep most of its performance while SGD starts to collapse towards the accuracy of a random classifier. We further noticed that filtering out hyperparameter configurations that lead to underperforming dense networks has almost no impact for SFW-trained models, indicating that the performance before pruning is largely unaffected by the regularization induced by the LMO. However, for SGD the opposite holds: strong regularization leads to decently performing networks in the high sparsity regime, however these hyperparameter configurations are unable to produce a well-performing model before pruning. In the low sparsity regime, these runs produce inacceptable results. Full results can be found in the appendix.

Low-Rank Matrix decomposition
To evaluate the capabilities of SFW for tensor decomposition as proposed in Section 2.2.3, we compare the Spectral-k-support norm approach to two natural baselines: nuclear norm penalty regularization with SGD and the corresponding proximal gradient descent (ProxGD) approach. Similarly to classical weight decay, the former computes the subgradient of a nuclear norm regularization term [64] and updates the weights accordingly. The latter performs an SGD update on the loss, followed by applying the soft-thresholding operator over the singular values of the parameters, a strategy also known as singular value thresholding [3,11]. Both approaches require the computation of the entire SVD, tradionally requiring O(nm min(n, m)) operations for an n × m matrix [1]. A detailed account of both methods can be found in Appendix A.2.3. We only apply these regularization mechanisms to the convolutional tensors (interpreted as matrices as previously outlined), while all other parameters are updated normally by SGD with weight decay. Similarly for SFW, we only constrain the convolutional tensors using the Spectral-k-support norm ball, while all other parameters are constrained to lie in appropriate L 2 -balls. At the end of training, we set the smallest singular values of each convolutional matrix to zero and replace the layer by two consecutive layers as described. The accuracies depicted correspond again to the on-average best hyperparameter configuration. A significant advantage of SFW in this regard is its efficiency, requiring only roughly half of the time of SGD and ProxGD, where we performed measurements on the same hardware, namely a 24-core Xeon Gold with Nvidia Tesla V100 GPU. Since the LMO requires the computation of the k largest singular pairs, the efficiency of SFW is clearly dependent on k. The on-average best configuration of SFW in this setting is given by k corresponding to ten percent of the singular values. We used a naive implementation of the k-SVD power method [6], noting that there are more sophisticated and faster algorithms [1].  This coupling between regularization strength and step size makes the tuning of the learning rate cumbersome. To decouple the tuning of the learning rate from the size of the feasible region, Pokutta et al. [54] propose two different learning rate rescaling mechanisms: diameter rescaling and gradient rescaling, the latter being used throughout our experiments in the previous section. While the former divides the learning rate by the L 2 -diameter D(C) of C, gradient rescaling rescales the update direction length to that of the batch gradient, i.e.,η t : It is however largely unclear what effect these normalization schemes have on both the convergence for regular training as well as the robustness to pruning, noting again that the learning rate of SFW explicitly controls the decay on non-activated parameters. To investigate the differences between gradient and diameter rescaling, we performed a large hyperparameter search over both k and τ for SFW with k-support norm constraints. Figure 2 shows the test accuracy vs. sparsity tradeoff curves directly before and after magnitude-pruning of the parameters. For each pruning sparsity, the best possible hyperparameter configuration is selected. The dashed lines show the corresponding accuracies pre-pruning when using the same hyperparameter configurations. We find that gradient rescaling consistently outperforms its diameter-based counterpart w.r.t. both dense as well as pruned test accuracy. While the robustness to pruning, indicated by the shaded areas, is more or less comparable (with slight advantages for gradient rescaling, since typically higher dense accuracies lead to larger drops), we notice that gradient rescaling is more consistent in its results (indicated by the min-max bands).
In our experiments, we found the denominator of gradient rescaling to influence the scale of the learning rate but not to be subject to much variation, whereas the batch gradient norm dynamically changes the learning rate over time. Figure 3 shows the evolution of ∇ t for two different radii of the k-support norm ball (with fixed k), where we compare to usual SGD training with weight decay. Note that for classical SGD the effective learning rate implicitly contains the batch gradient norm as well, since the gradient direction is not normalized. For both SGD and SFW, ∇ t is subject to noise and increases until 75% of the training process, despite the continuous decrease of the train loss. In fact, the batch gradient norm is not significantly smaller than at the start of training even though the loss converges. Similar observations have recently been made by Zhang et al. [73], who show a gap between classical optimization theory and the actual behaviour of deep neural networks. This behaviour might best be explainable by the presence of Batch-Normalization layers, whose interplay with weight decay has recently been analyzed by van Laarhoven [62] and Hoffer et al. [31]: layers preceding a Batch-Normalization are rescaling invariant, that is their output remains unchanged when multiplying all parameters by a scalar, however rescaling them inversely rescales the gradient norm in subsequent layers and iterations.  Weight decay continuously decreases the scale of the parameters and hence increases the scale of the batch gradient, where stronger decay of the former leads to stronger increase of the latter. Since in gradient rescaling the norm of the batch gradient also influences the strength of the decay of the parameters, this process has a self-accelerating dynamic. As visible in Figure 3, we observe amplification of the noise in the batch gradient norm and consequently in the learning rate, potentially being the cause for the higher generalization capabilities of gradient over diameter rescaling. More importantly however, the self-increasing dynamic of the learning rate results in larger steps towards the (sparse) vertices of the k-support norm ball, leading to a stronger decay on the previous parameter configuration, strongly discounting the remaining n − k directions. This increases the robustness to pruning, making gradient rescaling the method of choice in that setting.
With the following theorem, we lay the theoretical foundation by showing that incorporating the batch gradient norm into the learning rate leads to convergence of SFW at the specified rate, that is, the expected product of exact gradient norm and the Frank-Wolfe Gap G(θ t ) decays at a rate of O(T −1/2 ). The precise statement as well as a proof can be found in Appendix A.5. Theorem 4.1 (Convergence of gradient rescaling, informal). Assume that L is M-smooth and is G-Lipschitz and let η t = ∇ t η for appropriately chosen η and all 0 ≤ t < T . If θ a is chosen uniformly at random from the SFW iterates {θ i : 0 ≤ i < T }, then we have E [G(θ a ) · ∇L(θ a ) ] = O(T −1/2 ), where E denotes the expectation w.r.t. all the randomness present.

Discussion and Outlook
We proposed to combine the SFW algorithm with highly structured feasible regions in order to train once and obtain state-of-the-art dense performance as well as robustness to compression for a wide range of sparsities. Our experimental results show that the results of Miao et al. [50] can be significantly improved and we further propose novel approaches towards compression-aware training in the structured pruning as well as matrix-decomposition setting. We hope that our findings regarding the importance of the learning rate rescaling as well as Theorem 4.1 stimulate further research in the direction of compression-aware training with SFW.
However, we emphasize that our results hold primarily in the setting that we described, namely that of compression-aware training, where the training is sparsity-agnostic and retraining is prohibitive. Our goal was to show the versatility of SFW, which provides a suitable algorithmic framework for enforcing structure throughout training. If the sparsity can be incorporated into training, significantly more complex approaches can be applied, as outlined in the introduction.

Reproducibility is a crucial aspect of any work based on experiments.
Our code is available at github.com/ZIB-IOL/compression-aware-SFW. We further highlight some publicly available implementations we used.
For the computation of the theoretical speedup, we used the implementation of Blalock et al. [9] available at https://github.com/JJGO/shrinkbench. Regarding pruning methods that are not directly available through the Pytorch framework, we used the original implementation whenever possible and indicate so directly in the code.

A.2 Technical details
Whenever using SFW, we employ a momentum variant [54] and enforce local constraints, i.e., we constrain the parameters of each layer in a separate feasible region. As suggested by Pokutta et al. [54], we do not tune the radius τ of the feasible region C(τ ) directly but rather specify a scalar factor w > 0 and set the L 2 -diameter of the feasible region of each layer as where the expected L 2 -norm of the layer parameters θ are simply estimated by computing the mean among multiple default initializations. D is then in turn used to compute τ such that D(C(τ )) = D. This allows us to control the L 2 -regularization strength of each layer. We do not prune biases and batch-normalization parameters, as they only account for a small fraction of the total number of parameters yet are crucial for obtaining well-performing models [19]. Apart from the direct comparison between diameter and gradient rescaling, we use gradient rescaling throughout all experiments.
In the following, we state technical details of our approaches to unstructured pruning, structured pruning, and low-rank matrix decomposition.

A.2.1 Unstructured Weight Pruning
We employ the usual magnitude pruning, that is, we remove the parameters with smallest absolute value until we meet the desired level of compression. As we found it to work best among all algorithms, we prune globally, i.e. we select the smallest weights among all network parameters eligible for pruning. We note however that there exists several different magnitude-based selection approaches [cf. e.g. 22,26,42].

A.2.2 Structured Filter Pruning
As outlined in the main section, we follow the approach of Li et al. [44] and remove the convolutional filters with smallest L 1 -norm. Since each filter might correspond to a different number of parameters, depending on the convolutional layer it is located in, the L 1 -norm of filters is incomparable among different layers. We hence follow the local approach and prune the same amount in each convolutional layer until the desired sparsity is met. Since this will lead to the same theoretical speedup, independent of the algorithm at stake, we omit these values.
As a baseline comparison method we employ standard SGD with weight decay and group penalty regularization to encourage robustness to filter pruning. That is, similar to Alvarez and Salzmann [2], we minimize where we omit the weight decay term for the sake of legibility. Here, N f is the number of convolutional filters in the network and we penalize the L 2 -norm of each convolutional filter θ i , i = 1, . . . , N f .

A.2.3 Low-Rank Matrix decomposition
We extensively describe the rationale behind the decomposition of matrices with (preferably) low rank. Low-Rank matrix decomposition is centered around the idea of truncating the singular value decomposition (SVD). Let θ = U ΣV T = r i=1 u i σ i v T i be the SVD of rank-r matrix θ ∈ R n×m , where u i , v i are the singular vectors to singular values σ 1 ≥ . . . ≥ σ r > 0. The goal is now to find the best low-rank approximationθ ∈ R n×m to θ, namely to solve By the Eckart-Young-Mirsky theorem, a minimizer is given by where the smallest t − r singular values are truncated. If the singular values now decay rapidly, i.e., only the first t singular values contain most of the energy of θ, then this approximation can be applied as a post-processing step without much change in the output of a linear layer [17]. Consequently, a natural approach to ensure stability w.r.t. matrix decomposition is based on encouraging the parameter matrices to have low-rank throughout training, most prominently done by regularizing the nuclear norm of the weight matrix [3,60]. If t is small, then the number of FLOPs can be drastically reduced by decomposing a layer into multiple layers. For linear layers this is straightforward: let (θ, β) be weights and biases of a linear layer with SVD of θ as above. For input x, the layer computes being interpretable as the consecutive application of two linear layers: (Σ t V T t , 0) followed by (U t , β), possibly reducing the number of parameters from nm to t(n + m). For a four-dimensional convolutional tensor θ ∈ R n×c×d×d , where c is the number of in-channels, n the number of convolutional filters, and d is the spatial size, we cannot directly construct the SVD. However, we follow an approach similar to those of Alvarez and Salzmann [3] and Idelbayev and Carreira-Perpinán [32], interpreting θ as a (n × cd 2 ) matrix, whose truncated SVD decomposition allows us to replace the layer by two consecutive convolutional layers, the first one having t filters, c channels and spatial size d, followed by n filters, t channels and spatial size of one.
We describe the baselines for low-rank matrix decomposition in more detail. Both approaches solve the problem min where λ > 0 is a suitably chosen regularization parameter. The SGD approach does so by computing the subgradient of λ θ * [64] and performing an ordinary SGD update. On the other hand, ProxGD relies on the proximity operator of the nuclear norm, also known as singular value thresholding [3,11]. In each iteration this amount to first applying an SGD step to minimize L(θ t ) with learning rate η t obtainingθ t+1 , which is then followed by soft-thresholding the singular values ofθ t+1 to obtain θ t+1 , i.e., where u i , v i are the singular vectors ofθ t+1 to singular values σ i . In essence, singular values smaller than η t λ are cut-off, potentially resulting in θ t+1 having low rank. Table 2 shows the exact training configurations we used throughout all experiments, where we relied on a shortened training timeframe of 100 epochs with a linearly decaying learning rate, as suggested by Li et al. [45]. The indicated dense test accuracies are used to define which models to filter out, i.e., for CIFAR-10 we allow all runs that reach a dense test accuracy of at least 90%.

A.3 Experimental setup and extended results
In the following, we state the hyperparameter grids used as well as full tables and plots.

CIFAR-10: SGD
We only tune the weight decay factor λ ∈ {0, 1e-6, 5e-6, 1e-5, 5e-5, 1e-4, 5e-4, 1e-3, 5e-3, 1e-2}, since there are no additional hyperparameters. Figure 4 compares the two feasible regions in an even larger hyperparameter search. The rows correspond to the k-sparse polytope (above) and k-support norm ball (below), respectively. The left column shows a heatmap of the test accuracy before pruning. While both approaches lead to well performing models for a wide range of hyperparameter configurations (indicated as the radius multiplier w on the x-axis and k on the y-axis), the k-support norm ball reaches higher results and converges properly for all configurations at stake. The k-sparse polytope approach fails to yield adequately trained dense models when the radius is relatively small but k becomes larger, which is counter-intuitive, since larger k allows a larger fraction of the parameters to be activated. The right column shows the corresponding heatmap of the test accuracy right after pruning. Clearly, the proposed approach is robust to pruning for a wider hyperparameter range.  Figure 4: ResNet-18 on CIFAR-10: Contour plot when performing a large hyperparameter search over the radius and k of the feasible regions, where the first row corresponds to the k-sparse polytope and the second one corresponds to the k-support norm ball. The left column shows the test accuracy before pruning, while the right column shows the test accuracy after pruning. The k-support norm approach leads to better performing dense models given the hyperparameter search at stake, which in turn are more stable to pruning.
Miao et al. [50] showed that SFW (with k-sparse polytope constraints) outperforms SGD with weight decay, which in turn clearly, and unsurprisingly, outperforms the SFW-based approach when it is allowed to retrain. Our experiments indicate that while being less robust to pruning, SGD is able to reach on-par or better results after retraining, even when SFW is allowed to be retrained for the same amount of time. Leaving the domain of compression-aware training, this raises a more general question: in the case that retraining is not prohibited, is it beneficial to aim for robustness at pruning when trying to maximize the post-retraining accuracy? Figure 5 illustrates an experiment where we investigate this exact question by performing One-Shot IMP [26] to a sparsity of 95% and retraining for 10 epochs using LLR [76]. We tuned both the weight decay for regular retraining as well as the weight decay for the retraining phase. Non surprisingly, there is a weight decay sweet spot when it comes to maximizing the pre-pruning accuracy (left). The middle plot shows that higher weight decay typically leads to more robustness to pruning, however a too large weight decay hinders convergence of the dense model and might lower the performance after pruning. Surprisingly however, as depicted in the right plot showing the test accuracy after retraining, the optimal parameter configuration is the one that leads to the highest accuracy before pruning, which is also the least robust to pruning. This aligns with previous findings of Bartoldson et al. [5] and Zimmer et al. [76], who question the strive for pruning stability when retraining is possible. We once again highlight that retraining is prohibitive in the compression-aware training context.

CIFAR-100: SGD
We tune the weight decay factor λ 0 ∈ {1e-4, 2e-4, 5e-4, 1e-3} and the filter group penalty factor λ 1 ∈ {0, 1e-5, 5e-5, 1e-4, 2e-4, 5e-4, 1e-3, 2e-3, 5e-3}. Figure 7 show the results for filter pruning on CIFAR-10 and CIFAR-100, respectively. As outlined in the main section, we filter out runs that lead to too much performance degradation (b) Filter pruning: Comparison of the group-ksupport norm ball approach to SGD with weight decay and filter group penalties without removing runs that do not perform well before pruning. Figure 6: ResNet-18 on CIFAR-10: Accuracy-vs.-sparsity tradeoff curves for structured convolutional filter pruning. The plots show the parameter configuration with highest test accuracy after pruning when averaging over all sparsities at stake. The left plot does not include runs that perform poorly before pruning, i.e. those that lead to a degradation of at least five percent compared to usual SGD training with weight decay. The right plot additionally includes these runs. Filter pruning (without run filtering) SFW SGD (b) Filter pruning: Comparison of the group-ksupport norm ball approach to SGD with weight decay and filter group penalties without removing runs that do not perform well before pruning. Figure 7: WideResNet on CIFAR-100: Accuracy-vs.-sparsity tradeoff curves for structured convolutional filter pruning. The plots show the parameter configuration with highest test accuracy after pruning when averaging over all sparsities at stake. The left plot does not include runs that perform poorly before pruning, i.e. those that lead to a degradation of at least five percent compared to usual SGD training with weight decay. The right plot additionally includes these runs. before pruning, that is, runs that converge to an accuracy drop of at least five percent compared to the values obtained by usual SGD training (cf. Table 2). The left plots of aforementioned figures show the results when filtering out these poorly performing runs, while in the right plots the hyperparameter selection for the on-average best run includes the entire grid search. For CIFAR-10, we notice that the filtering has almost no impact for the SFW-trained models, as the curves in the left and right plot are the same. For SGD we find that strong regularization that is unable to produce well-performing dense models can lead to decently performing networks in the high sparsity regime, while failing in the low sparsity regime. For CIFAR-100, we find that SFW consistently outperforms SGD. When allowing underperforming runs to be included for the hyperparameter selection, both optimization strategies perform drastically better in the high sparsity regime, albeit at the price of deteriorating in the low sparsity regime. (b) Matrix decomposition: Comparison of the spectral-k-support norm ball approach to nuclear norm regularization approaches without removing runs that do not perform well before pruning. Figure 8: ResNet-18 on CIFAR-10: Accuracy-vs.-sparsity tradeoff curves for low-rank matrix decomposition. The plots show the parameter configuration with highest test accuracy after pruning when averaging over all sparsities at stake. The left plot does not include runs that perform poorly before pruning, i.e. those that lead to a degradation of at least five percent compared to usual SGD training with weight decay. The right plot additionally includes these runs.

A.4 Proofs of LMO constructions
In the following, we state the missing proofs of Lemma 2.1 and Lemma 2.2. Lemma A.1. Given ∇ t , let v t ∈ C k (τ ) such that where ∇ top k t is the vector obtained by setting to zero all n − k entries [∇ t ] j of ∇ t with j ∈ top k (|∇ t |). Then v t ∈ arg min v∈C k (τ ) v, ∇ t is a solution to Equation (2).
Proof. Recall that C k (τ ) = conv{v | v 0 ≤ k, v 2 ≤ τ }. By construction, all vertices v of C k (τ ) satisfy v 2 = τ and are k-Sparse, i.e., v 0 ≤ k. Note that being k-Sparse includes cases where more than n − k entries are zero. The minimum of Equation (2) is attained at one such v. Further recall the following reformulation of the euclidean inner product: v, ∇ t = v 2 ∇ t 2 cos(∠(v, ∇ t )) = τ ∇ t 2 cos(∠(v, ∇ t )), where ∠(v, ∇ t ) denotes the angle between v and ∇ t . This term is minimized as soon as the angle between v and ∇ t is maximal. If v was not required to be k-Sparse, i.e., v would be allowed to lie anywhere on the border of B 2 (τ ), the solution would clearly be given by −τ ∇ t / ∇ t 2 . However, since v is k-Sparse, the vector maximizing the angle to ∇ t is the one that is closest to −τ ∇ t / ∇ t 2 but is k-Sparse at the same time. This is exactly the one claimed.
Lemma A.2. Given ∇ t ∈ R n×m , let W t ∈ C σ k (τ ) such that where U k Σ k V T k is the truncated SVD of ∇ t such that only the k largest singular values are kept. Then W t ∈ arg min v∈C σ k (τ ) v, ∇ t is a solution to Equation (2).
Let W be some minimizer. Note that rescaling a matrix by a scalar has no effect on its rank. Let us hence assume that σ(W) 2 = α for some α > 0 and characterize W ∈ arg min rank(v)≤k v, ∇ t . Again, we have W, where ← − x is there vectorized form of matrix x. Since we can choose α ≤ τ , this term is minimal as soon as the angle ∠( ← − W, ← − ∇ t ) is maximal, i.e. cos(∠( ← − W, ← − ∇ t )) < 0 and α = τ , where we use the same euclidean-geometric interpretation as in the proof for Lemma 2.1 above. To obtain a maximal angle, we hence minimize the L 2 -distance between − ← − W and ← − ∇ t in compliance with the rank constraint. Since again − ← − W − ← − ∇ t 2 = − W − ∇ t F , the Eckart-Young-Mirsky theorem yields the claim, where we rescale appropriately to meet the Schattennorm constraint.

A.5 Convergence of SFW with gradient rescaling
Before priving the convergence of SFW with gradient rescaling as stated it informally in Theorem 4.1, we first recall some central definitions and assumptions.

A.5.1 Setting
Let Ω be the set of training datapoints from which we sample uniformly at random. In Equation (1) we defined a unique loss function i for each datapoint. In the following let (θ, ω i ) = i (θ) for ω i ∈ Ω. Similar to Reddi et al. [57] and Pokutta et al. [54], we define the SFW algorithm as follows, where the output θ a is chosen uniformly at random from all iterates θ 0 , . . . , θ T −1 .
We will use the same assumptions as Reddi et al. [57]. First of all, let us assume that L is M -smooth, that is for all x, y ∈ C, which implies the well-known inequality L(x) ≤ L(y) + ∇L(y), x − y + M 2 x − y 2 .
Further, we assume the function to be G-Lipschitz, that is for all x ∈ C and ω ∈ Ω we have A direct consequence is that the norm of the gradient estimator can be bounded as ∇ L(θ t ) ≤ G.