Original: Spartan: Differentiable Sparsity via Regularized Transportation
Revised: Spartan: Differentiable Sparsity via Regularized Transportation

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: We present Spartan, a method for training sparse neural network models witha predetermined level of sparsity.
Revised: We present Spartan, a method for training sparse neural network models witha predetermined level of sparsity.

Original: Spartan is based on a combination of twotechniques: (1) soft top-
Revised: Spartan is based on a combination of twotechniques: (1) soft top-

Original: k masking of low-magnitude parameters via a regularizedoptimal transportation problem and (2) dual averaging-based parameter updateswith hard sparsiﬁcation in the forward pass.
Revised: k masking of low-magnitude parameters via a regularizedoptimal transportation problem and (2) dual averaging-based parameter updateswith hard sparsiﬁcation in the forward pass.

Original: This scheme realizes an explorationexploitation tradeoff: early in training, the learner is able to explore various sparsitypatterns, and as the soft top- k approximation is gradually sharpened over the courseof training, the balance shifts towards parameter optimization with respect to aﬁxed sparsity mask.
Revised: This scheme realizes an explorationexploitation tradeoff: early in training, the learner is able to explore various sparsitypatterns, and as the soft top- k approximation is gradually sharpened over the courseof training, the balance shifts towards parameter optimization with respect to aﬁxed sparsity mask.

Original: Spartan is sufﬁciently ﬂexible to accommodate a varietyof sparsity allocation policies, including both unstructured and block structuredsparsity, as well as general cost-sensitive sparsity allocation mediated by linearmodels of per-parameter costs.
Revised: Spartan is sufﬁciently ﬂexible to accommodate a varietyof sparsity allocation policies, including both unstructured and block structuredsparsity, as well as general cost-sensitive sparsity allocation mediated by linearmodels of per-parameter costs.

Original: On ImageNet-1K classiﬁcation, Spartan yields 95%sparse ResNet-50 models and 90% block sparse ViT-B/16 models while incurringabsolute top-1 accuracy losses of less than 1% compared to fully dense training.
Revised: On ImageNet-1K classiﬁcation, Spartan yields 95%sparse ResNet-50 models and 90% block sparse ViT-B/16 models while incurringabsolute top-1 accuracy losses of less than 1% compared to fully dense training.

Original:  Introduction
Revised:  Introduction

Original: Sparse learning algorithms search for model parameters that minimize training loss while retainingonly a small fraction of non-zero entries.
Revised: Sparse learning algorithms search for model parameters that minimize training loss while retainingonly a small fraction of non-zero entries.

Original: Parameter sparsity yields beneﬁts along several axes: reduced model storage costs, greater computational and energy efﬁciency during training and inference,and potentially improved model generalization [32].
Revised: Parameter sparsity yields beneﬁts along several axes: reduced model storage costs, greater computational and energy efﬁciency during training and inference,and potentially improved model generalization [32].

Original: However, sparse training is a computationallychallenging task: for the general optimization problem of learning parameters θ ∈ R d subject to theconstraint that θ is k -sparse, the feasible set is the union of (cid:0) d k (cid:1) axis-aligned hyperplanes intersectingat the origin, each of dimension k .
Revised: However, sparse training is a computationallychallenging task: for the general optimization problem of learning parameters θ ∈ R d subject to theconstraint that θ is k -sparse, the feasible set is the union of (cid:0) d k (cid:1) axis-aligned hyperplanes intersectingat the origin, each of dimension k .

Original: This complex, nonconvex geometry of the feasible set compoundsthe existing difﬁculty of optimizing deep neural network models.
Revised: This complex, nonconvex geometry of the feasible set compoundsthe existing difﬁculty of optimizing deep neural network models.

Original: This work focuses on minimizing generalization error for sparse neural network models.
Revised: This work focuses on minimizing generalization error for sparse neural network models.

Original: To this end,we introduce Spartan , or Sparsity via Regularized Transportation —a sparse learning algorithm thatleverages an optimal transportation-based top- k masking operation to induce parameter sparsity duringtraining.
Revised: To this end,we introduce Spartan , or Sparsity via Regularized Transportation —a sparse learning algorithm thatleverages an optimal transportation-based top- k masking operation to induce parameter sparsity duringtraining.

Original: Spartan belongs to a family of “dense-to-sparse” algorithms that maintains a dense parametervector θ ∈ R d throughout training [47; 23], in contrast to “sparse-to-sparse” algorithms that adhereto a ˜ O ( k ) memory budget for representing the parameters of a k -sparse model [3; 30; 31; 13; 15].
Revised: Spartan belongs to a family of “dense-to-sparse” algorithms that maintains a dense parametervector θ ∈ R d throughout training [47; 23], in contrast to “sparse-to-sparse” algorithms that adhereto a ˜ O ( k ) memory budget for representing the parameters of a k -sparse model [3; 30; 31; 13; 15].

Original: While computational cost and memory usage at training time are important considerations, Spartanprimarily optimizes for performance at inference time.
Revised: While computational cost and memory usage at training time are important considerations, Spartanprimarily optimizes for performance at inference time.

Original: Intuitively, Spartan aims to achieve a controlled transition between the exploration and exploitationof various sparsity patterns during training.
Revised: Intuitively, Spartan aims to achieve a controlled transition between the exploration and exploitationof various sparsity patterns during training.

Original: In the exploration regime, the goal is for the learner toeasily transition between differing sparsity patterns in order to escape those that correspond to poorminima of the loss function.
Revised: In the exploration regime, the goal is for the learner toeasily transition between differing sparsity patterns in order to escape those that correspond to poorminima of the loss function.

Original: On the other hand, the goal of exploitation is for the learner to optimizemodel parameters given a ﬁxed sparsity pattern, or a small set of similar patterns.
Revised: On the other hand, the goal of exploitation is for the learner to optimizemodel parameters given a ﬁxed sparsity pattern, or a small set of similar patterns.

Original: This latter setting
Revised: This latter setting

Original: [Figure]
Revised: [Figure]

Original: avoids frequent oscillations between disparate sparsity patterns, which may be detrimental to theoptimization process.
Revised: avoids frequent oscillations between disparate sparsity patterns, which may be detrimental to theoptimization process.

Original: We operationalize this intuition through the use of a differentiable soft top-
Revised: We operationalize this intuition through the use of a differentiable soft top-

Original: k masking operationσ βk ( θ ) .
Revised: k masking operationσ βk ( θ ) .

Original: This function maps parameters θ to an approximately sparse output that suppresses lowmagnitude entries in θ .
Revised: This function maps parameters θ to an approximately sparse output that suppresses lowmagnitude entries in θ .

Original: We parameterize the soft top- k mask with a sharpness parameter β : atβ = 0 , σ βk simply scales the input by a constant, and as β → ∞ , the mapping reduces to hard top- kmagnitude-based selection (Figure 1 (a, b)).
Revised: We parameterize the soft top- k mask with a sharpness parameter β : atβ = 0 , σ βk simply scales the input by a constant, and as β → ∞ , the mapping reduces to hard top- kmagnitude-based selection (Figure 1 (a, b)).

Original: Soft top-
Revised: Soft top-

Original: k masking therefore constrains the iterates θ tto be close to the set of exactly k -sparse vectors; we give some geometeric intuition for the effectof this mapping in Figure 1 (c) and Figure 2.
Revised: k masking therefore constrains the iterates θ tto be close to the set of exactly k -sparse vectors; we give some geometeric intuition for the effectof this mapping in Figure 1 (c) and Figure 2.

Original: We implement soft top-
Revised: We implement soft top-

Original: k masking using a regularizedoptimal transportation formulation [11; 41] and demonstrate that this technique scales to networks onthe order of 10 8 parameters.
Revised: k masking using a regularizedoptimal transportation formulation [11; 41] and demonstrate that this technique scales to networks onthe order of 10 8 parameters.

Original: We evaluate Spartan on ResNet-50 [20] and ViT [14] models trained on the ImageNet-1K dataset.
Revised: We evaluate Spartan on ResNet-50 [20] and ViT [14] models trained on the ImageNet-1K dataset.

