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

Original: Abstract
Revised: Abstract

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

Original: Spartan is based on a combination of two techniques: (1) soft top- k masking of low-magnitude parameters via a regularized optimal transportation problem and (2) dual averaging-based parameter updates with hard sparsiﬁcation in the forward pass.
Revised: Spartan is based on a combination of two techniques: (1) soft top- k masking of low-magnitude parameters via a regularized optimal transportation problem and (2) dual averaging-based parameter updates with hard sparsiﬁcation in the forward pass.

Original: This scheme realizes an explorationexploitation tradeoff: early in training, the learner is able to explore various sparsity patterns, and as the soft top- k approximation is gradually sharpened over the course of 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 sparsity patterns, and as the soft top- k approximation is gradually sharpened over the course of training, the balance shifts towards parameter optimization with respect to a ﬁxed sparsity mask.

Original: Spartan is sufﬁciently ﬂexible to accommodate a variety of sparsity allocation policies, including both unstructured and block structured sparsity, as well as general cost-sensitive sparsity allocation mediated by linear models of per-parameter costs.
Revised: Spartan is sufﬁciently ﬂexible to accommodate a variety of sparsity allocation policies, including both unstructured and block structured sparsity, as well as general cost-sensitive sparsity allocation mediated by linear models 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 incurring absolute 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 incurring absolute 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 retaining only a small fraction of non-zero entries.
Revised: Sparse learning algorithms search for model parameters that minimize training loss while retaining only 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 [33].
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 [33].

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

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

Original: This work focuses on improving the generalization error of sparse neural network models.
Revised: This work focuses on improving the generalization error of sparse neural network models.

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

Original: Spartan belongs to the family of “parameter dense” training algorithms that maintains a dense parameter vector θ ∈ R d throughout training [48; 24], in contrast to “parameter sparse” training algorithms that adhere to a ˜ O ( k ) memory budget for representing the parameters of a ksparse model [3; 31; 32; 13; 16].
Revised: Spartan belongs to the family of “parameter dense” training algorithms that maintains a dense parameter vector θ ∈ R d throughout training [48; 24], in contrast to “parameter sparse” training algorithms that adhere to a ˜ O ( k ) memory budget for representing the parameters of a ksparse model [3; 31; 32; 13; 16].

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

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

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

Original: On the other hand, the goal of exploitation is for the learner to optimize model 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 optimize model 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 the optimization process.
Revised: avoids frequent oscillations between disparate sparsity patterns, which may be detrimental to the optimization process.

Original: We operationalize this intuition through the use of a differentiable soft top- k masking operation σ βk : R d → R d .
Revised: We operationalize this intuition through the use of a differentiable soft top- k masking operation σ βk : R d → R d .

Original: This function maps parameters θ to an approximately sparse output that suppresses low-magnitude entries in θ .
Revised: This function maps parameters θ to an approximately sparse output that suppresses low-magnitude 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- k magnitude-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- k magnitude-based selection (Figure 1 (a, b)).

Original: Soft top-
Revised: Soft top-

Original: k masking therefore constrains the iterates θ t to be close to the set of exactly k -sparse vectors, with the strength of this constraint mediated by the parameter β .
Revised: k masking therefore constrains the iterates θ t to be close to the set of exactly k -sparse vectors, with the strength of this constraint mediated by the parameter β .

Original: Figures 1 (c) and 2 give some geometeric intuition for the effect of this mapping.
Revised: Figures 1 (c) and 2 give some geometeric intuition for the effect of this mapping.

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

Original: k masking using a regularized optimal transportation formulation [11; 42] and demonstrate that this technique scales to networks on the order of 10 8 parameters in size.
Revised: k masking using a regularized optimal transportation formulation [11; 42] and demonstrate that this technique scales to networks on the order of 10 8 parameters in size.

Original: We evaluate Spartan on ResNet-50 [21] and ViT [15] models trained on the ImageNet-1K dataset.
Revised: We evaluate Spartan on ResNet-50 [21] and ViT [15] 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 accuracies than 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 accuracies than 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 at 97.5% sparsity, improving on the previous state-of-the-art by 0 .
Revised: 2% accuracy at 97.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 a 0 .
Revised: ViT-B/16 models reduce model storage size by 10 × and inference FLOPs by 7 . 4 × at the cost of a 0 .

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

Original: We further demonstrate that Spartan is effective for block structured pruning, a form of structured sparsity that is more amenable to acceleration on current GPU hardware than unstructured sparsity [18].
Revised: We further demonstrate that Spartan is effective for block structured pruning, a form of structured sparsity that is more amenable to acceleration on current GPU hardware than unstructured sparsity [18].

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

Original: 1% top-1 accuracy at 90% sparsity, compared to the baseline of 74 .
Revised: 1% top-1 accuracy at 90% sparsity, compared to the baseline of 74 .

Original: 1% at the same sparsity level.
Revised: 1% at the same sparsity level.

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 [48] and Top-
Revised: We show that Spartan interpolates between two existing sparse learning algorithms: iterative magnitude pruning [48] and Top-

Original: K Always Sparse Training [24].
Revised: K Always Sparse Training [24].

Original: • We empirically evaluate Spartan using ResNet-50 and ViT models trained on the ImageNet-1K dataset, demonstrating consistent improvements over the prior state-of-the-art.
Revised: • We empirically evaluate Spartan using ResNet-50 and ViT models trained on the ImageNet-1K dataset, demonstrating consistent improvements over the prior state-of-the-art.

Original: • We study the effect of Spartan’s hyperparameters on an exploration-exploitation tradeoff during training and on the ﬁnal accuracy of the trained models.
Revised: • We study the effect of Spartan’s hyperparameters on an exploration-exploitation tradeoff during training and on the ﬁnal accuracy of the trained models.

Original: We provide an open source implementation of Spartan at https://github.com/ facebookresearch/spartan .
Revised: We provide an open source implementation of Spartan at https://github.com/ facebookresearch/spartan .

Original: Notation.
Revised: Notation.

Original: We use R + and R ++ to denote the nonnegative real numbers and the strictly positive real numbers respectively.
Revised: We use R + and R ++ to denote the nonnegative real numbers and the strictly positive real numbers respectively.

Original: 1 d is the all-ones vector of dimension d .
Revised: 1 d is the all-ones vector of dimension d .

Original: Given a vector x ∈ R d , (cid:107) x (cid:107) p denotes
Revised: Given a vector x ∈ R d , (cid:107) x (cid:107) p denotes

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 on early work in the area of neural network pruning [23].
Revised: Our use of a magnitude-based criterion for sparse learning draws on early work in the area of neural network pruning [23].

