Back Razor: Memory-Efficient Transfer Learning by Self-Sparsified Backpropogation
Anonymous Author(s)
AffiliationAddress email
Abstract
Transfer learning from the model trained on large datasets to customized downstream tasks has been widely used as the pre-trained model can greatly boost thegeneralizability.
However, the increasing sizes of pre-trained models also lead to aprohibitively large memory footprints for downstream transferring, making themunaffordable for personal devices.
Previous work recognizes the bottleneck of thefootprint to be the activation, and hence proposes various solutions such as injectingspecific lite modules.
In this work, we present a novel and general frameworkcalled Back Razor , that can be plug-and-play applied to any pre-trained networkwithout changing its architecture.
The key idea of Back Razor is asymmetricsparsifying : pruning the activation stored for back-propagation, while keeping theforward activation dense.
It is based on the observation that the stored activationthat dominates the memory footprint is only employed on back-propagation.
Suchasymmetric pruning avoids affecting the forward computation, thus making moreaggressive pruning possible.
Furthermore, we conduct the theoretical analysisfor the convergence rate of Back Razor, showing that under mild conditions, ourmethod retains a similar convergence rate as vanilla SGD.
Extensive experimentson both Convolutional Neural Networks and Vision Transformers show that Back
Razor could yield up to 97% sparsity, saving 9.2x memory usage, without losingaccuracy.
The code is included in supplementary materials.
Introduction
Edge devices equipped with deep learning applications are becoming ubiquitous.
However, onlydeploying pre-trained models is not sufficient: these edge devices may keep collecting novel data thatgeneral pre-trained models have never seen before and fine-tuning on them would notably improvethe model’s performance.
A typical approach is to upload these newly collected data to cloud serversfor fine-tuning and download the trained model from the cloud to local devices.
However, sucha way has potential privacy issues due to data sensitivity and also imposes undesirable pressureson the Internet bandwidth, brought by transmitting both the collected data and trained models.
These challenges motivate researchers to explore on-device learning methods, i.e., how to fine-tunepre-trained models on memory-limited devices.
One of the main challenges of fine-tuning on edge devices is the memory constraints.
Figure 1exemplifies
the memory usage of training some popular networks: ResNet50 with a small batchsize of 16 can easily exceed the memory limitation of onboard devices.
The newly emerged Vision
Transformer (ViT) would incur even more memory consumption.
Recent studies show that theactivations stored for back-propagation take up a large proportion of memory cost during training.
Therefore, popular model compression techniques such as weight pruning and quantization are not highly efficient in reducing training memory footprints.
Parameter efficient training method like
Bitfit also fails to compress the activation to memory limitation (as shown in Figure 1).
Gradientcheckpointing can reduce the training memory by storing only a subset of activations, but requiresmore FLOPs at back-propagation due to the re-computation of discarded activations.
Using a smallerbatch size or half-precision for training can also lower the memory cost, but they incur a decrease inthe model’s accuracy and training speed.
[1] successfully avoided saving activation via freezing theweights for convolutional neural networks and only learning a newly added the lite residual module,thus no need to store the intermediate activations.
However, it comes at a certain price of the transferperformance.
[Figure]
A new stream of works studied howto directly compress the activations.
They mainly leverage the quantization technique that discretizes the activation tensors so these tensors canbe represented by efficient data formats and require less memory [2, 3, 4].
However, the activation pruning techniques have been less explored.
[5]applied random activation masks oneach layer in both forward and backward passes to reduce latency.
[6]integrated weight pruning and activation pruning by learning importancescores on neuron’s outputs.
Moreover,they adopt the symmetric scheme thatprunes the activation at both the forward and backward processes.
Sucha pruning scheme undermines the precision of models’ outputs, leading tosignificantly worse performance when the pruning ratio is high.
In this work, we present a general activation compression method called Back Razor.
Inspired bythe
observations that the activations are only stored for back-propagation, we adopt the asymmetricpruning scheme that only prunes the activation during the backward process.
Specifically, afteractivation is used by a subsequent layer, our method replaces it with a pruned activation which carriesa lower memory cost.
During the backward phase, the gradients are calculated with the approximatedactivation, but only with a small error between those with the exact activation.
Our contributions are summarized as following:
• We propose a memory-efficient fine-tuning method called Back Razor that prunes theactivation during back-propagation.
Our proposed method can be easily applied to variousneural networks, including popular convolutional neural networks and transformers.
Tothe best of our knowledge, this is the first work that studies pruning activation stored forback-propagation.•
We provide a theoretical analysis of our self-sparsified back-propagation method for Convolutional Neural Networks .
Specifically, we discover that under mild assumptions, ourmethod can share a similar convergence rate as vanilla SGD.•
We extensively verify the proposed Back Razor on both convolutional neural networks andTransformer for multiple datasets and show that the proposed method can achieve up to 97%sparsity with saving 9 . 2 × memory saving without losing accuracy.
By further verifying onGPU, the proposed method achieves 2 .
7 × on-device memory efficiency for ViT.
Related Works
2.1 Pruning
Pruning can be operated at different levels, e.g. , pruning on weight, weight gradients, activations, andactivation gradients.
Pruning on weights are most widely studied since it can accelerate the inferencephase [7, 8, 9]; pruning on weight gradients can help reduce communication cost especially in
Sparsify
Storage
Forward
Backward distributed training [10, 11]; pruning on activation gradients [12, 13, 14, 15] or activation itself havebeen less addressed since they are mostly dynamic during training.
[5] randomly generated staticsparse masks and applied them on each layer in both forward and backward passes.
[6] proposedintegral pruning that prunes both activation and weights by consecutively learning and applying twomasks on weights and activations.
However these works are adopting the symmetric pruning scheme,i.e. , pruning the activation at both the forward and the backward process.
Pruning the activationsat each layer brings errors, and they will accumulate as the forwarding proceeds.
In this work, weperform a simple yet meaningful twist: only prune the activations used for back-propagation.
2.2 Memory Efficient Transfer Learning
Transfer Learning from a model pre-trained on large datasets to small datasets has been widelyused for many applications [16, 17, 18, 19].
In contrast to the pre-training that is conducted onlarge clusters, transfer learning always happens on personal computing devices with limited memoryconstrains.
To fit the large model to small devices, a common choice is to fine-tune the last severallayers of the model [20, 16, 21, 22].
However, it can be a sub-optimal choice, especially for thedownstream tasks with large distribution differences.
In contrast, fine-tuning the whole can yieldbetter accuracy in most cases [1, 23].
However, it can easily lead to memory outages.
On the otherhand, some parameter efficient transfer learning [24, 25, 26, 27] works only requires tuning lessthan 1% of the parameters.
But this parameter efficiency only brings a mild memory efficiency.
Themost relevant prior work is TinyTL [1], it recognizes the main bottleneck of memory is in activationand introduces a new trainable module that only requires down-sampled activation.
However, itstill focuses on exact gradient calculation while we employ approximate gradient calculation, whichenables compressing activation without changing the model architecture.
2.3 Propagation Activation Compression
Few previous works have explored the approximate activation techniques.
Most of them focus on thequantization methods [2, 3].
[4] also explored the JPEG compression.
Following works improve thequantization via fine-grained group division [28, 29].
Recently, activation compression is extended to
Vision transformer [30] and Graph Neural Networks [31].
However, to the best of our knowledge,none of them have explored the compressing via sparsifying, which is an orthogonal direction withexisting methods.
Method
3.1 Activation is the bottleneck of memory
The forward of neural networks can be express as following
[Equation]
where x and y is the input and label, respectively.
L is the loss between prediction f
( x ; θ ) andlabel y calculated by loss function L .
f i and θ i denote the function and parameters of i th layer, respectively.
The network is composed of l layers.
We further denote the output of i th layer (a.k.aactivation) as z i
=
f i ( z i − 1 , θ i ) and z 1 = f 1 ( x ; θ 1 ) .
When conducting back-propagation, ∂L∂ z i and∂L∂ θ
i are required to be calculated for each activation.
Taking linear layer with z i
=
W i
z i −
+ b i ( θ i =
[ W i , b i ] ) as an example, the gradient required to becalculated are [1]:
[Equation]
As shown in Equation 2, the activation z
i is employed for computing one of the gradients.
Thiscreates the need of saving the activation of the forward process.
Otherwise, extra FLOPs are requiredfor re-computing the activation.
The size of activation is always larger than the weight as weight isalways shared across different patches or tokens.
Also, the size of activation can increase linearly withthe batch size.
Moreover, in the transformer, the length of input tokens can also lead to a quadraticincrease of the memory.
Discussion about freezing weight.
TinyTL [1] pointed out that we can avoid saving activationvia freezing the weight W i as other terms do not require the participation of z i .
This is true formost architectures of convolutional neural networks.
However, the emergence of ViT breaks thisassumption.
Many operations (e.g. self-attention, Softmax, GeLU [32]) in transformer would involvethe activation for computing ∂L∂ z i , which cannot be avoided if previous layers requires updating.
To tackle this challenge, we design a more general method: directly sparsifying the activation forbackpropagation.
3.2 Back-propagation activation self-sparsification
In contrast to the activation sparsification [5, 6] that prunes the activation of both forward andbackward, our proposed Back Razor leverages the back-propagation activation sparsification: itonly prunes the activation that is used for backward while keeping the forward activation dense asshown in Figure 2.
The normal backward process can be formally defined as∂L
[Equation]
where h ( z ) i ( z i , θ i ) :
= ∂
z i +1 ∂
z i and h ( θ ) i
( z i , θ i ) :
= ∂
z i +
∂
θ
i , representing the derivatives of z i +1 withrespect to z
i
and θ i , respectively.
For Back Razor, the backward process can be expressed as
[Equation]
where ˜ z i denotes the pruned activation z i .
We inherit the mathematical forms behind h ( z ) i and h ( θ ) ifrom the normal backward to preserve the gradient flow, but calculate the gradients on weights andactivations with the pruned layerwise activation ˜ z i .
With the above backward, there is no need forsaving the dense activation z i .
Instead, we store a sparse version ˜
z
i that is more memory efficient.
The detailed optimization algorithm with the proposed Back Razor can be found in Algorithm 1.
Inthe forward pass, we would compute with dense activation while only saving its sparse version.
Inthe backward, we employ the sparse activation for computing the gradient.
We use a simple pruning method for backward activation pruning: prune the smallest magnitudeactivations.
The smallest magnitude values are usually playing less significant roles than those withlarger magnitude.
This pruning method is also widely used for many pruning works [33].
In practice,we would prune all the activations to a fixed sparsity of λ by setting all the values under k th smallestvalue as zero, where k = λn , n is the total number of values for the activation.
As the ranking canbe slow for large activations, instead of ranking the whole activation, we rank the activation of eachsample and use the sample-wise threshold for pruning.
This largely shortens the ranking time.
Also,we noted that the thresholds across different samples are very close, which means it is functionallysimilar to global ranking.
When saving the sparse matrix, we save it into two parts: a bitmap that indicates the position ofthe non-zero elements and a smaller dense matrix that contains the values of the non-zero elements.
Algorithm
Initialize Initialize
Calculate
The bitmap has the same shape as the origin sparse tensor, but it is much smaller since it is in booltype.
The memory cost of this format is n 8 + λnc type byte, where c type is the count of bytes for theemployed data type.
For the previous activation pruning method, the error that comes from pruning could graduallyaccumulate layer by layer in the forward, causing the deviation from the pre-trained model andleading to performance degradation.
In contrast, the proposed Back Razor ensures the correctnessof the forward.
Moreover, in the backward progress, the pruning error of many layers would notaccumulate: the activations of those layers are only used for computing the gradient of itself (e.g.linear layer, batch/layer normalization layer).
It’s also worth noting that the proposed Back Razorsaves the same amount of memory as activation pruning under the same pruning ratio.
Comparison with recent parameter-efficient transfer methods.
Recent studies reveal that finetuning can be parameter efficient.
For example, diffPrune [24] points out that the fine-tuningperformance can match fully fine-tuning via only modifying 0.5% of the weights.
However, itemploys Gumbel-Softmax [34] to make the masked parameters differentiable, which means thedense activations are required for updating the weight.
Therefore, it cannot yield memory efficiency.
On the other hand, other works also reveal that fine-tuning part of the model can yield competitiveperformance [25, 26, 27, 35].
Though it can lead to memory saving, we empirically verified that
Back Razor can yield a higher saving ratio.
3.3 Theoretical Results
Now we study the convergence of Back Razor.
We simplify the notion and define L to be the trainingloss of a convolutional neural network that takes the model’s weights as the input.
θ t = { W i } Li =1 asthe flattened vector of the model’s parameters at step t , and g t be the gradient calculated with all thesamples in the training dataset.
Furthermore, we denote the stochastic gradient by ˜ g t , which is thegradient calculated with one sample from the training set.
We make several classical assumptions [36]:
Assumption 1.
Let g ( θ ) be the gradient of the objective L at point θ , then for all x
and y we assume there exists a non-negative constant vector β that
[Equation]
Assumption 2.
For each iteration t , the stochastic gradient ˜ g t satisfies the following conditions:
[Equation]
where n is the number of parameters of models and σ 2 j is a constant representing the variance forcoordinate j.
Assumption 3.
The objective L is bounded below by L ∗ .
These assumptions are discussed in Appendix A. We first prove the convergence for linear networks.
Lemma 1.
For linear neural networks, the gradient of parameters can decomposed by ˜ g t = ˜
A t ˜ z t ,and approximated gradient after activation pruning can decomposed by ˜ g ′ t = ˜
A t ˜ z ′ t , where ˜ z t is theexact activation and ˜ z ′ t is the pruned activations.
In linear networks, ∂L∂ z i only depends on the model’s weights after the i th layer, so pruning theactivation stored for backward does not change the its gradient.
Therefore in Lemma 1, ˜
A t isexclusively constructed by activations’ gradients, and not changed after activation pruning.
We call ˜
A t the transformation matrix for activations at step t .
Theorem 1 (Convergence) .
If Assumption 1-3 hold, ( p + αβ ) K a < 1 where p is the pruning ratioand K a is the squared condition number of the transformation matrix, then
[Equation]
where the T is the number of iterations, β = ∥ β ∥ ∞ , θ 0 is the initial weights and σ 2 = (cid:80) nj =1 σ 2 j .
The Eqn. 5 consists of two terms: the first term will converge to zero as the number of iterations Tgoes to infinity.
The second term will not vanish, but its value can be controlled by using larger batchsize.
Intuitively, it proves that our algorithm reach to the neighborhood of a stationary point, andthe radius of the neighborhood is bounded by the gradient variance.
The gradient variance can bereduced by larger batch sizes: when using a batch size of B , the variance is reduced by B .
Remark.
The key factor for theoretical convergence is ( p + αβ ) K a .
αβ can be reduced by usingsmaller learning rate.
The theoretical convergence will not be guaranteed when the transformationmatrix is ill-conditioned.
We next discuss the convergence for non-linear networks, and more specifically, we focus on analyzingconvolutional neural networks with a common CONV - BN - ReLU structure [15].
The activation pruninghappens after the convolutional layer, where the gradient on weights and activations can be calculatedas [15]:
[Equation]
which suggests the gradient on z i will not be affected if we use the pruned activation
˜ z i to obtain thederivatives on θ i when back-warding through the convolutional layer.
The activations are neitherstored nor pruned for BN and ReLU operations, so the Lemma 1 holds for convolutional neuralnetworks.
So under the same assumptions, we achieve similar convergence situation for CNNs.
Wepresent more details in Appendix.
We lastly discuss the Transformer structure.
The self-attentions inside the Transformer do not meetthe Lipchitz condition [37], therefore we do not provide theoretical analysis for them.
Howeverwe empirically observe that Back Razor is also applicable on Vision Transformer [38] with decentperformance (refer to Section 4.3)
Experiment4.1 Setting
Dataset Following the common practice [24, 39, 40], we employ supervised pretrain of Imagenet [41].
Specifically, we employ ImageNet1K and Imagenet22K for convolutional neural networks and ViT,respectively.
For downstream fine-tuning, we consider four datasets: CIFAR10, CIFAR100 [42],
Aircraft [43] and Pets [44].
For all datasets, we employ resolution of 224 × 224 following [1].
Platform Our experiments are implemented with Pytorch [45] and conducted on 1080 Ti or V
GPUs.
Memory Measuring We by default report the theoretical memory usage following [1].
Wealso report the memory usage measured on the GPU at Section 4.3.
[Table]
4.2 Back Razor with Convolutional Neural NetworksSetting
We choose ProxylessNAS-Mobile [47] as the backbone following TinyTL [1].
It is worthnoting that we keep the origin setting of ProxylessNAS-Mobile instead of modifying the architectureas in TinyTL given Back Razor is applicable for any architecture.
We also follow the most oftraining settings of TinyTL for fair comparison: The fine-tuning epochs and batch size is set as 50epochs and 8, respectively.
The model is optimized with adam [48] optimizer and cosine learninglr schedule [49].
The initial learning rate is tuned for each dataset.
We freeze all the parameters of
Batch Normalization layers [50] as they occupy more memory than convolutional layers while havingmuch fewer parameters to update.
By freezing them, we could free the occupied memory with lessinfluence on the fine-tuning performance.
We also use the origin activation for ReLU6, which ismemory efficient (equal to the memory cost of a bitmap with the same shape).
The Back Razor isthen applied to convolutional layers and the fully-connected classification head.
Comparison In Table 1, we compare the proposed Back Razor with previous transfer learningmethods on CNN.
Specifically, it includes i)
FT-Last: fine-tuning the last fully connected layer of thenetwork, which is also referred as linear classification head.
ii) FT-Norm+Last: fine-tuning the batchnormalization layers plus the classification head iii) FT-Full: fine-tuning the full model.
Also, weinclude two different variants of the previous State-of-The-Art method [1]: TinyTL and TinyTL@320(TinyTL@320 employs a larger resolution image of 320).
Compared with FT-Full, the proposed Back Razor@ 90% use 8.7x times smaller memory whileachieving comparable performance.
The variance for BackRazor@90% with respect to differentrandom seeds is also comparable with FT-Full and they are both small.
In contrast, FT-Norm+Lasttakes 4.6x time more memory while being [0.7%, 17.5%, 2.1%, 3.1%] worse than the proposed Back
Razor@ 90% in [Pets, Aircraft, CIFAR10, CIFAR100], respectively.
Though FT-Last does yieldslightly less training memory (32MB v.s. 42MB), the performance largely degrades: the accuracyis lower by [1.6%, 42.7%, 11.0%, 14.5%] than the proposed Back Razor@ 90% in [Pets, Aircraft,
CIFAR10, CIFAR100], respectively.
Compared with the previous State-of-The-Art method TinyTL, the proposed Back Razor@ 90% iscomparable in terms of the training memory usage (37MB v.s 42MB).
Meanwhile, it yields higheraccuracy of [1.1%, 12.2%, 1.0%, 1.9%] at [Pets, Aircraft, CIFAR10, CIFAR100], respectively.
Whencompared with TinyTL@320, which employs a larger resolution, the proposed Back Razor@ 90%saves 1.5 times of the memory while still yielding an improvement of [0.0%, 5.3%, 0.8%, 1.8%] at[Pets, Aircraft, CIFAR10, CIFAR100], respectively.
It’s also worth noting that the proposed Back Razor method can even achieve more aggressive sparsity.
For instance, the performance of Back Razor is [ 92 . 9% , 93 .
5% ] for pruning ratioof [ 90% , 95% , 97% , 99% ], respectively.
The performance maintains the same level even with thepruning ratio of 97% and only marginally drops at the pruning ratio of 99% .
A similar trend is also
[Figure]
Pruning Ratio
A cc u r acy
M e m o r y
U sa g e // 0 .
[Figure] [Figure]
[Figure] [Figure]
Ratio
A cc u r acy
[Figure]
U sa
[Figure]
observed in other datasets.
However, the high pruning ratio does not bring further memory reductionsince the main bottleneck of memory is no longer the accumulated activations.
Instead, it is the modelparameters (of 11.6MB) and one large forward activation (of 24.5MB).
Can we prune at forward pass We further compare the performance difference between activationpruning (prune both forward and backward) with the proposed Back Razor (only prune backward)in Figure 3.
We implement the activation pruning via directly applying the mask of the Back Razorto the forward pass.
It can be observed that the model can easily collapse (the model can hardlyconverge and only have accuracy below 5%) for activation pruning even with a mild pruning ratio of70%.
In contrast, for the proposed Back Razor, the performance can be at a high level even with anaggressive pruning ratio of 99% ( the performance is even higher than activation prune at a sparsityof 50%).
This demonstrates that the backward activation can be more sparse than that of the forward.
Convergence speed in practice As shown in Figure 4, the convergence speed in transfer experimentsis comparable with each other.
In the early stage of the training, the convergence speed of the proposed
Back Razor is even faster.
And in the last stage of the training, the FT-Full is only marginally higherthan Back Razor.
4.3 Back Razor with Vision Transformer
Setting In this section, we study the proposed Back Razor on ViT [38].
The patch size is by defaultset as 16.
The finetuning employ the standard SGD optimizer with cosine learning rate decay.
Thetraining steps are fixed as 20k and the initial learning rate is tuned for each dataset.
We employ alarger batch size of 128 following the common practice for accelerating training [38, 51, 30].
[Figure]
[Table]
Comparison In Table 2, we compare the proposed Back Razor with previous transfer learningmethods on ViT. It includes the FT-Last and FT-Full with the same definition of Section 4.3.
Moreover,we further compare the parameter efficient training method Bitfit [25] and the backpropagationquantization method Mesa [25].
Compared with FT-full, the proposed Back Razor@ 80% yields comparable performance across thedatasets while achieving 3 . 9 × memory saving.
Though FT-Last employs smaller train memory thanthe proposed method, the accuracy drops by [2.2%, 2.1%, 6.4%] in [Pets, CIFAR10, CIFAR100],respectively, compared to the proposed Back Razor@ 80% .
The Bitfit achieves a competitive performance via tuning less than 0.1% of the parameters.
However,it can only reduce the memory saving of 2 .
× memory given it requires saving some large activationsin full (e.g. attention map).
In contrast, the proposed Back Razor can reduce the size of any activations.
While Mesa has similar memory usage and performance compared with the proposed BackRazor, themethod of mesa is orthogonal with the proposed Back Razor and by combining them together, wecan achieve 7 . 5 × memory saving with comparable performance to FT-Full.
size of 128 (test in CIFAR100).
We skip Bitfit here as it cannot yield hardware memory efficiency on GPU with
Pytorch.
[Table]
When employing a larger pruning ratio in ViT, the memory
[Table]
usage can be more efficient and the performance would
[Table]
not incur large degradation.
Especially, for CIFAR10, the
[Table]
performance only drop by 0.2% compared to FT-Full at
[Table]
a pruning ratio of 95%, demonstrating the proposed Back
[Table]
Razor can achieve an aggressive pruning ratio in ViT.
[Table]
Hardware memory efficiency Table 3 illustrates the hardware memory usage for different methods.
When testing on the modern computing device GPU, the Back
Razor@ 80% can yield 2 .
7 × memory efficiency.
By com-
[Table]
bining with Mesa, its efficiency can achieve 3 . 4 × .
The
[Table]
theoretical memory usage aligns with real memory usage.
As the theoretical memory usage is not achieved, there is
[Table]
still optimizing space for the implementation.
[Table]
Conclusion
[Table]
In this work, we propose Back Razor, a memory-efficienttransfer learning method, which is also the first asymmetricpruning method that only sparsifies the back-propagation activations while keeping the forward.
Through extensive experiments in both Convolutional Neural Networks and Vision Transformer acrossmultiple datasets, we demonstrated the proposed Back Razor can achieve state-of-the-art memoryefficiency with an aggressive pruning ratio.
Moreover, it can be combined with the propagationquantization method and yield better memory efficiency.
Limitation and future work
In this work, we mainly study the transfer classification task.
Our futureworks targets extending Back Razor to various downstream tasks, including semantic segmentation,dense prediction, and language understanding.
Checklist
For all authors...
(a) Do the main claims made in the abstract and
introduction
accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[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)?
Yes
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?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
The paper itself comes with a license.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
We do not curate/release new assets.
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
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?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A
Theory
A.1 Assumptions
We provide grounds for assumptions in this section.
Assumption 1 is a fine-grained version of the standard Lipschitz smoothness, i.e. ,
[Equation]
We follow [52] to make this assumption.
Assumption 2 combines the assumption of the unbiased gradient and bounded variance, which isalso standard in literature.
The standard variance bound on the gradient can be recovered by definingσ 2 = (cid:80)
j σ 2 j
and we can see that E [(˜ g t − g t ) 2 ] ≤ σ 2 .
Assumption 3 defines a lower bound for objective value, which is necessary for convergence tostationary points.
A.2 Proof of Theorem
Let L k be training objective at step k ,
i.e. , L k
=
L ( θ k ) .
Let m be the pruning mask for activation˜ z k , then we have ˜ z
k
= diag ( m )˜
z k .
Proof.
The difference between the two objectives L k +1 − L k :
[Equation]
Assumption
[Equation]
which is based on the SGD update rule θ k +1 − θ k
= − α k ˜ g ′ k where ˜ g k represents the prunedstochastic gradient.
For the first term, we have
[Equation]
and
[Equation]
Assume ˜
A k has positive singular values
[Equation]
we have
[Equation]
where p := n − Kn is the pruning ratio and K a := max λ 2 (˜ A k ) min λ 2 (˜ A k ) is the squared condition number of ˜ A k .
This is because1.
[Equation]
since m is generated by the TopK operation.3.
[Equation]
The second term can be bounded by
[Equation]
Combining these two items together:
[Equation]
Next we find the expected improvement at time k + 1 :
[Equation]
and by extending the expectation over the trajectory and summing over T iterations we finish theproof.