Original: On
Revised: On

Original: ResNet-50, we ﬁnd that sparse models trained with Spartan achieve higher generalization accuraciesthan those trained with existing methods at sparsity levels of 90% and above.
Revised: ResNet-50, we ﬁnd that sparse models trained with Spartan achieve higher generalization accuraciesthan those trained with existing methods at sparsity levels of 90% and above.

Original: In particular, we train
Revised: In particular, we train

Original: ResNet-50 models to 76 .
Revised: ResNet-50 models to 76 .

Original: 5% top-1 validation accuracy at 95% sparsity and to 74 .
Revised: 5% top-1 validation accuracy at 95% sparsity and to 74 .

Original: 2% accuracy at97.5% sparsity, improving on the previous state-of-the-art by 0 .
Revised: 2% accuracy at97.5% sparsity, improving on the previous state-of-the-art by 0 .

Original: 6% and 1 . 6% respectively.
Revised: 6% and 1 . 6% respectively.

Original: Our sparse
Revised: Our sparse

Original: ViT-B/16 models reduce model storage size by 10 × and inference FLOPs by 7 . 4 × at the cost of a0 .
Revised: ViT-B/16 models reduce model storage size by 10 × and inference FLOPs by 7 . 4 × at the cost of a0 .

Original: 6% accuracy reduction relative to DeiT-B [37].
Revised: 6% accuracy reduction relative to DeiT-B [37].

Original: We further demonstrate that Spartan is effectivefor block structured pruning, a form of structured sparsity that is more amenable to acceleration oncurrent GPU hardware than unstructured sparsity [17].
Revised: We further demonstrate that Spartan is effectivefor block structured pruning, a form of structured sparsity that is more amenable to acceleration oncurrent GPU hardware than unstructured sparsity [17].

Original: On a ViT-B/16 model with 16 × 16 blockstructured pruning, Spartan achieves 79 .
Revised: On a ViT-B/16 model with 16 × 16 blockstructured pruning, Spartan achieves 79 .

Original: 1% top-1 accuracy at 90% sparsity.
Revised: 1% top-1 accuracy at 90% sparsity.

Original: To summarize, we make the following contributions in this paper:
Revised: To summarize, we make the following contributions in this paper:

Original: • We present Spartan, a sparsity-inducing training algorithm based on a soft top- k masking operation.
Revised: • We present Spartan, a sparsity-inducing training algorithm based on a soft top- k masking operation.

Original: We show that Spartan interpolates between two existing sparse learning algorithms:iterative magnitude pruning [47] and Top-
Revised: We show that Spartan interpolates between two existing sparse learning algorithms:iterative magnitude pruning [47] and Top-

Original: K Always Sparse Training [23].• We empirically evaluate Spartan using ResNet-50 and ViT models trained on theImageNet-1K dataset, demonstrating consistent improvements over the prior state-of-the-art.• We study the effect of Spartan’s hyperparameters on an exploration-exploitation tradeoffduring training and on the ﬁnal accuracy of the trained models.
Revised: K Always Sparse Training [23].• We empirically evaluate Spartan using ResNet-50 and ViT models trained on theImageNet-1K dataset, demonstrating consistent improvements over the prior state-of-the-art.• We study the effect of Spartan’s hyperparameters on an exploration-exploitation tradeoffduring training and on the ﬁnal accuracy of the trained models.

Original: Notation.
Revised: Notation.

Original: We use R + to denote the set of nonnegative real numbers.
Revised: We use R + to denote the set of nonnegative real numbers.

Original: 1 denotes the all-ones vector.(cid:107) x (cid:107) p denotes the p -norm of x , and (cid:107) x (cid:107) 0 is the number of nonzero entries in x .
Revised: 1 denotes the all-ones vector.(cid:107) x (cid:107) p denotes the p -norm of x , and (cid:107) x (cid:107) 0 is the number of nonzero entries in x .

Original: We write x (cid:12) y todenote elementwise multiplication of x and y .
Revised: We write x (cid:12) y todenote elementwise multiplication of x and y .

Original: The operator Π k ( x ) :=
Revised: The operator Π k ( x ) :=

Original: arg min y : (cid:107) y (cid:107) 0 ≤ k (cid:107) x − y (cid:107) 2denotes
Revised: arg min y : (cid:107) y (cid:107) 0 ≤ k (cid:107) x − y (cid:107) 2denotes

Original: Euclidean projection onto the set of k -sparse vectors.
Revised: Euclidean projection onto the set of k -sparse vectors.

Original:  Related Work
Revised:  Related Work

Original: Neural network pruning.
Revised: Neural network pruning.

Original: Our use of a magnitude-based criterion for sparse learning draws onearly work in the area of neural network pruning [22].
Revised: Our use of a magnitude-based criterion for sparse learning draws onearly work in the area of neural network pruning [22].

Original: Magnitude-based pruning is computationallycheap relative to alternative criteria that rely on ﬁrst- or second-order information [32; 26; 19], and isa perhaps surprisingly performant option despite its simplicity [16].
Revised: Magnitude-based pruning is computationallycheap relative to alternative criteria that rely on ﬁrst- or second-order information [32; 26; 19], and isa perhaps surprisingly performant option despite its simplicity [16].

Original: More generally, Spartan buildson a substantial body of previous work that aims to jointly optimize sparsity patterns and modelparameters for deep neural networks [42; 8; 18; 28; 27; 47, inter alia ; see, e.g. , 21 for a survey].
Revised: More generally, Spartan buildson a substantial body of previous work that aims to jointly optimize sparsity patterns and modelparameters for deep neural networks [42; 8; 18; 28; 27; 47, inter alia ; see, e.g. , 21 for a survey].

Original: Spartan as a generalization of existing methods.
Revised: Spartan as a generalization of existing methods.

Original: We highlight two particularly relevant methods inthe literature: the iterative magnitude pruning (IMP) method from Zhu and Gupta [47] (Algorithm 1),and Top-KAST , or Top- K Always Sparse Training [23] (Algorithm 2).
Revised: We highlight two particularly relevant methods inthe literature: the iterative magnitude pruning (IMP) method from Zhu and Gupta [47] (Algorithm 1),and Top-KAST , or Top- K Always Sparse Training [23] (Algorithm 2).

Original: At the extreme points of thesharpness parameter β , Spartan’s parameter update procedure reduces to those of IMP and Top-KAST.
Revised: At the extreme points of thesharpness parameter β , Spartan’s parameter update procedure reduces to those of IMP and Top-KAST.

Original: We can thus view Spartan as a method that generalizes and smoothly interpolates between this pair ofalgorithms.
Revised: We can thus view Spartan as a method that generalizes and smoothly interpolates between this pair ofalgorithms.

Original: As β → ∞ , Spartan is equivalent to IMP: this approach sparsiﬁes parameters using atop- k magnitude-based binary mask, and entries that were masked in the forward pass receive zerogradient in the backward pass.
Revised: As β → ∞ , Spartan is equivalent to IMP: this approach sparsiﬁes parameters using atop- k magnitude-based binary mask, and entries that were masked in the forward pass receive zerogradient in the backward pass.

Original: At β = 0 , Spartan reduces to Top-KAST: this method again sparsiﬁesparameters by magnitude with a binary mask, but unlike IMP, all entries are updated in the backwardpass using the gradient of the loss with respect to the masked parameters.
Revised: At β = 0 , Spartan reduces to Top-KAST: this method again sparsiﬁesparameters by magnitude with a binary mask, but unlike IMP, all entries are updated in the backwardpass using the gradient of the loss with respect to the masked parameters.

Original: The Top-KAST update isthus an application of the straight-through gradient method [5; 9], otherwise known as lazy projectionor dual averaging (DA) in optimization [33; 40; 2].
Revised: The Top-KAST update isthus an application of the straight-through gradient method [5; 9], otherwise known as lazy projectionor dual averaging (DA) in optimization [33; 40; 2].

Original: Algorithm
Revised: Algorithm

Original: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , step size η t
Revised: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , step size η t

Original: Output: parameters θ t + 1: ˜ θ t =
Revised: Output: parameters θ t + 1: ˜ θ t =