Original: Magnitude-based pruning is computationally cheap relative to alternative criteria that rely on ﬁrst- or second-order information [33; 27; 20], and is a perhaps surprisingly performant option despite its simplicity [17].
Revised: Magnitude-based pruning is computationally cheap relative to alternative criteria that rely on ﬁrst- or second-order information [33; 27; 20], and is a perhaps surprisingly performant option despite its simplicity [17].

Original: More generally, Spartan builds on a substantial body of previous work that aims to jointly optimize sparsity patterns and model parameters for deep neural networks [43; 8; 19; 29; 28; 48, inter alia ; see, e.g. , 22 for a survey].
Revised: More generally, Spartan builds on a substantial body of previous work that aims to jointly optimize sparsity patterns and model parameters for deep neural networks [43; 8; 19; 29; 28; 48, inter alia ; see, e.g. , 22 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 sparse training methods in the literature: the iterative magnitude pruning (IMP) method (Algorithm 1, [48]), and Top- K Always Sparse Training , or Top-KAST (Algorithm 2, [24]).
Revised: We highlight two particularly relevant sparse training methods in the literature: the iterative magnitude pruning (IMP) method (Algorithm 1, [48]), and Top- K Always Sparse Training , or Top-KAST (Algorithm 2, [24]):

Original: At the extreme points of the sharpness 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 of algorithms.

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

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

Original: The Top-KAST update is thus an application of the straight-through gradient method [5; 9], otherwise known as lazy projection or dual averaging (DA) in optimization [34; 41; 2; 14].

Original: See Sec.

Original: 5 for further discussion of this connection.

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: 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 ∇ L (˜ θ t ) 3: return θ t +
Revised: θ t − η t ∇ L (˜ θ t ) 3: return θ t +

Original: Note that both Algorithms 1 and 2 compute the forward pass through the network using the sparsiﬁed parameters ˜ θ , obtained by setting all but the top- k entries of θ by magnitude to zero.
Revised: Note that both Algorithms 1 and 2 compute the forward pass through the network using the sparsiﬁed parameters ˜ θ , obtained by setting all but the top- k entries of θ by magnitude to zero.

Original: Smooth approximations.
Revised: Algorithms 1 and 2 differ in the presence of the ∇ ˜ θ t term in line 2. This is a d × d diagonal matrix where ( ∇ ˜ θ t ) ii = iff θ t,i was not set to zero by the projection Π k , and ( ∇ ˜ θ t ) ii = 0 otherwise. At the extreme points of the sharpness parameter β , Spartan’s parameter update procedure reduces to those of IMP and Top-KAST.

Original: Our use of the soft top- k operation is related to prior methods that use the logistic sigmoid function as a differentiable approximation to  the step function for sparse training [30; 36; 1]. These approaches similarly regulate the sharpness of the approximation with a temperature parameter that scales the input logits to the sigmoid function.
Revised: We can thus view Spartan as a method that generalizes and smoothly interpolates between this pair of algorithms. As β → ∞ , Spartan is equivalent to IMP: this approach sparsiﬁes parameters using a top- k magnitude-based binary mask, and entries that were masked in the forward pass receive zero gradient in the backward pass. At β = 0 , Spartan reduces to Top-KAST: this method again sparsiﬁes parameters by magnitude with a binary mask, but unlike IMP, all entries are updated in the backward pass using the gradient of the loss with respect to the masked parameters.

Original: A distinctive characteristic Input: parameters θ t , loss function L ( θ ) , sparsity budget k , sharpness parameter β , step size η t
Revised: The Top-KAST update is thus an application of the straight-through gradient method [5; 9], otherwise known as lazy projection or dual averaging (DA) in optimization [34; 41; 2; 14]. See Sec.

Original: Output: parameters θ t + 1: σ βk ( θ t ) =
Revised: 5 for further discussion on this point.

Original: θ t (cid:12) softtopk( | θ t | , k, β ) (cid:46) apply soft masking 2: ˜ θ t =
Revised: Smooth approximations. Our use of the soft top- k operation is related to prior methods that use the logistic sigmoid function as a differentiable approximation to the step function for sparse training [30; 36; 1].

Original: Π k (cid:16) σ βk ( θ t ) (cid:17) (cid:46) project onto k -sparse set 3: θ t +1 = θ t − η t ∇ σ βk ( θ t ) ∇ L (˜ θ t ) (cid:46) compute dual averaging update 4: return θ t + of Spartan is that the parameter k directly controls the degree of sparsity of the mask; this differs from previous approaches involving the logistic sigmoid approximation that only indirectly control the degree of sparsity using an L 1 penalty term.
Revised: These approaches similarly regulate the sharpness of the approximation with a temperature parameter that scales the input logits to the sigmoid function. A distinctive characteristic of Spartan is that the parameter k directly controls the degree of sparsity of the mask; this differs from previous approaches involving the logistic sigmoid approximation that only indirectly control the 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 [39; 26; 40] and sparse training methods speciﬁcally applied to vision transformers [7; 6].
Revised: These include structured pruning of pre-trained transformers for NLP tasks [39; 26; 40] and sparse training methods speciﬁcally applied to vision transformers [7; 6].

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

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

Original: We begin by motivating our proposed approach.
Revised: We begin by motivating our proposed approach.

Original: A key advantage of the dual averaging method (Algorithm 2) over IMP (Algorithm 1) is that it mitigates the issue of gradient sparsity.
Revised: A key advantage of the dual averaging method (Algorithm 2) over IMP (Algorithm 1) is that it mitigates the issue of gradient sparsity.

Original: In the IMP update, only those parameters that were not masked in the forward pass receive a nonzero gradient, resulting in slow progress at high levels of sparsity.
Revised: In the IMP update, only those parameters that were not masked in the forward pass receive a nonzero gradient, resulting in slow progress at high levels of sparsity.

Original: In contrast, dual averaging applies a dense update to the parameter vector θ in each iteration—intuitively, these dense updates can help “revive” parameters that are below the top- k magnitude threshold and are consequently masked to zero.
Revised: In contrast, dual averaging applies a dense update to the parameter vector θ in each iteration—intuitively, these dense updates can help “revive” parameters that are below the top- k magnitude threshold and are consequently masked to zero.

Original: Dual averaging iterates are thus able to more quickly explore the space of possible sparsity patterns.
Revised: Dual averaging iterates are thus able to 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 to instability in optimization, thus hampering the ﬁnal performance of the model.
Revised: On the other hand, large variations in the sparsity patterns realized post-projection can lead to instability in optimization, thus hampering the ﬁnal performance of the model.

Original: For instance, TopKAST empirically beneﬁts from the application of additional sparsity-inducing regularization on the parameters [24 (Sec. 2.3), 37].
Revised: For instance, TopKAST empirically beneﬁts from the application of additional sparsity-inducing regularization on the parameters [24 (Sec. 2.3), 37].

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

Original: Each iteration of Spartan consists of the following two steps (Algorithm 3): (1) an approximately k -sparse masking of the model parameters using the soft top- k operator, and (2) dual averaging-based 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 approximately k -sparse masking of the model parameters using the soft top- k operator, and (2) dual averaging-based parameter updates with the set of exactly k -sparse vectors as the feasible set.

Original: This procedure aims to combine the advantages of Algorithms 1 and 2 within a single parameter update scheme.
Revised: This procedure aims to combine the advantages of Algorithms 1 and 2 within a single parameter update scheme.

Original: Note that the Spartan parameter update is essentially identical to Algorithm 2, save for the inclusion of the intermediate soft top-
Revised: Note that the Spartan parameter update is essentially identical to Algorithm 2, save for the inclusion of the intermediate soft top-

Original: k masking step.
Revised: k masking step.

Revised: Algorithm

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

Revised: Output: parameters θ t + 1: σ βk

Revised: [Equation]

Original: In the following, we begin by describing the softtop- k masking operation and address the issues of scalability and of applying soft top-
Revised: [Equation] (cid:46) apply soft masking 2: ˜ [Equation] =

Revised: [Equation]

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

Revised: [Equation]

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

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

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

Original: We subsequently discuss how the update procedure outlined in Algorithm 3 can be incorporated within a complete training loop.
Revised: We then discuss how the 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 sharpness parameter β ≥ 0 , and outputs m ∈ [0 , 1] d , a smoothed version of the top- k indicator vector.
Revised: softtopk takes as input a vector v ∈ R d , a budget parameter k > 0 and a sharpness parameter β ≥ 0 , and outputs m ∈ [0 , 1] d , a smoothed version of the top- k indicator vector.

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 introduce a mild generalization of the soft top- k operator from Xie et al. [42] by incorporating a strictly positive cost vector c . In particular, we require that the output mask m ∈ [0 , 1] d satisﬁes the budget constraint c T m = k .This abstraction is useful for modeling several notions of heterogeneous
Revised: softtopk( | θ | , k, β ) .

Original: Input: values v ∈ R d , costs c ∈ R d ++ , budget k ∈ (0 , 1 Td c ] , sharpness parameter β ≥ 0 , max.
Revised: Input: values v ∈ R d , costs c ∈ R d ++ , budget k ∈ (0 , 1 Td c ] , 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, m 0 = 1 d (cid:46) normalize values v by costs c 2: for t = 1 , . .
Revised: Output: mask m ∈ [0 , 1] d , dual variables µ ∈ R , ν ∈ R d 1: z = βv/c, m 0 = 1 d (cid:46) normalize values v by costs c 2: for t = 1 , . .

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

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

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

Original: exp( z + µ t + ν t − log c ) (cid:46) compute mask 6: if | v T ( m t − m t − 1 ) | < (cid:15) | v T m t − 1 | then (cid:46) check stopping criterion 7:
Revised: exp( z + µ t + ν t − log c ) (cid:46) compute mask 6: if | v T ( m t − m t − 1 ) | < (cid:15) | v T m t − 1 | then (cid:46) check stopping criterion 7:

Original: return m t , µ t , ν t 8: return m T , µ T , ν T
Revised: return m t , µ t , ν t 8: return m T , µ T , ν T

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

Original: outputs g ∈ R d , mask m ∈ [0 , 1] d , costs c , budget k , sharpness parameter β
Revised: outputs g ∈ R d , mask m ∈ [0 , 1] d , costs c , budget k , sharpness parameter β

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

Original: inputs 1: a  c  m  i 2: return 
Revised: inputs 1: a 1 = (cid:80) i g i m i (1 − m i ) , a 2 = (cid:80) i c i m 2 i 2: return βm (cid:12) ( 1 d − m ) (cid:12) ( g/c − a 1 / ( k − a 2 ))

Original: )) parameter costs within a model; for instance, parameters that are repeatedly invoked within a network have a higher associated computational cost.
Revised: We introduce a mild generalization of the soft top- k operator from Xie et al. [42] by incorporating a strictly positive cost vector c . In particular, we require that the output mask m ∈ [0 , 1] d satisﬁes the budget constraint c T m = k . This abstraction is useful for modeling several notions of heterogeneous parameter costs within a model; for instance, parameters that are repeatedly invoked within a network have a higher associated computational cost.