Original: Π k ( θ t ) 2: θ t +

Revised: Π k ( θ t ) 2: θ t +

Original: =
Revised: =

Original: θ t − η t ∇ ˜ θ t ∇ L (˜ θ t ) 3: return θ t +
Revised: θ t − η t ∇ ˜ θ t ∇ L (˜ θ t ) 3: return θ t +

Original: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , step size η t
Revised: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , step size η t

Original: Output: parameters θ t + 1: ˜ θ t =
Revised: Output: parameters θ t + 1: ˜ θ t =

Original: Π k ( θ t ) 2: θ t +

Revised: Π k ( θ t ) 2: θ t +

Original: =
Revised: =

Original: θ t − η t ∇ L (˜ θ t ) 3: return θ t +
Revised: θ t − η t ∇ L (˜ θ t ) 3: return θ t +

Original: Smooth approximations.
Revised: Smooth approximations.

Original: Our use of the soft top- k operation is related to prior methods thatuse the logistic sigmoid function as a differentiable approximation to the step function for sparsetraining [29; 35; 1].
Revised: Our use of the soft top- k operation is related to prior methods thatuse the logistic sigmoid function as a differentiable approximation to the step function for sparsetraining [29; 35; 1].

Original: These approaches similarly regulate the sharpness of the approximation with atemperature parameter that scales the input logits to the sigmoid function.
Revised: These approaches similarly regulate the sharpness of the approximation with atemperature parameter that scales the input logits to the sigmoid function.

Original: A distinctive characteristicof Spartan is that the parameter k directly controls the degree of sparsity of the mask; this differsfrom previous approaches involving the logistic sigmoid approximation that only indirectly controlthe degree of sparsity using an L 1 penalty term.
Revised: A distinctive characteristicof Spartan is that the parameter k directly controls the degree of sparsity of the mask; this differsfrom previous approaches involving the logistic sigmoid approximation that only indirectly controlthe degree of sparsity using an L 1 penalty term.

Original: Transformer-speciﬁc methods.
Revised: Transformer-speciﬁc methods.

Original: Several approaches are specialized to transformer-based architectures.
Revised: Several approaches are specialized to transformer-based architectures.

Original: These include structured pruning of pre-trained transformers for NLP tasks [38; 25; 39] andsparse training methods speciﬁcally applied to vision transformers [7; 6].
Revised: These include structured pruning of pre-trained transformers for NLP tasks [38; 25; 39] andsparse training methods speciﬁcally applied to vision transformers [7; 6].

Original: In contrast, Spartan is ageneral-purpose sparse training algorithm that is designed to be agnostic to the model architecture.
Revised: In contrast, Spartan is ageneral-purpose sparse training algorithm that is designed to be agnostic to the model architecture.

Original: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , sharpness parameter β , step size η t
Revised: Input: parameters θ t , loss function L ( θ ) , sparsity budget k , sharpness parameter β , step size η t

Original: Output: parameters θ t +
Revised: Output: parameters θ t +

Original: [Equation]
Revised: [Equation]

Original: (cid:46) apply soft masking
Revised: (cid:46) apply soft masking

Original: [Equation]
Revised: [Equation]

Original: (cid:46) project onto k -sparse set 3:
Revised: (cid:46) project onto k -sparse set 3:

Original: [Equation]
Revised: [Equation]

Original: ) (cid:46) compute dual averaging update 4: return θ t +
Revised: ) (cid:46) compute dual averaging update 4: return θ t +

Original:  Sparsity via Regularized Transportation
Revised:  Sparsity via Regularized Transportation

Original: Each iteration of Spartan consists of the following two steps (Algorithm 3): (1) an approximatelyk -sparse masking of the model parameters using the soft top- k operator, and (2) dual averagingbased parameter updates with the set of exactly k -sparse vectors as the feasible set.
Revised: Each iteration of Spartan consists of the following two steps (Algorithm 3): (1) an approximatelyk -sparse masking of the model parameters using the soft top- k operator, and (2) dual averagingbased parameter updates with the set of exactly k -sparse vectors as the feasible set.

Original: This updateprocedure aims to combine the advantages of optimization via dual averaging and those of smoothapproximation-based algorithms.
Revised: This updateprocedure aims to combine the advantages of optimization via dual averaging and those of smoothapproximation-based algorithms.

Original: The key advantage of the dual averaging method is that it mitigates the issue of gradient sparsity: in astandard gradient update, only those parameters that were not masked in the forward pass receive anonzero gradient, resulting in slow progress at high sparsities.
Revised: The key advantage of the dual averaging method is that it mitigates the issue of gradient sparsity: in astandard gradient update, only those parameters that were not masked in the forward pass receive anonzero gradient, resulting in slow progress at high sparsities.

Original: Thus, dual averaging iterates are ableto more quickly explore the space of possible sparsity patterns.
Revised: Thus, dual averaging iterates are ableto more quickly explore the space of possible sparsity patterns.

Original: On the other hand, large variations in the sparsity patterns realized post-projection can lead toinstability in optimization, thus hampering the ﬁnal performance of the model; for instance, TopKAST empirically beneﬁts from the addition of sparsity-inducing regularizers on its iterates [23; 36].
Revised: On the other hand, large variations in the sparsity patterns realized post-projection can lead toinstability in optimization, thus hampering the ﬁnal performance of the model; for instance, TopKAST empirically beneﬁts from the addition of sparsity-inducing regularizers on its iterates [23; 36].

Original: This issue motivates our use of a soft top- k approximation as a mechanism for controlling the stabilityof our training iterates.
Revised: This issue motivates our use of a soft top- k approximation as a mechanism for controlling the stabilityof our training iterates.

Original: In the following, we begin by describing the soft top- k masking operation and address the issues ofscalability and of applying soft top-
Revised: In the following, we begin by describing the soft top- k masking operation and address the issues ofscalability and of applying soft top-

Original: k masking to structured sparsity.
Revised: k masking to structured sparsity.

Original: We subsequently discuss howthe update procedure outlined in Algorithm 3 can be incorporated within a complete training loop.
Revised: We subsequently discuss howthe update procedure outlined in Algorithm 3 can be incorporated within a complete training loop.

Original: 3.1 Soft Top- k Masking
Revised: 3.1 Soft Top- k Masking

Original: Our soft top- k masking scheme is based on the soft top- k operator softtopk( v, k, β ) described by
Revised: Our soft top- k masking scheme is based on the soft top- k operator softtopk( v, k, β ) described by

Original: Xie et al.
Revised: Xie et al.

Original: softtopk takes as input a vector v ∈ R d , a budget parameter k > 0 and a sharpnessparameter β ≥ 0 , and outputs a vector m ∈ [0 , 1] d .
Revised: softtopk takes as input a vector v ∈ R d , a budget parameter k > 0 and a sharpnessparameter β ≥ 0 , and outputs a vector m ∈ [0 , 1] d .

Original: By using the magnitudes of the parameter vectorθ as a measure of the “value” of each entry, we obtain a soft top- k magnitude pruning operator σ βk ( · )by masking the entries of the parameter vector with the output of the soft top- k operator:
Revised: By using the magnitudes of the parameter vectorθ as a measure of the “value” of each entry, we obtain a soft top- k magnitude pruning operator σ βk ( · )by masking the entries of the parameter vector with the output of the soft top- k operator:

Original: [Equation]
Revised: [Equation]

Original: We generalize the soft top- k operator described by Xie et al.
Revised: We generalize the soft top- k operator described by Xie et al.

Original: [41] by incorporating a strictly positivecost vector c .
Revised: [41] by incorporating a strictly positivecost vector c .

Original: In particular, we require that the output mask m ∈ [0 , 1] d satisﬁes the budget constraintc T m ≤ k .
Revised: In particular, we require that the output mask m ∈ [0 , 1] d satisﬁes the budget constraintc T m ≤ k .

Original: This abstraction is useful for modeling several notions of heterogeneous costs associatedwith different sparsity patterns within a model.
Revised: This abstraction is useful for modeling several notions of heterogeneous costs associatedwith different sparsity patterns within a model.

Original: For instance, parameters that are repeatedly invokedwithin a network have a higher associated computational cost: Evci et al.
Revised: For instance, parameters that are repeatedly invokedwithin a network have a higher associated computational cost: Evci et al.

Original: [15] observe that the FLOPcount of a ResNet-50 model at a ﬁxed sparsity can vary by more than a factor of 2 × due to variationsin the level of sparsity among convolutional layers with differing output dimensions.
Revised: [15] observe that the FLOPcount of a ResNet-50 model at a ﬁxed sparsity can vary by more than a factor of 2 × due to variationsin the level of sparsity among convolutional layers with differing output dimensions.

Original: To derive the cost-sensitive soft top- k operator, we begin by expressing the optimal mask as thesolution to the following linear program (LP):
Revised: To derive the cost-sensitive soft top- k operator, we begin by expressing the optimal mask as thesolution to the following linear program (LP):

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

Original: k
Revised: k

Original: [Equation]
Revised: [Equation]

Original: Input: values v ∈ R d , costs c ∈ R d + , budget k ∈ (0 , c T 1 ] , sharpness parameter β ≥ 0 , max.
Revised: Input: values v ∈ R d , costs c ∈ R d + , budget k ∈ (0 , c T 1 ] , sharpness parameter β ≥ 0 , max.

Original: iterations T , tolerance (cid:15) , initial dual variable µ
Revised: iterations T , tolerance (cid:15) , initial dual variable µ

Original: Output: mask m ∈ [0 , 1] d , dual variables µ ∈ R , ν ∈ R d 1: z = βv/c (cid:46) normalize values v by costs c 2: m 0 =3: for t = 1 , . .
Revised: Output: mask m ∈ [0 , 1] d , dual variables µ ∈ R , ν ∈ R d 1: z = βv/c (cid:46) normalize values v by costs c 2: m 0 =3: for t = 1 , . .

Original: , T do 4: ν t = log c − log( 1 + exp( z + µ t − 1 ))
Revised: , T do 4: ν t = log c − log( 1 + exp( z + µ t − 1 ))

Original: (cid:46) normalize mask entries 5: µ t = log k − log (cid:80)
Revised: (cid:46) normalize mask entries 5: µ t = log k − log (cid:80)

Original: i exp( z i + ν t,i ) (cid:46) normalize sum to be equal to k 6: m t =
Revised: i exp( z i + ν t,i ) (cid:46) normalize sum to be equal to k 6: m t =

Original: exp( z + µ t + ν t − log c ) (cid:46) compute mask 7: if | v T ( m t − m t − 1 ) | < (cid:15) | v T m t − 1 | then 8: return m t , µ t , ν t 9: return
Revised: exp( z + µ t + ν t − log c ) (cid:46) compute mask 7: if | v T ( m t − m t − 1 ) | < (cid:15) | v T m t − 1 | then 8: return m t , µ t , ν t 9: return

Original: m T , µ T , ν T
Revised: m T , µ T , ν T

Original: Input: gradient w.r.t.
Revised: Input: gradient w.r.t.

Original: outputs g , mask m , costs c , budget k , sharpness parameter β
Revised: outputs g , mask m , costs c , budget k , sharpness parameter β

Original: Output: gradient w.r.t.
Revised: Output: gradient w.r.t.

Original: inputs 1: a = 1 T ( g (cid:12) m (cid:12) ( 1 − m )) 2: b = 1 T ( c (cid:12) m 2 ) 3: γ = a/ ( k − b ) 4: return βm (cid:12) ( 1 − m ) (cid:12) ( g/c − γ ) (cid:46) gradient vanishes when m i = 

Revised: inputs 1: a = 1 T ( g (cid:12) m (cid:12) ( 1 − m )) 2: b = 1 T ( c (cid:12) m 2 ) 3: γ = a/ ( k − b ) 4: return βm (cid:12) ( 1 − m ) (cid:12) ( g/c − γ ) (cid:46) gradient vanishes when m i =

Original: or m i =
Revised: or m i =

Original: Deferring the details to the Appendix, we rewrite this LP in the following form with y :
Revised: Deferring the details to the Appendix, we rewrite this LP in the following form with y :

Original: = c (cid:12) x :maximize y,y (cid:48)
Revised: = c (cid:12) x :maximize y,y (cid:48)

Original: ∈ R d + ( v/c ) T y (2) subject to y + y (cid:48)
Revised: ∈ R d + ( v/c ) T y (2) subject to y + y (cid:48)

Original: [Equation]
Revised: [Equation]

Original: We identify this LP as an optimal transportation problem with the variables ( y, y (cid:48) ) and proﬁt matrixdeﬁned by the normalized values v/c .
Revised: We identify this LP as an optimal transportation problem with the variables ( y, y (cid:48) ) and proﬁt matrixdeﬁned by the normalized values v/c .

Original: Note that when c = 1 , this LP reduces to the original softtop-
Revised: Note that when c = 1 , this LP reduces to the original softtop-

Original: k OT problem from Xie et al.
Revised: k OT problem from Xie et al.

Original: In order to efﬁciently approximate the solution to Problem 2, we add entropic regularization withweight 1 /β and solve the resulting regularized problem using the Sinkhorn-Knopp algorithm [11; 4;12].
Revised: In order to efﬁciently approximate the solution to Problem 2, we add entropic regularization withweight 1 /β and solve the resulting regularized problem using the Sinkhorn-Knopp algorithm [11; 4;12].

Original: Algorithms 4 and 5 describe the forward and backward passes of the cost-sensitive soft top- koperator respectively.
Revised: Algorithms 4 and 5 describe the forward and backward passes of the cost-sensitive soft top- koperator respectively.

Original: The expression for the gradient in Algorithm 5 follows from Theorem 3 in
Revised: The expression for the gradient in Algorithm 5 follows from Theorem 3 in

Original: Xie et al.
Revised: Xie et al.

Original: [41] with some algebraic simpliﬁcation.
Revised: [41] with some algebraic simpliﬁcation.

Original: We remark that this closed-form backward pass isapproximate in the sense that it assumes that we obtain the optimal dual variables µ, ν in the forwardpass; in practice, we do not encounter issues when using approximate values of the dual variableswith tolerance (cid:15)
Revised: We remark that this closed-form backward pass isapproximate in the sense that it assumes that we obtain the optimal dual variables µ, ν in the forwardpass; in practice, we do not encounter issues when using approximate values of the dual variableswith tolerance (cid:15)

Original: = 0 .
Revised: = 0 .

Original: 01 and a maximum iteration count of T = 100 in Algorithm 4.
Revised: 01 and a maximum iteration count of T = 100 in Algorithm 4.

Original: Scalability.
Revised: Scalability.

Original: Since we apply soft top-
Revised: Since we apply soft top-

Original: k masking to high-dimensional parameters in each iteration oftraining, it is important to minimize the computational overhead of Sinkhorn iteration in the forwardpass.
Revised: k masking to high-dimensional parameters in each iteration oftraining, it is important to minimize the computational overhead of Sinkhorn iteration in the forwardpass.

Original: We observe that in order to compute the hard top- k projection step in dual averaging (as in
Revised: We observe that in order to compute the hard top- k projection step in dual averaging (as in

Original: Algorithm 3), it is necessary to ﬁnd the index i k of the k th largest entry of | θ | .
Revised: Algorithm 3), it is necessary to ﬁnd the index i k of the k th largest entry of | θ | .

Original: By using the value− β | θ i k | /c
Revised: By using the value− β | θ i k | /c

Original: i k to initialize the dual variable µ 0 , we can accelerate the convergence of Sinkhorn iteration.
Revised: i k to initialize the dual variable µ 0 , we can accelerate the convergence of Sinkhorn iteration.

Original: Concretely, we demonstrate in Sec.
Revised: Concretely, we demonstrate in Sec.

Original: 4.3 that our implementation of Spartan incurs a per-iterationruntime overhead of approximately 5% over standard dense ResNet-50 training.
Revised: 4.3 that our implementation of Spartan incurs a per-iterationruntime overhead of approximately 5% over standard dense ResNet-50 training.