Original: As a concrete example, Evci et al.
Revised: As a concrete example, Evci et al.

Original: [16] observed that the FLOP count of a ResNet-50 model at a ﬁxed global sparsity can vary by a factor of over 2 × due to differences in the individual sparsities of convolutional layers with differing output dimensions.
Revised: [16] observed that the FLOP count of a ResNet-50 model at a ﬁxed global sparsity can vary by a factor of over 2 × due to differences in the individual sparsities of convolutional layers with differing output dimensions.

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

Original: [Equation]
Revised: [Equation]

Original: Deferring the derivation to the Appendix, we can rewrite this LP as the following optimal transportation (OT) problem:
Revised: Deferring the derivation to the Appendix, we can rewrite this LP as the following optimal transportation (OT) problem:

Original: [Equation]
Revised: [Equation]

Original: with the cost matrix C :
Revised: with the cost matrix C :

Original: =
Revised: =

Original: [ − v/c, 0] .
Revised: [ − v/c, 0] .

Original: We can recover the mask from Y as m i
Revised: We can recover the mask from Y as m i

Original: = Y
Revised: = Y

Original: i 1 /c i .
Revised: i 1 /c i .

Original: Now following [42; 11], we introduce entropic regularization to obtain a smoothed version of the top- k operator:
Revised: Now following [42; 11], we introduce entropic regularization to obtain a smoothed version of the top- k operator, with the degree of smoothing controlled by β :

Original: [Equation]
Revised: [Equation]

Original: where H ( Y ) :
Revised: where H ( Y ) :

Original: = − (cid:80) ij Y ij (log Y ij − 1)  .
Revised: = − (cid:80) ij Y ij (log Y ij − 1) is the entropy of Y .