Original: Structured sparsity.
Revised: Structured sparsity.

Original: To implement block structured sparsity, we instantiate one mask variable perblock and mask all parameters within that block with the same mask variable.
Revised: To implement block structured sparsity, we instantiate one mask variable perblock and mask all parameters within that block with the same mask variable.

Original: To compute the mask,we use the
Revised: To compute the mask,we use the

Original: 3.2 Training with Spartan
Revised: 3.2 Training with Spartan

Original: Training schedule.
Revised: Training schedule.

Original: When training with Spartan, we divide the training process into a warmup phase,an intermediate phase, and a ﬁne-tuning phase.
Revised: When training with Spartan, we divide the training process into a warmup phase,an intermediate phase, and a ﬁne-tuning phase.

Original: In our experiments, the warmup phase consists of theﬁrst 20% of epochs, the intermediate phase the next 60%, and the ﬁne-tuning phase the ﬁnal 20%.
Revised: In our experiments, the warmup phase consists of theﬁrst 20% of epochs, the intermediate phase the next 60%, and the ﬁne-tuning phase the ﬁnal 20%.

Original: Inthe warmup phase, we linearly anneal the global sparsity of the model from a value of 1 at the startof training to the target level of sparsity.
Revised: Inthe warmup phase, we linearly anneal the global sparsity of the model from a value of 1 at the startof training to the target level of sparsity.

Original: Throughout the intermediate phase, we maintain the targetsparsity level, and during the ﬁne-tuning phase, we ﬁx the sparsity mask to that used at the start ofﬁne-tuning.
Revised: Throughout the intermediate phase, we maintain the targetsparsity level, and during the ﬁne-tuning phase, we ﬁx the sparsity mask to that used at the start ofﬁne-tuning.

Original: This training schedule is similar to those used in prior work [36].
Revised: This training schedule is similar to those used in prior work [36].

Original: Exploration vs. exploitation as a function of β .
Revised: Exploration vs. exploitation as a function of β .

Original: From the start of training up to the start ofthe ﬁne-tuning phase, we linearly ramp the sharpness parameter β from an initial value of 1 to theﬁnal value β max .
Revised: From the start of training up to the start ofthe ﬁne-tuning phase, we linearly ramp the sharpness parameter β from an initial value of 1 to theﬁnal value β max .

Original: We interpret the spectrum of updates parameterized by β as being characteristicof an exploration-exploitation tradeoff.
Revised: We interpret the spectrum of updates parameterized by β as being characteristicof an exploration-exploitation tradeoff.

Original: In Figure 3, we illustrate this phenomenon by plotting
Revised: In Figure 3, we illustrate this phenomenon by plotting

Original: Pearson correlation coefﬁcients between sparsity masks at different stages of training.
Revised: Pearson correlation coefﬁcients between sparsity masks at different stages of training.

Original: We observethat iterative magnitude pruning converges on the ﬁnal mask relatively early in training, which isindicative of insufﬁcient exploration of different sparsity patterns—consequently, this model achieveslower validation accuracy than the remaining models.
Revised: We observethat iterative magnitude pruning converges on the ﬁnal mask relatively early in training, which isindicative of insufﬁcient exploration of different sparsity patterns—consequently, this model achieveslower validation accuracy than the remaining models.

Original: The three models trained with Spartan eachramp β from 1 at the start of training to β max at epoch 80.
Revised: The three models trained with Spartan eachramp β from 1 at the start of training to β max at epoch 80.

Original: The value of β max correlates well bothwith the Pearson correlation with the ﬁnal mask, and with the Pearson correlation between masksat the end of subsequent epochs.
Revised: The value of β max correlates well bothwith the Pearson correlation with the ﬁnal mask, and with the Pearson correlation between masksat the end of subsequent epochs.

Original: Spartan thus interpolates between the high-exploration regimeof standard dual averaging and the low-exploration regime of iterative magnitude pruning.
Revised: Spartan thus interpolates between the high-exploration regimeof standard dual averaging and the low-exploration regime of iterative magnitude pruning.

Original: In thisexample, the intermediate setting β max = 10 achieves the highest top-1 validation accuracy of 74 .
Revised: In thisexample, the intermediate setting β max = 10 achieves the highest top-1 validation accuracy of 74 .

Original: 6% ,with Top-KAST at 73 .
Revised: 6% ,with Top-KAST at 73 .

Original: 5% and IMP at 68 . 2% .
Revised: 5% and IMP at 68 . 2% .

Original:  Empirical Evaluation
Revised:  Empirical Evaluation

Original: In this section, we report the results of our sparse training experiments on the ImageNet-1K datasetwith two standard architectures: ResNet-50 [20] and ViT-B/16 [14], consisting of 25.6M and 86.4Mparameters respectively.
Revised: In this section, we report the results of our sparse training experiments on the ImageNet-1K datasetwith two standard architectures: ResNet-50 [20] and ViT-B/16 [14], consisting of 25.6M and 86.4Mparameters respectively.

Original: On ResNet-50, we evaluate only unstructured pruning, whereas on ViT-B/16,we evaluate both unstructured and block structured pruning.
Revised: On ResNet-50, we evaluate only unstructured pruning, whereas on ViT-B/16,we evaluate both unstructured and block structured pruning.

Original: We subsequently present empiricalstudies of the sensitivity of Spartan to the value of β , the effect of running Spartan without the dualaveraging step, and the computational overhead of soft top- k masking.
Revised: We subsequently present empiricalstudies of the sensitivity of Spartan to the value of β , the effect of running Spartan without the dualaveraging step, and the computational overhead of soft top- k masking.

Original: 4.1 ImageNet-1K Classiﬁcation
Revised: 4.1 ImageNet-1K Classiﬁcation

Original: In all our experiments, we run Spartan with the training schedule described in Section 3.2.
Revised: In all our experiments, we run Spartan with the training schedule described in Section 3.2.

Original: We trainand evaluate our models on the ImageNet-1K dataset with the standard training-validation split andreport means and standard deviations over 3 independent trials.
Revised: We trainand evaluate our models on the ImageNet-1K dataset with the standard training-validation split andreport means and standard deviations over 3 independent trials.

Original: We provide additional details on ourexperimental setup in the Appendix.
Revised: We provide additional details on ourexperimental setup in the Appendix.

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

Original: ResNet-50 experimental setup.
Revised: ResNet-50 experimental setup.

Original: For consistency with our baselines, we use standard data augmentation with horizontal ﬂips and random crops at 224 × 224 resolution.
Revised: For consistency with our baselines, we use standard data augmentation with horizontal ﬂips and random crops at 224 × 224 resolution.

Original: For all Spartan runs, we useβ max = 10 , which we selected based on models trained at 95% accuracy.
Revised: For all Spartan runs, we useβ max = 10 , which we selected based on models trained at 95% accuracy.

Original: We sparsify the parametersof all linear and convolutional layers with a global sparsity budget, excluding bias parameters andthe parameters of batch normalization layers.
Revised: We sparsify the parametersof all linear and convolutional layers with a global sparsity budget, excluding bias parameters andthe parameters of batch normalization layers.

Original: Our baselines are iterative magnitude pruning [47],
Revised: Our baselines are iterative magnitude pruning [47],

Original: RigL with the Erdos-Renyi-Kernel (ERK) sparsity distribution [15], Soft Threshold Weight Reparameterization (STR) [24], probabilistic masking (ProbMask) [46], OptG [45], and Top-KAST with
Revised: RigL with the Erdos-Renyi-Kernel (ERK) sparsity distribution [15], Soft Threshold Weight Reparameterization (STR) [24], probabilistic masking (ProbMask) [46], OptG [45], and Top-KAST with

Original: Powerpropagation and ERK [23; 36].
Revised: Powerpropagation and ERK [23; 36].

Original: We additionally re-run the most performant baseline method,
Revised: We additionally re-run the most performant baseline method,

Original: Top-KAST, using a reimplementation in our own codebase.
Revised: Top-KAST, using a reimplementation in our own codebase.

Original: For Top-KAST, we exclude the ﬁrstconvolutional layer from pruning, following [36; 15].
Revised: For Top-KAST, we exclude the ﬁrstconvolutional layer from pruning, following [36; 15].

Original: We use mixed precision training with a batchsize of 4096 on 8 NVIDIA A100 GPUs.
Revised: We use mixed precision training with a batchsize of 4096 on 8 NVIDIA A100 GPUs.

Original: ViT experimental setup.
Revised: ViT experimental setup.

Original: We use the ViT-B architecture with 16 × 16 patches at 224 × 224 inputresolution.
Revised: We use the ViT-B architecture with 16 × 16 patches at 224 × 224 inputresolution.

Original: We augment the training data using RandAugment [10], MixUp [44] and CutMix [43].
Revised: We augment the training data using RandAugment [10], MixUp [44] and CutMix [43].

Original: Our ViT models are trained from random initialization, without any pretraining.
Revised: Our ViT models are trained from random initialization, without any pretraining.

Original: We set β max = 20for Spartan with unstructured sparsity, and β max = 320 and β max = 640 for Spartan with 16 × 16and 32 × 32 blocks respectively.
Revised: We set β max = 20for Spartan with unstructured sparsity, and β max = 320 and β max = 640 for Spartan with 16 × 16and 32 × 32 blocks respectively.

Original: These latter settings are the values of β max for the unstructuredcase scaled up by factors of 16 and 32 —since each block averages the magnitudes of B 2 entries, weexpect the variance of the block values to correspondingly decrease by a factor of B 2 , and we thuscompensate by scaling up β by B .
Revised: These latter settings are the values of β max for the unstructuredcase scaled up by factors of 16 and 32 —since each block averages the magnitudes of B 2 entries, weexpect the variance of the block values to correspondingly decrease by a factor of B 2 , and we thuscompensate by scaling up β by B .

Original: In the block structured case, we exclude the input convolutionallayer and the output classiﬁcation head from pruning since their parameter dimensions are notdivisible by the block size.
Revised: In the block structured case, we exclude the input convolutionallayer and the output classiﬁcation head from pruning since their parameter dimensions are notdivisible by the block size.

Original: We use mixed precision training with a batch size of 4096 on 16 NVIDIA
Revised: We use mixed precision training with a batch size of 4096 on 16 NVIDIA

Original: A100 GPUs across 2 nodes.
Revised: A100 GPUs across 2 nodes.

Original: [Table]
Revised: [Table]

Original: Table 1 lists the top-1 validation accuracies achieved by fully dense ResNet-50 and
Revised: Table 1 lists the top-1 validation accuracies achieved by fully dense ResNet-50 and

Original: ViT-B/16 models.
Revised: ViT-B/16 models.

Original: Table 2 reports validation accuracies for ResNet-50 models at 80%, 90%, 95%,97.5% and 99% sparsity, and Table 3 reports validation accuracies for ViT at 90% sparsity.
Revised: Table 2 reports validation accuracies for ResNet-50 models at 80%, 90%, 95%,97.5% and 99% sparsity, and Table 3 reports validation accuracies for ViT at 90% sparsity.

Original: In the
Revised: In the

Original: Appendix, we additionally report measurements of inference-time FLOP costs for our sparse modelsand the results of experiments with FLOP-sensitive pruning.
Revised: Appendix, we additionally report measurements of inference-time FLOP costs for our sparse modelsand the results of experiments with FLOP-sensitive pruning.

Original: For ResNet-50 models, we ﬁnd that Spartan outperforms all our baselines across all training durationsat sparsity levels of 90% and above.
Revised: For ResNet-50 models, we ﬁnd that Spartan outperforms all our baselines across all training durationsat sparsity levels of 90% and above.

Original: In particular, Spartan achieves a mean top-1 accuracy within1% of fully dense training at 95% sparsity.
Revised: In particular, Spartan achieves a mean top-1 accuracy within1% of fully dense training at 95% sparsity.

Original: We observe that additional epochs of sparse trainingconsistently improves the ﬁnal generalization accuracy; in contrast, validation accuracy peaks at200 epochs for dense training.
Revised: We observe that additional epochs of sparse trainingconsistently improves the ﬁnal generalization accuracy; in contrast, validation accuracy peaks at200 epochs for dense training.

Original: This trend persists at 800 training epochs, where Spartan achieves67 .
Revised: This trend persists at 800 training epochs, where Spartan achieves67 .

Original: 18 ± 0 .
Revised: 18 ± 0 .

Original: 13 top-1 accuracy at 99% sparsity.
Revised: 13 top-1 accuracy at 99% sparsity.

Original: For the Top-KAST baseline, we omit results at 99%sparsity due to training instability.
Revised: For the Top-KAST baseline, we omit results at 99%sparsity due to training instability.

Original: For ViT-B/16, Spartan outperforms Top-KAST for both unstructured and block structured pruning.
Revised: For ViT-B/16, Spartan outperforms Top-KAST for both unstructured and block structured pruning.

Original: We observe a particularly large improvement over Top-KAST in the block structured case, where
Revised: We observe a particularly large improvement over Top-KAST in the block structured case, where

Original: Spartan improves absolute validation accuracy by 7 . 8% for 32 × 32 block structured pruning.
Revised: Spartan improves absolute validation accuracy by 7 . 8% for 32 × 32 block structured pruning.

Original: Forunstructured pruning, Spartan achieves comparable accuracy to SViTE [7] ( 81 .
Revised: Forunstructured pruning, Spartan achieves comparable accuracy to SViTE [7] ( 81 .

Original: 2% for Spartan vs.81 .
Revised: 2% for Spartan vs.81 .

Original: 3% for SViTE), but with 30% higher sparsity ( 90% for Spartan vs. 60% for SViTE).
Revised: 3% for SViTE), but with 30% higher sparsity ( 90% for Spartan vs. 60% for SViTE).

Original: In exchangefor a 0 .
Revised: In exchangefor a 0 .

Original: 6% reduction in accuracy relative to DeiT-B [37], Spartan reduces model storage cost by 10 ×and the FLOP cost of inference by 7 . 4 × .
Revised: 6% reduction in accuracy relative to DeiT-B [37], Spartan reduces model storage cost by 10 ×and the FLOP cost of inference by 7 . 4 × .

Original: 4.2 Sensitivity and Ablation Analysis
Revised: 4.2 Sensitivity and Ablation Analysis

Original: Figure 4 (left) shows the effect of ablating the dual averaging step in Spartan— i.e. , omitting hardtop- k projection in the forward pass—over a range of settings of β max for 95% sparse ResNet-50models trained for 100 epochs.
Revised: Figure 4 (left) shows the effect of ablating the dual averaging step in Spartan— i.e. , omitting hardtop- k projection in the forward pass—over a range of settings of β max for 95% sparse ResNet-50models trained for 100 epochs.

Original: The dashed line shows top-1 accuracy with Top-KAST.
Revised: The dashed line shows top-1 accuracy with Top-KAST.

Original: For Spartantraining without dual averaging, we compute a hard top- k mask at the end of epoch 80, and as withstandard Spartan training, we ﬁx this mask until the end of training.
Revised: For Spartantraining without dual averaging, we compute a hard top- k mask at the end of epoch 80, and as withstandard Spartan training, we ﬁx this mask until the end of training.

Original: In the absence of top- k projection,we ﬁnd that accuracy increases with increasing β max up to β max = 80 ; at lower settings of β max ,the higher approximation error of soft masking is detrimental to the ﬁnal accuracy of the model.
Revised: In the absence of top- k projection,we ﬁnd that accuracy increases with increasing β max up to β max = 80 ; at lower settings of β max ,the higher approximation error of soft masking is detrimental to the ﬁnal accuracy of the model.

Original: Incontrast, the use of top- k projection in the forward pass mitigates this mismatch between training andinference and improves the ﬁnal accuracy of the sparse model.
Revised: Incontrast, the use of top- k projection in the forward pass mitigates this mismatch between training andinference and improves the ﬁnal accuracy of the sparse model.

Original: 4.3 Computational Overhead
Revised: 4.3 Computational Overhead