Original: In this standard form, we can efﬁciently solve this regularized OT problem using the Sinkhorn-Knopp algorithm [11; 4; 12].
Revised: In this standard form, we can efﬁciently solve this regularized OT problem using the Sinkhorn-Knopp algorithm [11; 4; 12].

Original: Algorithms 4 and 5 describe the forward and backward passes of the resulting cost-sensitive soft top- k operator (see Appendices B and C for their derivations).
Revised: Algorithms 4 and 5 describe the forward and backward passes of the resulting cost-sensitive soft top- k operator (see Appendices B and C for their derivations).

Original: The expression for the gradient in
Revised: The expression for the gradient in

Original: Algorithm 5 follows from Theorem 3 in Xie et al.
Revised: Algorithm 5 follows from Theorem 3 in Xie et al.

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

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

Original: = 0 .
Revised: = 0 .

Original: 01 and a maximum iteration count of
Revised: 01 and a maximum iteration count of

Original: T = 100 in Algorithm 4.
Revised: 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 of training, it is important to minimize the computational overhead of Sinkhorn iteration in the forward pass.
Revised: k masking to high-dimensional parameters in each iteration of training, it is important to minimize the computational overhead of Sinkhorn iteration in the forward pass.

Original: We observe that in order to compute the hard top- k projection step in dual averaging (as in Algorithm 3), it is necessary to ﬁnd the index i k of the k th largest entry of | θ |  .
Revised: We note that in order to compute the hard top- k projection step in dual averaging (as in Algorithm 3), it is necessary to ﬁnd the index i k of the k th largest entry of | θ | /c .

Original: By using the heuristic initialization µ 0 = − β | θ i k | /c
Revised: Given the index i k , we can use the heuristic initialization µ 0 = − β | θ i k | /c

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

Original: This heuristic is based on the fact that | θ i k | /c
Revised: This heuristic is based on the observation that | θ i k | /c

Original: i k is the threshold value of the normalized value vector | θ | /c in hard top- k masking.
Revised: i k is the threshold value of the normalized value vector | θ | /c in hard top- k masking.

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

Original: 4.3 that our implementation of Spartan incurs a per-iteration runtime overhead of approximately 5% over standard dense ResNet-50 training.
Revised: 4.3 that our implementation of Spartan incurs a per-iteration runtime 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 per block and mask all parameters within that block with the same mask variable.
Revised: To implement block structured sparsity, we instantiate one mask variable per block and mask all parameters within that block with the same mask variable.

Original: To compute the mask, we use the sum of the magnitudes of the entries in each block as the corresponding value v i .
Revised: To compute the mask, we use the sum of the magnitudes of the entries in each block as the corresponding value v i .

Original: In a standard pruning setting with uniform costs across all parameters, the corresponding cost c i is simply the total number of entries in the block.
Revised: In a standard pruning setting with uniform costs across all parameters, the corresponding cost c i is simply the total number of entries in the block.

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: In the warmup phase, we linearly anneal the global sparsity of the model from a value of 1 at the start of training to the target level of sparsity.
Revised: In the warmup phase, we linearly anneal the global sparsity of the model from a value of 1 at the start of training to the target level of sparsity.

Original: Throughout the intermediate phase, we maintain the target sparsity 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 target sparsity 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 [37].
Revised: This training schedule is similar to those used in prior work [37].

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 of the ﬁ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 of the ﬁ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 characteristic of an exploration-exploitation tradeoff.
Revised: We interpret the spectrum of updates parameterized by β as being characteristic of 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 observe that iterative magnitude pruning converges on the ﬁnal mask relatively early in training, which is indicative of insufﬁcient exploration of different sparsity patterns—consequently, this model achieves lower validation accuracy than the remaining models.
Revised: We observe that iterative magnitude pruning converges on the ﬁnal mask relatively early in training, which is indicative of insufﬁcient exploration of different sparsity patterns—consequently, this model achieves lower validation accuracy than the remaining models.

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

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

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

Original: In this example, the intermediate setting β max = 10 achieves the highest top-1 validation accuracy of 74 .
Revised: In this example, 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: [Table] [Table]
Revised: [Table] [Table]

Original:  Empirical Evaluation
Revised:  Empirical Evaluation

Original: In this section
Revised: In this section

Original: [Table]
Revised: [Table]

Original: we report the results of our sparse training experiments on the ImageNet-1K dataset with two standard architectures: ResNet-50 [21] and ViT-B/16 [15], consisting of 25.6M and 86.4M parameters respectively.
Revised: we report the results of our sparse training experiments on the ImageNet-1K dataset with two standard architectures: ResNet-50 [21] and ViT-B/16 [15], consisting of 25.6M and 86.4M parameters respectively.

Original: On ResNet-50, we evaluate only unstructured pruning
Revised: On ResNet-50, we evaluate only unstructured pruning

Original: [Table]
Revised: [Table]

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

Original: We subsequently present empirical studies of the sensitivity of Spartan to the value of β , the effect of running Spartan without the dual averaging step
Revised: We subsequently present empirical studies of the sensitivity of Spartan to the value of β , the effect of running Spartan without the dual averaging step

Original: [Table]
Revised: [Table]

Original: and the computational overhead of soft top- k masking.
Revised: 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
Revised: In all our experiments

Original: [Table]
Revised: [Table]

Original: we run Spartan with the training schedule described in Section 3.2.
Revised: we run Spartan with the training schedule described in Section 3.2.

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

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

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 parameters of all linear and convolutional layers with a global sparsity budget, excluding bias parameters and the parameters of batch normalization layers.
Revised: We sparsify the parameters of all linear and convolutional layers with a global sparsity budget, excluding bias parameters and the parameters of batch normalization layers.

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

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

Original: Powerpropagation and ERK [24; 37].
Revised: Powerpropagation and ERK [24; 37].

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

Original: [Table]
Revised: [Table]

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 ﬁrst convolutional layer from pruning (following [37; 16]) and we use fully dense backward passes ( i.e. , with a backward sparsity of 0), since this setting achieved the highest accuracy in prior work [24].
Revised: For Top-KAST, we exclude the ﬁrst convolutional layer from pruning (following [37; 16]) and we use fully dense backward passes ( i.e. , with a backward sparsity of 0), since this setting achieved the highest accuracy in prior work [24].

Original: We use mixed precision training with a batch size of 4096 on 8 NVIDIA A100 GPUs.
Revised: We use mixed precision training with a batch size 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 input resolution.
Revised: We use the ViT-B architecture with 16 × 16 patches at 224 × 224 input resolution.

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

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 =for Spartan with unstructured sparsity, and β max = 320 and β max = 640 for Spartan with 16 ×and 32 × 32 blocks respectively.
Revised: We set β max =for Spartan with unstructured sparsity, and β max = 320 and β max = 640 for Spartan with 16 ×and 32 × 32 blocks respectively.

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

Original: In the block structured case, we exclude the input convolutional layer and the output classiﬁcation head from pruning since their parameter dimensions are not divisible by the block size.
Revised: In the block structured case, we exclude the input convolutional layer and the output classiﬁcation head from pruning since their parameter dimensions are not divisible 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: Results.
Revised: Results.

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 report additional measurements of inference FLOP costs for our sparse models and the results of experiments with FLOP-sensitive pruning.
Revised: Appendix, we report additional measurements of inference FLOP costs for our sparse models and the results of experiments with FLOP-sensitive pruning.

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

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

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

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

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: We note that the accuracy improvements in the Spartan-trained
Revised: We note that the accuracy improvements in the Spartan-trained

Original: ResNet-50 models relative to Top-KAST are not a result of increased FLOP counts at a given level of sparsity, as can be seen from the FLOP measurements reported in Appendix E.
Revised: ResNet-50 models relative to Top-KAST are not a result of increased FLOP counts at a given level of sparsity, as can be seen from the FLOP measurements reported in Appendix E.

Original: For ViT-B/16 models, Spartan outperforms Top-KAST for both unstructured and block structured pruning.
Revised: For ViT-B/16 models, 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 Spartan improves absolute validation accuracy by 7 . 8% for 32 × 32 block structured pruning.
Revised: We observe a particularly large improvement over Top-KAST in the block structured case, where Spartan improves absolute validation accuracy by 7 . 8% for 32 × 32 block structured pruning.

Original: For unstructured pruning, Spartan achieves comparable accuracy to SViTE [7] ( 81 . 2% for Spartan vs.
Revised: For unstructured pruning, Spartan achieves comparable accuracy to SViTE [7] ( 81 . 2% for Spartan vs.

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 exchange for a 0 .
Revised: In exchange for a 0 .

Original: 6% reduction in accuracy relative to DeiT-B [38], 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 [38], 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 hard top- k projection in the forward pass—over a range of settings of β max for 95% sparse ResNetmodels trained for 100 epochs.
Revised: Figure 4 (left) shows the effect of ablating the dual averaging step in Spartan— i.e. , omitting hard top- k projection in the forward pass—over a range of settings of β max for 95% sparse ResNetmodels 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 Spartan training without dual averaging, we compute a hard top- k mask at the end of epoch 80, and as with standard Spartan training, we ﬁx this mask until the end of training.
Revised: For Spartan training without dual averaging, we compute a hard top- k mask at the end of epoch 80, and as with standard 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 , thehigher detrimental
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 detrimental

Original: 4.3 Computational Overhead
Revised: 4.3 Computational Overhead

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

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

Original: We use random inputs and labels to avoid incurring data movement costs.
Revised: We use random inputs 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 from the previous iteration (“dual caching”), 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 from the previous iteration (“dual caching”), and Sinkhorn iteration with µ initialized using the value − β | θ i k | /c

Original: i k , where we compute the index
Revised: i k , where we compute the index

Original: i k of the k th largest entry of the normalized value vector | θ | /c by sorting its entries in each iteration (“sorting”).
Revised: i k of the k th largest entry of the normalized value vector | θ | /c by sorting its entries in each iteration (“sorting”).

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

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

Original: In our remaining experiments, we use the simple sorting-based approach since it corresponds to the lowest relative overhead over standard dense training (approximately 5%).
Revised: In our remaining experiments, we use the simple sorting-based approach since 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 computing the mask is independent of the batch size.
Revised: We note that this relative overhead decreases as the batch size increases since the cost of computing the mask is independent of the batch size.

Revised: Spartan is therefore best suited to training with large batch sizes since we amortize the cost of mask updates over the size of the batch.

Original:  Discussion
Revised:  Discussion

Original: Extensions.
Revised: Extensions.

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

Original: In this approach, we would compute masks using a set of auxiliary parameters φ ∈ R d instead of the magnitudes | θ | : σ βk ( θ ; φ ) =
Revised: In this approach, we would compute masks using a set of auxiliary parameters φ ∈ R d instead of the magnitudes | θ | : σ βk ( θ ; φ ) =

Original: θ (cid:12) softtopk( φ, k, β ) , and similarly for hard projection.
Revised: θ (cid:12) softtopk( φ, k, β ) , and similarly for hard projection.

Original: We remark that while this variant requires additional memory during training to store the value parameters, there is no additional cost during inference since the sparsity mask is ﬁxed at the end of training.
Revised: We remark that while this variant requires additional memory during training to store the value parameters, there is no additional cost during 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 passes during training, it incurs higher memory and computational costs in each iteration thansparse-tosparse methods like RigL [16].
Revised: Since Spartan retains a dense parameter vector and computes dense backward passes during training, it incurs higher memory and computational costs in each iteration than methods like RigL [16] that use both a sparse parameters and sparse backward passes.

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

Original: However, we do not compare total training FLOP costs in our empirical evaluation.
Revised: However, we do not 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 crude linear cost models.
Revised: A further limitation is that cost-sensitive pruning with Spartan is only compatible with relatively crude linear cost models of the form c T m , where c is the cost vector and m is the mask.

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

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

Original: For example, the cost models that we consider here cannot encode the notion that coherently pruning an entire row or column of a parameter matrix will yield more cost savings than pruning random entries of the matrix.
Revised: For example, the cost models that we consider here cannot encode the notion that coherently pruning an entire row or column of a parameter matrix will yield more cost savings than pruning random entries of the matrix.

Original: Optimization with Dual Averaging.
Revised: Optimization with Dual Averaging.

Original: As discussed in Section 3, Spartan and Top-KAST are both applications of the dual averaging method [34; 41].
Revised: As discussed in Section 3, Spartan and Top-KAST are both applications of the dual averaging method [34; 41].

Original: The empirically observed effectiveness of dual averaging in sparse training is reminiscent of its success in a related domain where continuous optimization is complicated by discrete constraints: neural network quantization.
Revised: The empirically observed effectiveness of dual averaging in sparse training is reminiscent of its success in a related domain where continuous optimization is complicated by discrete constraints: neural network quantization.

Original: In this area, the
Revised: In this area, the

Original: BinaryConnect algorithm
Revised: BinaryConnect algorithm

Original: [9], which trains binary neural networks with parameters in {± 1 } d , is considered to be a foundational method.
Revised: [9], which trains binary neural networks with parameters in {± 1 } d , is considered to be a foundational method.

Original: As observed by Bai et al.
Revised: As observed by Bai et al.

Original: [2], BinaryConnect is itself a special case of dual averaging, and the theoretical underpinnings of this connection to dual averaging were further analyzed by Dockhorn et al.
Revised: [2], BinaryConnect is itself a special case of dual averaging, and the theoretical underpinnings of this connection to dual averaging were further analyzed by Dockhorn et al.

Original: These parallels with neural network quantization suggest that similar ideas may well apply towards deepening our understanding of sparse learning.
Revised: These parallels with neural network quantization suggest that similar ideas may well apply towards deepening our understanding of sparse learning.

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 expected to continue growing in tandem with the rising prevalence of large transformer-based architectures in vision and NLP applications.
Revised: At present, the total energy footprint associated with serving these models in production is expected to continue growing in tandem with the rising prevalence of large transformer-based architectures in vision and NLP applications.

Original: Research towards improving the energy efﬁciency of deep neural networks is therefore an important counterbalancing force against increasing resource usage by these models.
Revised: Research towards improving the energy efﬁciency of deep neural networks is therefore an important counterbalancing force against increasing resource usage by these 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 hardware acceleration 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 hardware acceleration 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 update schemes: 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 update schemes: 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 improved generalization accuracy for sparse convolutional and transformer models, particularly at higher levels of sparsity.
Revised: We show that there exists an intermediate regime between these two methods that yields improved generalization accuracy for sparse convolutional and transformer models, particularly at higher levels of sparsity.

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

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

Original: Acknowledgements and Disclosure of Funding
Revised: Acknowledgements and Disclosure of Funding

Original: We are grateful to Trevor Gale for his feedback on this work.
Revised: We are grateful to Trevor Gale for his feedback on this work.

Original: We also thank our anonymous reviewers for their valuable suggestions on improving the manuscript.
Revised: We also thank our anonymous reviewers for their valuable suggestions on improving the manuscript.

Original: This work was funded by Meta.
Revised: This work was funded by Meta.

Original: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See Section 5.

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See Section 5.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A] (b) Did you include complete proofs of all theoretical results?

Original: [N/A] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See Section 4 and the Appendix.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] 4.

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?