Original: We evaluate the computational overhead of Spartan over standard dense training by measuring wallclock training times for a ResNet-50 model (Figure 4, right).
Revised: We evaluate the computational overhead of Spartan over standard dense training by measuring wallclock training times for a ResNet-50 model (Figure 4, right).

Original: Our benchmark measures runtime on asingle NVIDIA A100 GPU over 50 iterations of training with a batch size of 256.
Revised: Our benchmark measures runtime on asingle NVIDIA A100 GPU over 50 iterations of training with a batch size of 256.

Original: We use randominputs and labels to avoid incurring data movement costs.
Revised: We use randominputs and labels to avoid incurring data movement costs.

Original: We compare three approaches: standard
Revised: We compare three approaches: standard

Original: Sinkhorn iteration, Sinkhorn iteration with the dual variable µ initialized using its ﬁnal value fromthe previous iteration, and Sinkhorn iteration with µ initialized using the value − β | θ i k | /c
Revised: Sinkhorn iteration, Sinkhorn iteration with the dual variable µ initialized using its ﬁnal value fromthe previous iteration, and Sinkhorn iteration with µ initialized using the value − β | θ i k | /c

Original: i k , wherethe index
Revised: i k , wherethe index

Original: i k of the k th largest entry of | θ | is computed by sorting the vector | θ | in each iteration.
Revised: i k of the k th largest entry of | θ | is computed by sorting the vector | θ | in each iteration.

Original: Standard Sinkhorn iteration incurs higher overheads as β increases—this is due to the additionaliterations required to reach convergence as the regularized OT problem more closely approximates theoriginal LP.
Revised: Standard Sinkhorn iteration incurs higher overheads as β increases—this is due to the additionaliterations required to reach convergence as the regularized OT problem more closely approximates theoriginal LP.

Original: We ﬁnd that dual caching and sorting both prevent this growth in runtime over the rangeof β values that we tested.
Revised: We ﬁnd that dual caching and sorting both prevent this growth in runtime over the rangeof β values that we tested.

Original: In our remaining experiments, we use the simple sorting-based approachsince it corresponds to the lowest relative overhead over standard dense training (approximately 5%).
Revised: In our remaining experiments, we use the simple sorting-based approachsince it corresponds to the lowest relative overhead over standard dense training (approximately 5%).

Original: We note that this relative overhead decreases as the batch size increases since the cost of computingthe mask is independent of the batch size.
Revised: We note that this relative overhead decreases as the batch size increases since the cost of computingthe mask is independent of the batch size.

Original:  Discussion
Revised:  Discussion

Original: As an alternative to magnitude-based pruning, we may also apply Spartan in conjunction withlearned value parameters , as in methods such as Movement Pruning [34].
Revised: As an alternative to magnitude-based pruning, we may also apply Spartan in conjunction withlearned value parameters , as in methods such as Movement Pruning [34].

Original: In this approach, wewould compute masks using a set of auxiliary parameters φ ∈ R d instead of the magnitudes | θ | :σ βk ( θ ; φ ) = θ (cid:12) softtopk( φ, k, β ) , and similarly for hard projection.
Revised: In this approach, wewould compute masks using a set of auxiliary parameters φ ∈ R d instead of the magnitudes | θ | :σ βk ( θ ; φ ) = θ (cid:12) softtopk( φ, k, β ) , and similarly for hard projection.

Original: We remark that while this variantrequires additional memory during training to store the value parameters, there is no additional costduring inference since the sparsity mask is ﬁxed at the end of training.
Revised: We remark that while this variantrequires additional memory during training to store the value parameters, there is no additional costduring inference since the sparsity mask is ﬁxed at the end of training.

Original: Limitations.
Revised: Limitations.

Original: Since Spartan retains a dense parameter vector and computes dense backward passesduring training, it incurs higher memory and computational costs in each iteration than sparse-tosparse methods like RigL [15].
Revised: Since Spartan retains a dense parameter vector and computes dense backward passesduring training, it incurs higher memory and computational costs in each iteration than sparse-tosparse methods like RigL [15].

Original: Nevertheless, we note that in terms of total computational or energycost over a full training run, Spartan may remain a competitive option as it requires fewer iterationsof training to reach a given accuracy threshold relative to sparse-to-sparse methods.
Revised: Nevertheless, we note that in terms of total computational or energycost over a full training run, Spartan may remain a competitive option as it requires fewer iterationsof training to reach a given accuracy threshold relative to sparse-to-sparse methods.

Original: However, we donot compare total training FLOP costs in our empirical evaluation.
Revised: However, we donot compare total training FLOP costs in our empirical evaluation.

Original: A further limitation is that cost-sensitive pruning with Spartan is only compatible with relatively crudelinear cost models.
Revised: A further limitation is that cost-sensitive pruning with Spartan is only compatible with relatively crudelinear cost models.

Original: This restriction arises due to the requirements of the regularized OT formulationused to compute the soft top- k mask.
Revised: This restriction arises due to the requirements of the regularized OT formulationused to compute the soft top- k mask.

Original: In particular, this limitation precludes the use of cost modelsinvolving interaction terms such as those arising from spatially coherent sparsity patterns.
Revised: In particular, this limitation precludes the use of cost modelsinvolving interaction terms such as those arising from spatially coherent sparsity patterns.

Original: Societal Impacts.
Revised: Societal Impacts.

Original: Inference with deep neural network models is a computationally intensive process.
Revised: Inference with deep neural network models is a computationally intensive process.

Original: At present, the total energy footprint associated with serving these models in production is expectedto continue growing in tandem with the rising prevalence of large transformer-based architecturesin vision and NLP applications.
Revised: At present, the total energy footprint associated with serving these models in production is expectedto continue growing in tandem with the rising prevalence of large transformer-based architecturesin vision and NLP applications.

Original: Research towards improving the energy efﬁciency of deep neuralnetworks is therefore an important counterbalancing force against increasing resource usage bythese models.
Revised: Research towards improving the energy efﬁciency of deep neuralnetworks is therefore an important counterbalancing force against increasing resource usage bythese models.

Original: The development of sparse learning algorithms is particularly relevant to these efforts,and we expect that the impact of these approaches will further increase as sparsity-aware hardwareacceleration becomes more widely available.
Revised: The development of sparse learning algorithms is particularly relevant to these efforts,and we expect that the impact of these approaches will further increase as sparsity-aware hardwareacceleration becomes more widely available.

Original:  Conclusions & Future Work
Revised:  Conclusions & Future Work

Original: In this work, we describe a sparse learning algorithm that interpolates between two parameter updateschemes: standard stochastic gradient updates with hard masking and the dual averaging method.
Revised: In this work, we describe a sparse learning algorithm that interpolates between two parameter updateschemes: standard stochastic gradient updates with hard masking and the dual averaging method.

Original: We show that there exists an intermediate regime between these two methods that yields improvedgeneralization accuracy for sparse convolutional and transformer models, particularly at higher levelsof sparsity.
Revised: We show that there exists an intermediate regime between these two methods that yields improvedgeneralization accuracy for sparse convolutional and transformer models, particularly at higher levelsof sparsity.

Original: While we have demonstrated promising empirical results with our proposed method, thelearning dynamics of stochastic optimization for deep networks under sparsity constraints remainsrelatively poorly understood from a theoretical standpoint.
Revised: While we have demonstrated promising empirical results with our proposed method, thelearning dynamics of stochastic optimization for deep networks under sparsity constraints remainsrelatively poorly understood from a theoretical standpoint.

Original: There thus exists ample opportunity forfurther work towards better understanding sparse learning algorithms, which may in turn inspirefuture algorithmic advances in this area.
Revised: There thus exists ample opportunity forfurther work towards better understanding sparse learning algorithms, which may in turn inspirefuture algorithmic advances in this area.

Original: Checklist
Revised: A Derivation of the Optimal Transportation LP

Revised: Here, we show how the original top- k LP with costs c can be straightforwardly rewritten in the formof an optimal transportation problem. For a given value vector v ∈ R d , cost vector c (cid:31) 0 and budgetk , the top- k LP is:maximize [Equation]