Original: [Yes] (b) Did you mention the license of the assets?

Original: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

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

Original: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A] A Derivation of the Optimal Transport LP (Problem 2)
Revised: A Derivation of the Optimal Transport LP (Problem 2)

Original: Here, we show how the original top- k LP with costs c can be straightforwardly rewritten in the form of an optimal transportation problem.
Revised: Here, we show how the original top- k LP with costs c can be straightforwardly rewritten in the form of an optimal transportation problem.

Original: For a given value vector v ∈ R d , cost vector c ∈ R d ++ , and budget k , the top- k LP is: maximize m ∈ R d v T m subject to 0 (cid:22) m (cid:22) 1 d , c T m = k.
Revised: For a given value vector v ∈ R d , cost vector c ∈ R d ++ , and budget k , the top- k LP is: maximize m ∈ R d v T m subject to 0 (cid:22) m (cid:22) 1 d , c T m = k.

Original: Deﬁne y :
Revised: Deﬁne y :

Original: = c (cid:12) m and substitute to obtain: maximize
Revised: = c (cid:12) m and substitute to obtain: maximize

Original: y ∈ R d ( v/c ) T y subject to 0 (cid:22) y (cid:22) c, 1 T y = k.
Revised: y ∈ R d ( v/c ) T y subject to 0 (cid:22) y (cid:22) c, 1 Td y = k.

Original: Now eliminate the upper bound constraint by introducing additional variables y (cid:48) to give: maximize y,y (cid:

Revised: Now eliminate the upper bound constraint by introducing additional variables y (cid:48) to give: maximize y,y (cid:

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

Original: =
Revised: =

Original: c, 1 Td y = k.
Revised: c, 1 Td y = k.

Original: Finally, we rewrite this problem in the form of a standard OT problem by introducing the variable
Revised: Finally, we rewrite this problem in the form of a standard OT problem by introducing the variable

Original: Y :=
Revised: Y :=

Original: [ y, y (cid:48) ] and the cost matrix C :=
Revised: [ y, y (cid:48) ] and the cost matrix C :=

Original: [ − v/c, 0] to yield: minimize
Revised: [ − v/c, 0] to yield: [Equation]

Original: Y ∈ R
Revised: Y [Equation]

Original: d × 2 + (cid:88) ij C ij Y ij subject to Y 1 2 = c, 1 d Y =
Revised: =

Original: [ k, 1 Td c − k ] , where we obtain the second column constraint by combining y + y (cid:48)
Revised: c, [Equation] where we obtain the second column constraint by combining y + y (cid:48)

Original: =
Revised: =

Original: c and 1 Td y = k .
Revised: c and 1 Td y = k .

Original: B Derivation of Soft Top- k Forward Pass (Algorithm 4)
Revised: B Derivation of Soft Top- k Forward Pass (Algorithm 4)

Original: Problem 3 is an entropy regularized optimal transport problem with cost matrix C :=
Revised: Problem 3 is an entropy regularized optimal transport problem with cost matrix C :=

Original: [ − v/c, 0] ∈ R d ×, row marginals c ∈ R d , and column marginals [ k, 1 Td c − k ] .
Revised: [ − v/c, 0] ∈ R d ×, row marginals c ∈ R d , and column marginals [ k, 1 Td c − k ] .

Original: By Lemma 2 in [11], the optimal solution to this problem can be written in the form diag(exp ν ) exp( − βC )diag(exp[ µ, µ (cid:48) ]) , with dual variables µ, µ (cid:48) ∈ R and ν ∈ R d .
Revised: By Lemma 2 in [11], the optimal solution to this problem can be written in the form diag(exp ν ) exp( − βC )diag(exp[ µ, µ (cid:48) ]) , with dual variables µ, µ (cid:48) ∈ R and ν ∈ R d .

Original: Moreover, we can compute a sequence of iterates converging to an optimal collection of dual variables using Sinkhorn iteration.
Revised: Moreover, we can compute a sequence of iterates converging to an optimal collection of dual variables using Sinkhorn iteration.

Original: Note that the dual variables µ, µ (cid:48) , and ν are only unique up to an additive constant: for any µ, µ (cid:48) , ν, and δ , we have diag(exp ν ) exp( − βC )diag(exp[ µ, µ (cid:48) ])
Revised: Note that the dual variables µ, µ (cid:48) , and ν are only unique up to an additive constant: for any µ, µ (cid:48) , ν, and δ , we have diag(exp ν ) exp( − βC )diag(exp[ µ, µ (cid:48) ])

Original: =
Revised: =

Original: diag(exp( ν + δ )) exp( − βC )diag(exp[ µ − δ, µ (cid:48) − δ ]) .
Revised: diag(exp( ν + δ )) exp( − βC )diag(exp[ µ − δ, µ (cid:48) − δ ]) .

Original: We use this degree of freedom to ﬁx µ (cid:48)
Revised: We use this degree of freedom to ﬁx µ (cid:48)

Original: = 0 .
Revised: = 0 .

Original: This yields the following Sinkhorn updates:
Revised: This yields the following Sinkhorn updates:

Original: [Equation]
Revised: [Equation]

Original: as stated in Algorithm 4.
Revised: as stated in Algorithm 4.

Revised: Note that eliminating µ (cid:48) reduces the cost of computing the forward pass by roughly 1 / 3 .

Original: C Derivation of Soft Top- k Backward Pass (Algorithm 5)
Revised: C Derivation of Soft Top- k Backward Pass (Algorithm 5)

Original: The gradient of the loss with respect to the input values v of the soft top- k function follows from
Revised: The gradient of the loss with respect to the input values v of the soft top- k function follows from

Original: Theorem 3 of [42] with some algebraic manipulation.
Revised: Theorem 3 of [42] with some algebraic manipulation.

Original: We restate the theorem below with the notation used in this paper.
Revised: We restate the theorem below with the notation used in this paper.

Original: In the following, let ¯ q :=
Revised: In the following, let ¯ q :=

Original: q : − 1 denote the vector q with the last entry removed, and let ¯ Y :
Revised: q : − 1 denote the vector q with the last entry removed, and let ¯ Y :

Original: = Y · , : − 1 denote the matrix Y with the last column removed.
Revised: = Y · , : − 1 denote the matrix Y with the last column removed.

Original: Theorem [42].
Revised: Theorem [42].

Original: Let the solution Y ∈ R N × M of the entropy regularized optimal transport problem with cost matrix C , row marginals p and column marginals q be given by:
Revised: Let the solution Y ∈ R N × M of the entropy regularized optimal transport problem with cost matrix C , row marginals p and column marginals q be given by:

Original: [Equation]
Revised: [Equation]

Original: where ν ∗ ∈ R N , and µ ∗ ∈ R M are optimal dual variables.
Revised: where ν ∗ ∈ R N , and µ ∗ ∈ R M are optimal dual variables.

Original: Then dν ∗ dC and dµ ∗ dC are given by:
Revised: Then dν ∗ dC and dµ ∗ dC are given by:

Original: [Equation]
Revised: [Equation]

Original: where − H − 1 D ∈ R ( N + M − 1) × N × M , 0 ∈ R 1 × N × M , and
Revised: where − H − 1 D ∈ R ( N + M − 1) × N × M , 0 ∈ R 1 × N × M , and

Original: D (cid:96)ij = β (cid:26) δ (cid:96)i Y ij
Revised: D (cid:96)ij = β (cid:26) δ (cid:96)i Y ij

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: N
Revised: N

Original: [Equation]
Revised: [Equation]

Original: ij
Revised: ij

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Specializing this theorem to our setting with N = d and M = 2 , and recalling that the soft top- k mask values are given by m i
Revised: Specializing this theorem to our setting with N = d and M = 2 , and recalling that the soft top- k mask values are given by m i

Original: = Y
Revised: = Y

Original: i 1 /c
Revised: i 1 /c

Original: i , we obtain:
Revised: i , we obtain:

Original: [Equation]
Revised: [Equation]

Original: Therefore:
Revised: Therefore:

Original: [Equation]
Revised: [Equation]

Original: Taking the derivative of the mask with respect to the ﬁrst column of the cost matrix
Revised: Taking the derivative of the mask with respect to the ﬁrst column of the cost matrix

Original: [Equation]
Revised: [Equation]

Original: we have:
Revised: we have:

Original: [Equation]
Revised: [Equation]

Original: Finally, using our deﬁnition of the cost matrix as C · 1 = − v/c and applying the loss gradient g := (cid:0) dLdm (cid:1) T , we obtain:
Revised: Finally, using our deﬁnition of the cost matrix as C · 1 = − v/c and applying the loss gradient g := (cid:0) dLdm (cid:1) T , we obtain:

Original: [Equation]
Revised: [Equation]

Original: as desired.
Revised: as desired.

Original: D Training Details
Revised: D Training Details

Original: In Tables 4 and 5, we detail the hyperparameters used for our training runs.
Revised: In Tables 4 and 5, we detail the hyperparameters used for our training runs.

Original: These hyperparameters were derived from https://github.com/pytorch/vision/blob/
Revised: These hyperparameters were derived from https://github.com/pytorch/vision/blob/

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

Original: E FLOP Measurements
Revised: E FLOP Measurements

Original: Due to differences in the computational cost associated with individual parameters
Revised: Due to differences in the computational cost associated with individual parameters

Original: [Table]
Revised: [Table]

Original: the sparsity fraction does not map 1-to-1 to the fraction of FLOPs required for inference.
Revised: the sparsity fraction does not map 1-to-1 to the fraction of FLOPs required for inference.

Original: Tables 6 andgive FLOP costs for our sparse models as a percentage of the FLOP cost of the corresponding dense model.
Revised: Tables 6 andgive FLOP costs for our sparse models as a percentage of the FLOP cost of the corresponding dense model.

Original: We performed our FLOP measurements using the open source tool available at https://github.com/sovrasov/flops-counter.pytorch/ .
Revised: We performed our FLOP measurements using the open source tool available at https://github.com/sovrasov/flops-counter.pytorch/ .

Original: We count multiply and add operations as one FLOP each.
Revised: We count multiply and add operations as one FLOP each.

Original: We remark that there exists some inconsistency in the literature regarding this convention, with some prior work using multiplyaccumulate (MAC) counts and FLOP counts interchangeably.
Revised: We remark that there exists some inconsistency in the literature regarding this convention, with some prior work using multiplyaccumulate (MAC) counts and FLOP counts interchangeably.

Original: To convert the base FLOP counts listed below for ResNet-50 and ViT-B/16 to MACs, we can simply divide the given counts by 2.
Revised: To convert the base FLOP counts listed below for ResNet-50 and ViT-B/16 to MACs, we can simply divide the given counts by 2.

Original: In Table 6, the ResNet-50 FLOP counts for Top-KAST are slightly higher than those for Spartan, possibly due to the exclusion of the input convolutional layer from pruning in the case of Top-KAST.
Revised: In Table 6, the ResNet-50 FLOP counts for Top-KAST are slightly higher than those for Spartan, possibly due to the exclusion of the input convolutional layer from pruning in the case of Top-KAST.

Original: In particular, this demonstrates that the accuracy improvements seen in Spartan-trained models over those trained using Top-KAST do not correlate with an increase in their FLOP costs.
Revised: In particular, this demonstrates that the accuracy improvements seen in Spartan-trained models over those trained using Top-KAST do not correlate with an increase in their FLOP costs.

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

Original: F FLOP-Sensitive Pruning
Revised: F FLOP-Sensitive Pruning

Original: 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 of a convolutional layer where the output has size N × N along its spatial dimensions.
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 of a convolutional layer where the output has size N × N along its spatial dimensions.

Original: We evaluate two valuation functions: v 1 ( c i , θ i ) =
Revised: We evaluate two valuation functions: v 1 ( c i , θ i ) =

Original: c i | θ i | and v 0 .
Revised: c i | θ i | and v 0 .

Original: 5 ( c i , θ i ) =
Revised: 5 ( c i , θ i ) =

Original: √ c i | θ i | .
Revised: √ c i | θ i | .

Original: v 1 results in the same pruning order as in standard pruning, but with a FLOP budget constraint instead of the usual sparsity budget.
Revised: v 1 results in the same pruning order as in standard pruning, but with a FLOP budget constraint instead of the usual sparsity budget.

Original: v 0 .
Revised: v 0 .

Original: 5 assigns a lower value to the parameters of convolutional layers, and results in networks where the parameters of convolutional layers are preferentially pruned.
Revised: 5 assigns a lower value to the parameters of convolutional layers, and results in networks where the parameters of convolutional layers are preferentially pruned.

Original: We use β max = 10 for v 1 and β max = 160 for v 0 .
Revised: We use β max = 10 for v 1 and β max = 160 for v 0 .

Original: 5 to compensate for the relatively smaller scale of the normalized values v/c in the soft top- k forward pass (Algorithm 4).
Revised: 5 to compensate for the relatively smaller scale of the normalized values v/c in the soft top- k forward pass (Algorithm 4).

Original: Table 8 gives the top-1 accuracy, FLOP percentage, and sparsity percentage for each of these valuation functions.
Revised: Table 8 gives the top-1 accuracy, FLOP percentage, and sparsity percentage for each of these valuation functions.

Original: Spartan yields models with identical FLOP percentages of 5 .
Revised: Spartan yields models with identical FLOP percentages of 5 .

Original: 76% , which is slightly higher than the budgeted value of 5% —this discrepancy is due to the additional cost of the normalization layers and activation functions in the network.
Revised: 76% , which is slightly higher than the budgeted value of 5% —this discrepancy is due to the additional cost of the normalization layers and activation functions in the network.

Original: Most notably, there is a substantial difference in the sparsity percentages realized by these valuation functions.
Revised: Most notably, there is a substantial difference in the sparsity percentages realized by these valuation functions.

Original: As expected, v 0 .
Revised: As expected, v 0 .

Original: 5 preferentially sparsiﬁes the parameters of convolutional layers and yields denser fully connected layers, resulting in lower sparsity overall.
Revised: 5 preferentially sparsiﬁes the parameters of convolutional layers and yields denser fully connected layers, resulting in lower sparsity overall.

Original: [Table]
Revised: [Table]

Original: G Additional Experiments
Revised: G Additional Experiments

Original: In Table 9, 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 the parameters of all convolutional and fully-connected layers with the exception of bias terms (prune all).
Revised: In Table 9, 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 the parameters of all convolutional and fully-connected layers with the exception of bias terms (prune all).

Original: Top-KAST (excl.
Revised: Top-KAST (excl.

Original: input conv.)
Revised: input conv.)

Original: denotes the Top-KAST variant used in the experiments presented in the main text, where we exclude the input convolutional layer from pruning.
Revised: denotes the Top-KAST variant used in the experiments presented in the main text, where we exclude the input convolutional layer from pruning.

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

Original: [Table]
Revised: [Table]

Original: H Learned Sparsity Patterns
Revised: H Learned Sparsity Patterns

Original: We observe a qualitative difference in the distribution of per-layer sparsities between ViT-B/ models trained with unstructured sparsity and those trained with block structured sparsity (Figure 5).
Revised: We observe a qualitative difference in the distribution of per-layer sparsities between ViT-B/ models trained with unstructured sparsity and those trained with block structured sparsity (Figure 5).

Original: In particular, the output projections of self-attention layers under block structured pruning are signiﬁcantly more dense in the later blocks of the network relative to unstructured pruning.
Revised: In particular, the output projections of self-attention layers under block structured pruning are signiﬁcantly more dense in the later blocks of the network relative to unstructured pruning.

Original: The reasons for this difference are not immediately clear to us, and we leave further investigation of this phenomenon to future work.
Revised: The reasons for this difference are not immediately clear to us, and we leave further investigation of this phenomenon to future work.

Original: Block structured pruning produces coherent sparsity patterns in ViT-B/16 models.
Revised: Block structured pruning produces coherent sparsity patterns in ViT-B/16 models.

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

Original: This matrix maps vectors of dimension 768 to query, key, and value embedding vectors, each of dimension 768 .
Revised: This matrix maps vectors of dimension 768 to query, key, and value embedding vectors, each of dimension 768 .

Original: We observe that the training process yields similar sparsity patterns in the query and key embedding submatrices, which correspond to the left and center panels in the visualization for each layer.
Revised: We observe that the training process yields similar sparsity patterns in the query and key embedding submatrices, which correspond to the left and center panels in the visualization for each layer.

Original: This is an intuitively reasonable property since the self-attention layer computes inner products of the query and key embeddings in order to construct attention maps.
Revised: This is an intuitively reasonable property since the self-attention layer computes inner products of the query and key embeddings in order to construct attention maps.

Original: We note that this symmetry emerges purely as a result of the optimization process; we did not incorporate any prior knowledge into Spartan regarding the role of particular entries of the weight matrices subject to sparsiﬁcation.
Revised: We note that this symmetry emerges purely as a result of the optimization process; we did not incorporate any prior knowledge into Spartan regarding the role of particular entries of the weight matrices subject to sparsiﬁcation.