Original: For all authors...
Revised: R [Equation] subject to 0 (cid:22) x

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (cid:22) 1 , c T x = k. Deﬁne the variables y as y : = c (cid:12) x and substitute to obtain:maximize y ∈ R ( v/c ) T y subject to 0 (cid:22) y (cid:22) c, 1 T y =

Original: contributions and scope?
Revised: k. Now eliminate the upper bound constraint by introducing additional variables y (cid:48) to give:maximize y,y (cid:48) [Equation] + ( v/c )

Original: [Yes]
Revised: [Equation]

Original: (b) Did you describe the limitations of your work?
Revised: y subject to y + y (cid:48) = c, 1 T y =

Revised: k, 1 T y (cid:48)

Original: [Yes] See Section 5.
Revised: = 1 T c [Equation]

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: k, which can be recognized as an optimal transportation problem in the variables y and y (cid:48) . B Training Details In Tables 4 and 5, we detail the hyperparameters used for our training runs

Original: [Yes] See
Revised: [Table]

Original: Section 5.
Revised: [Table] [Table] C FLOP Measurements Due to differences in the computational cost associated with individual parameters, the sparsityfraction does not map 1-to-1 to the fraction of FLOPs required for inference.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: Tables 6 and 7give FLOP costs for our sparse models as a percentage of the FLOP cost of the correspondingdense model. We performed our FLOP measurements using the open source tool available athttps://github.com/sovrasov/flops-counter.pytorch/ . We count multiply and add operations as one FLOP each. There is some inconsistency in the literatureregarding this convention, with some prior work using multiply-accumulate (MAC) counts and FLOP counts interchangeably.

Original: them?
Revised: To convert the base FLOP counts listed below for ResNet-50 and ViT-B/16to MACs, we can simply divide the given counts by 2. In Table 6, the ResNet-50 FLOP counts for Top-KAST are slightly higher than those for Spartan. This is due to the exclusion of the input convolutional layer from pruning in the case of Top-KAST. [Table]

Original: [Yes]
Revised: [Table]

Original: If you are including theoretical results...
Revised: We demonstrate FLOP-sensitive pruning with Spartan on ResNet-50 using the following cost model:assign a cost of 1 to each parameter of a fully connected layer, and a cost of N 2 to each parameter ofa convolutional layer where the output has size N × N along its spatial dimensions.

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: We evaluatetwo valuation functions: v 1 ( c i , θ i ) = c i | θ i | and v 0 . 5 ( c i , θ i ) = √ c i | θ i | .

Original: [N/A]
Revised: v 1 results in the same pruningorder as in standard pruning, but with a FLOP budget constraint instead of the usual sparsity budget.v 0 . 5 assigns a lower value to the parameters of convolutional layers, and results in networks wherethe parameters of convolutional layers are preferentially pruned. We use β max = 10 for v 1 andβ max = 160 for v 0 .

Original: (b) Did you include complete proofs of all theoretical results?
Revised: 5 to compensate for the relatively smaller scale of the normalized values v/c

Original: [N/A]3.
Revised: inthe soft top- k forward pass (Algorithm 4). Table 8 gives the top-1 accuracy, FLOP percentage, and sparsity percentage for each of these valuationfunctions. Spartan yields models with identical FLOP percentages of 5 .

Original: If you ran experiments...
Revised: 76% , which is slightly higherthan the budgeted value of 5% —this discrepancy is due to the additional cost of the normalizationlayers and activation functions in the network. Most notably, there is a substantial difference in thesparsity percentages realized by these valuation functions. As expected, v 0 . 5 preferentially sparsiﬁesthe parameters of convolutional layers and yields denser fully connected layers, resulting in lowersparsity overall.

Original: (a) Did you include the code, data, and instructions needed to reproduce the main exper-
Revised: [Table] E Additional Experiments In Table [Table]

Original: imental results (either in the supplemental material or as a URL)?
Revised: we compare Spartan against two additional variants of the Top-KAST baseline: TopKAST with the Erdos-Renyi-Kernel (ERK) sparsity distribution, and with pruning applied to theparameters of all convolutional and fully-connected layers with the exception of bias terms (pruneall). Top-KAST (excl. input conv.) denotes the Top-KAST variant used in the experiments presentedin the main text, where we exclude the input convolutional layer from pruning.

Revised: We ﬁnd that there issome small variation in the measured top-1 validation accuracies, but our conclusion that Spartanimproves generalization at higher levels of sparsity is unchanged

Original: [No] We did not
Revised: [Table] [Table]

Revised: F Learned Sparsity Patterns

Original: include code alongside our submission since code release has not yet been approved;
Revised: We observe a qualitative difference in the distribution of per-layer sparsities between ViT-B/16models trained with unstructured sparsity and those trained with block structured sparsity (Figure 5). In particular, the output projections of self-attention layers under block structured pruning aresigniﬁcantly more dense in the later blocks of the network relative to unstructured pruning. Thereasons for this difference are not immediately clear to us, and we leave further investigation of thisphenomenon to future work. Block structured pruning produces coherent sparsity patterns in ViT-B/16 models.

Original: however, we intend to eventually open source our code on GitHub.
Revised: In Figure 6,we visualize the magnitudes of the weight matrices corresponding to the input projection of each self-attention layer in a ViT-B/16 model trained with Spartan using 32 × 32 block structured pruning. This matrix maps vectors of dimension 768 to query, key, and value embedding vectors, each ofdimension 768 . We observe that the training process yields similar sparsity patterns in the query andkey embedding submatrices, which correspond to the left and center panels in the visualization foreach layer. This is an intuitively reasonable property since the self-attention layer computes innerproducts of the query and key embeddings in order to construct attention maps.

Original: (b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
Revised: We note that thissymmetry emerges purely as a result of the optimization process; we did not incorporate any priorknowledge into Spartan regarding the role of particular entries of the weight matrices subject tosparsiﬁcation. G Derivation of Soft Top- k Backward Pass (Algorithm 5) The gradient of the loss with respect to the input values v of the soft top- k function follows from

Original: [Yes] See Section 4 and the Appendix.
Revised: Theorem 3 of [41] with some algebraic manipulation.

Original: (c) Did you report error bars (e
Revised: We restate the theorem below with the notationused in this paper. In the following, let ¯ q := q : − 1 denote the vector q with the last entry removed, andlet ¯ Y :

Original: g., withrespect to the random seed after running experiments multiple times)?
Revised: = Y · , : − 1 denote the matrix Y with the last column removed.

Original: [Yes]
Revised: Theorem [41].

Original: (d) Did you include the total amount ofcompute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: Let the solution Y ∈ R N × M of the entropic regularized optimal transport problemwith cost matrix C , row marginals p and column marginals q be given by:

Original: [Yes]
Revised: [Equation]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: where ν ∗ ∈ R N , and µ ∗ ∈ R M are optimal dual variables. Then dν ∗ dC and dµ ∗ dC are given by:

Original: [Yes]
Revised: [Equation]

Original: (b) Did you mention the license of the assets? [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: where − H − 1 D ∈ R ( N + M − 1) × N × M , 0 ∈ R 1 × N × M , and D (cid:96)ij = β (cid:26) δ (cid:96)i Y ij

Original: [N/A]
Revised: [Equation]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: (cid:96) = 1 , . .

Original: [N/A]
Revised: N [Equation]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
Revised: [Equation] N

Original: [N/A]
Revised: [Equation]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: Specializing this theorem to our setting with N = d and M = 2 , and recalling that the soft top- kmask values are given by m i = Y

Original: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
Revised: i 1 /c i , we obtain: [Equation] Therefore:

Original: [N/A]
Revised: [Equation]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: Taking the derivative of the mask with respect to the ﬁrst column of the cost matrix, we have:

Original: [N/A]
Revised: [Equation]

Original: (c) Did you include the estimated hourly wage paid to participants and  the total amount spent on participant compensation?
Revised: Finally, using our deﬁnition of the cost matrix as C · 1 = − v/c and applying the loss gradientg := (cid:0) dLdm (cid:1) T , we obtain:

Original: [N/A]
Revised: [Equation]

Revised: as desired.