Efficient Backpropagation with Variance Controlled Adaptive Sampling
Authors: Ziteng Wang
Jianfei Chen
Jun Zhu
Venue: ICLR 2024 poster (2024)
Keywords: efficient training algorithms
stochastic gradient descent
importance sampling
variance reduction
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Self-adaptive sampling method that reduces the compute cost of backpropagation (BP) by up to 73.87%, with mirrored performance of exact training under theoretical guarantee.

Abstract: Sampling-based algorithms, which eliminate "unimportant" computations during forward and/or backpropagation (BP), offer potential solutions to accelerate neural network training. However, since sampling introduces approximations to training, such algorithms may not consistently maintain accuracy across various tasks. In this work, we introduce a variance-controlled adaptive sampling (VCAS) method designed to minimize the computational load of BP. VCAS computes an unbiased stochastic gradient with fine-grained layerwise importance sampling in data dimension for activation gradient calculation and leverage score sampling in token dimension for weight gradient calculation. To preserve accuracy, we control the additional variance introduced by learning the sample ratio jointly with model parameters during training. We assessed VCAS on multiple fine-tuning and pre-training tasks in both vision and natural language domains. On all the tasks, VCAS can preserve the original training loss trajectory and validation accuracy with an up to 73.87% FLOPs reduction of BP and 49.58% FLOPs reduction of the whole training process. The implementation is available at https://github.com/thu-ml/VCAS.

Published as a conference paper at ICLR 2024
EFFICIENT BACKPROPAGATION WITH VARIANCE-CONTROLLED ADAPTIVE SAMPLING
Ziteng Wang, Jianfei Chen', Zhu Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua- Tsinghua-Bosch Bosch Joint ML Center, Tsinghua University
wangzited e230mm 1s. s tannguua. edu. cn; (jianfeic, dcsz))atsinghun edu.cn
ABSTRACT
Sampling- campling-based based algorithms, which eliminate "unimportant" computations dur- ing forward and/or back propagation (BP), offer potential solutions to accelerate neural network training. However, since sampling introduces approximations training, such algorithms may not consistently maintain accuracy across various tasks. In this work, we introduce variance-contrroledd adaptive sampling (VCAS) method designed to accelerate BP. VCAS computes an unbiased stochastic gradi- ent with fine-grained layerwise importance sampling in data dimension for acti- vation gradient calculation and leverage score sampling in token dimension for weight gradient calculation. To preserve accuracy, we control the additional vari- ance by learning the sample ratio jointly with model parameters during training. We assessed VCAS on multiple fine-tuning and pre-training tasks in both vision and natural language domains. On all the tasks, VCAS can preserve the original training loss trajectory and validation accuracy with an up to 73.87% FLOPs re- duction of BP and 49.58% FLOPs reduction of the whole training process. The implementation is available at https://github.comm htt.ttuu..oo/tuuuuuueeeaea Sarttuuumm//AAA
INTRODUCTION
Training neural networks can be computationally inten- sive. Contemporary networks typically employ stochastic gradient methods (Bottou et al. 2018) for training, which iteratively process batches of data compute stochas- tic gradients through forward propagation (FP) and back propagation (BP) techniques (Rumelhart
Exact
1.0
Train Loss
SB
0.8
UB
al.] 1986).
FP+BP are costly, as they need to process every datum in the batch and every connection in the network. result-
0.6
VCAS
0.4
ing in multiplicative time complexity of batch size and model size. Such time complexity becomes increasingly problematic in the era of big data and big models.
0.2
0
10000 20000 30000
steps
Data samples are not equally important. Some might be easy for the network to learn, while others might be ex- tremely hard. Training can be accelerated by utilizing this disparity, focusing the available computational resources on more pivotal samples. At a high level, this can be achieved by further sampling the batch with higher keep probability of more important samples. The computa- tional overhead is consequently diminished, in proportion to the quantity of retained samples. Various methods are proposed assess the importance of samples, including
Figure 1: VCAS mirrors the conver- gence trajectory with exact training with FL.OPs redution of 41.56%. Other methods like SB (Jiang et al.) 2019 and UB (Katharopoulos & Fleuret| 2018) fail with similar FLOPs reduction.
meta-learning methods (Fan et al.] 2017: Coleman et al. 2019 Mindermann et al. 2022 loss-baser based
methods (Loshchilov & Hutter 2015 Chang et al. 2001 Jiang al. 22019] Ouyang et al. 2022).
Corresponding author.

---

Published as a conference paper at ICLR 2024
and gradient norm based methods (Needell al., 2014 Zhao & Zhang 2015 Johnson & Guestrin 2018 Katharopoulos & Fleuret 2018).
Alain et al. 2015
While such methods seem promising, ness. Misjudging the importance can hamper convergence, potentially leading to degraded accuracy and even longer training time than uniform sampling. Moreover, the optimal sample ratio is influ- enced by data distribution, which differs between tasks and is challenging to determine in advance. In general, there is a "no-free-lunch" phenomenon (Kaddour et al. 2023), where aggressive sam- pling often comes at the cost of reduced robustness.
concern of sampling- based methods is their robust-
In this work, we propose robust variance-controlledd adaptive sampling (VCAS) algorithm for deep learning under the stochastic optimization framework. VCAS computes a cost-effective ap- proximated stochastic gradient (ASG) by partially conducting backpropagation for specific data and tokens. This ASG is unbiased, and we have developed an adaptive sampling method to meticulously control the variance of the ASG, aligning it with the original stochastic gradient's variance. Con- sequently, convergence remains largely unaffected, with our method mirroring the progression of exact algorithms,, as delineated in Fig.
Unlike previous methods, VCAS construct the ASG in a fine-grained manner. Rather than dropping samples one-time in a whole, VCAS gradually drops more samples when backpropagating from topmost to bottommost network layers, the gradient getting sparser. Furthermore, VCAS also more aggressively drops data in finer granularity of tokens rather than samples when computing the weight gradients. VCAS can achieve smaller variance under given computational budget compared to coarse grained sampling on the data dimension.
We evaluate VCAS multiple finetuning and pre-training tasks of language models and vision transformers. VCAS can preserve the original training loss trajectory and the validation accuracy on all tasks, while adaptively determining the computational saving depending on the difficulty of the task. VCAS can reduce the computational cost of backpropagation by up to 73.87%, and reduce the overall training computation by up to 49.58%.
2 RELATED WORK
Methods focusing on the difference of data, known as online batch selection (Loshchilov & Hutter, 2015), can be mainly categorized into three classes: meta learning methods, loss based methods and gradient norm based methods. this section we will discuss these three ways separately and briefly introduce other orthogonal efficient training methods.
Meta Learning Methods. Some works formulate data sampling into an optimization problem and train a separate meta predictor to solve it. Fan al. (2017) use deep reinforcement learning to train an agent for data selection. CColman et al. (2019) and Mindermann et al.] (2022) train a separate cheaper model with similar architecture for guidance. However, training meta predictor will intro- duce further overhead and it's a non-trivial learning task with more uncertainty introduced for weak theoretical guarantee.
Loss Based Methods. Loss is a natural indicator of the importance of different data, Loshchilov & Hutter 2015) maintains history of losses and develops a sophisticated distribution based on the value or rank of loss. Jiang et al. (2019) and Ouyang et al. (2022) simplify it with sampling distribution proportion
the percentile of Toss in the history. Chang et al. (2017) broadens the history to every datum and proposes to sample by the variance of prediction probability directly linked with previous losses. Dong et al. (202T) provides another method of minimizing the L2 norm between the sampled loss and the exact counterpart. [hhah et al. (2020) samples the smallest loss for robustness to outliers. /Zhang al. (023) ensembles several loss methods with a preset sample ratio and varies the weights assigned to these methods adaptively. Simple and effective they may be, the loss based methods are heuristic and always need a hyperparameter of sample ratio to tune for different tasks, violating the goal of efficient training.
Gradient Norm Based Methods. Previous works have proved that the optimal data sampling dis- tribution for SGD is proportional to the gradient nerdeNeedell et al.) 2014: Zhao & Zhang 2015). But calculating the gradient norm is prohibitive since it needs a full process of backpropagation. To solve this problem, Alain et al. (2015) applies distributed training with many workers calculating
2

---

Published as a conference paper at ICLR 2024
Layer
260
V200
22(1-1
SampleAcot
h9
Layer +
Layer
flatee
2m
FerD
Sample War
g
NT
Figure 2: The computing diagram of backpropagation with VCAS in every layer. We use light blue squares to represent small gradient entries and orange for large ones. White squares are discarded by sampling. The upper line calculates activation gradient and the lower for weight gradient. Please refer to Sec. for notations.
this importance score in parallel. Johnson Guestrin (2018) uses a sedond-order approximation of gradient norm with history maintained. Closely related
work, Kathaaaaaaooooolloo & Fleuret (2018)
develops pure online algorithm by constructing an upper bound of gradient norm to sample with much cheaper computation. These methods usually more expensive but have relatively strong theoretical guarantees. So we follow this way in our activation sampling.
Orthogonal Efficient Training Methods. Data pruning (Paul et al. 2021 Fayyaz al. 2022 fo-
cuses on filtering less informative data before the whole training. Architecture pruning like layer dropping (Huang al. 2016 Zhang & He 2020) and token dropping (Hou al. 2022 Yao al., 2022 et al. 2022) modifies the architecture to make models faster to train with modest affect performance. Mixed precision training and quantization (Micikevicius et al.) 2018 Chen al. 2021 Liu et al. 2022) change the training procedure to use low-precision in calculation for accel- eration. Sparsity (Hoefler et al., 2021) focuses on pruning near-zero values in weights. activations, or gradients to achieve a low (Rihan & Aamodt] (220) and low memory footprint(Nikdan al. 2023), yet is usually hard to bring a wall-clock time reduction like us due the lack of hardware (ooport(NVIDIA 2021). All these works are orthogonal to our work since we focus on the computation approximation of a certain model architecture on certain dataset with a certain training procedure to bring real training acceleration.
pparsity(Hoefler
3 VARIANCE-CONTROLLE
SAMPLING AS STOCHASTIC OPTIMIZATION
In this section, we present high-level overview of our sampling algorithm stochastic opti- mization. Consider the learning problem of a model f(X;0) parameterized by 0 on dataset
PP with loss function ((.:).. Define the learning objective as D = [(X,, 1))
( E3 ((;
where the expectation is taken over all possible batches B (X, from D. The model parameters can be learned by stochastic optimization algorithms (Bottou al. 2018) with a stochastic gradient (SG) g(; 13) (f 0). y), which is unbiased approximation of VeL(O).
However, computing the stochastic gradient can be still too expensive, since requires the full for- ward and back propagation, which iterate over all model parameters and all data in the batch. We build a cheap stochastic approximation g(:: B3 c) of the SG. which we refer as approximated stochas- tic gradient (ASG). ASG only computes the backpropagation partially, and is therefore cheaper than the SG. The randomness in the computing procedure of ASG is captured by e. We ensure that ASG unbiased: E, (0; B, <)] g0; B).

---

Published as a conference paper at ICLR 2024
Backpropagation
Backpropagation,
Backpropagation
1e-10
30
30
30
100200 Luou DEj6
25
8
25
25
1e-8
20
2
20
-
20
16-6
Datum
153
15
15
10
10
10
1e-4
08
1e2
10
8
4
10
0
10
8
2
0
Layer Index
Layer Index
Layer Index
(a) 100-th iter.
(b) 300-th iters
3000-th iters
Figure 3: Gradient distribution over different layer and iterations of BERT-base finetuning on SST2 (6315 iterations in total). The normalized gradient norm of each datum is shown in the heatmaps. Black solid lines are the 95% percentile. Data above the lines are likely to be dicarded by VCAS.
With an unbiased SG, stochastic optimization algorithms are guaranteed to converge stationary point of Eq. (T). while the converge speed depends on the variance (cf. Bottou et al. (2018)). There- fore, if the variance of the ASG be controlled to the similar variance level of SG, substituting the SG with ASG should have little impact to the convergence behavior. In fact, by the law total variance (Chung 2001), the variance of ASG can be decoupled as
Var [g( B, e)] Var [g(0; ; B)] + Es [Var, [9(0; [[0;B,e)]]
where the first term the intrinsic variance SG caused by subsampling batches from the dataset, and the second term is the additional variance incurred by ASG. In the subsequent sections, we will discuss our constructions of the ASG, which incurs negligible additional variance compared to SG.
FINE-GRAINED SAMPLING
Here present variance-controoleed adaptive sampling (VCAS), specific construction of the ASG. We compute ASG by approximating the backpropagation in fine-grained manner, and speed up matrix multiplications with importance sampling on the data dimension.
Assume a batch X of shape N x T' x K, where N is the batch size, T is the number of tokens of each datum, and K the dimensionality. For an I-layer network the model f(X;0) can be described by the following forward propagation procedure: Z X, z())
f()) (Z(-1); ((V) (X::;) 0) z(L), where 7() and gil) are the activation and parameters of the 1-th layer, and = (0(D) ) The SG can computed by back-propagation in the following form: Vz(-1> = ( (V z (11 8(0). Voc = =g(")
()) (Vza) z(=1)),())) where V z0 and V Voon de- note the activation weight gradient, (h) and (d) denote the function that calculates input / weight gradient of layer with the output gradient, layer input and weight. The SG g(0; B) = (V (Voo))f. )
As illustrated by Fig. 3 the activation gradients V z are sparse: the gradient (V 2un); is close to zero for most sample i, except for few important samples. Such sparsity becomes more prominent as backpropagating lower layers and the training progresses. To speed up computation, we add samplers in the backpropagation graph:
z SampleA c.p (Vz(), Vz = = (v (V20; (11 2(1-1),(0))
.00)
V z SampleWe... Ena (4
(00,,z111) (Fz0,2((1)) z(-1)), V ( (vz0;; 2(1-1) 100))
(2)
The sampler SampleA c.r (.) randomly filter out unimportant data from the activation gradient, the keep ratio is Pt, with the randomness captured by The sampler is applied for each layer, so the activation gradient becomes increasingly sparse when backpropagating from the L-th layer to the first layer. The sampler SampleW, su. (-) filters (data, token) pairs specifically for weight gradient calculation, with a keep ratio in and the randomness &t- With these samplers, we only need to compute backpropagation for the retained data / token, the computational cost is reduced. The sampling procedure is illustrated in Fig. which constructs an unbiased ASG g(:: (; B, e, p,v) (V W with V oco defined Eq. and & (Ei)= P (p)f V (vi)==1-
4

---

Published as a conference paper at ICLR 2024
4.1 ACTIVATION GRADIENT
We apply unbiased low-variance approximation to the activation gradient to speed up subsequent computation. For an activation gradient tensor G of shape N T x K, sample
G = SampleA (G) G (m((, p) & 10 1),
where o is element-wise product, and & is tensor outer product. The mask m E RN is random Bernoulli vector: p)c,p); = Bern (p:; p where Σi Pr Np, and Bern(p; denotes Bernoulli random number generator with probability p and randomness c. Since Efmm(c, = 1, Vi, the approximation is unbiased: E[G] G. The sampler zeros out the gradient for all the data whose m(c, p): O. The amount of retained data is Np expectation. With the sampler, we only need to compute backpropagation for retained data, so the cost is times lower.
The variance of the approximation is Var G = Ei 1- [|G.lle- where we define the variance
of a random tensor element-wise as Var G - Ligh Var Cissa and G. denotes the i-th matrix of
-
G in the dimension. We compute the keep probability (p:) to minimize the variance, deriving a distribution proportional to the gradient norm of each datum: p, oc ||G.]|p. Minimizing the variance of the activation gradient not necessarily minimize the variance of ASG, which is the gradient of parameters. Nevertheless, this is useful heuristic which empirically achieves low variance is revealed by Katharopoulos & Fleuret 2018), and the ASG variance will be carefully controlled by adaptive algorithm, shall see soon in Sec.
4.2 WRIGHT GRADIENT
We can accelerate the computation of weight gradient for linear layers by sampling in both data and
token dimensions. Consider the approximate back propagation of linear layer Z = g0) :
Z SampleA (Vzi)),
V z SampleWer.. E1,P1 (vz0, (u-1), V Vic
V z (
in matrix form, where reshape the activation/gradiinnss to NT x K, and V 2 is already sampled matrix with only NTp, pi non-zero rows in expectation. However, V z is only sampled in the data dimension. In fact, even (V 2 )s is retained for some datum i, might still have some rows (i.e.,
tokens) which are close to zero. We can further sample
20 =SampleW En ($20) (0-)) = = 200 ()
where the mask R. is random Bernoulli vector, and is all-one vector: m(E.v); Bern(g;:)) c)/ge qi, where ENT qi NTppl. The variance is
NT
Var
V
pm
Σ
qi
V
z()
2
- 11
4
The minimal variance solution is q as leverage score sampling in randomized numerical linear algebra (Drineas & Mahoney 2018
2
This sampling method is also known
5 ADAPTING SAMPLE RATIOS
The question remained is how to set the sample ratios (pt)_1 and (v)) There a tradeoff:: lowering the sample ratio reduces the computational cost, but increases the variance. As discussed in Sec.B] this ratio should be set to ensure that the additional variance of ASG is marginal compared to the original variance of SG. Adapting the sample ratio is nontrivial since the gradient sparsity pattern vary across layers and vary over time during training. In this section, we present an adaptation algorithm to control the variance during the entire training trajectory.
First, we introduce single hyperparameter to control the sample ratios (pt) for all layers. Intuitively, when the gradient norm (|[G&ll becomes sparser, we can aggressively utilize
5

---

Published as a conference paper at ICLR 2024
smaller keep ratio i to maximize speedup. Therefore, compute pt based on the sparsity pt of the gradient norm sequence:
n pr(s) = mn/(n/N1 Σ |IG.l 2
Σ [G..lpp) pr(s) = max maxps Ps 35
(4)
where s E [0, 1] is hyperparameter on how much gradient norm is preserved. It's shown in Fig. that gradient norm grows sparser with layer, yielding descending trend of for I from L to 1. Thus it's reasonable to construct a monotone increasing sequence of {pr)/=1 based on {pi)f=
By law of total variance, we can decompose the variance of ASG as
Var [g(0; B, C, v)] Var (0(0; B3] + Ea(Var, [g(0; B, c, p(s))]] + BB,,, [g(0; B, e, P. v]]. where we write g(0; e, :- Be(g(0; e, p, to be the ASG without the sampler for weight gradient computation. The three variance terms are the SG variance, the variance introduced by approximately computing activation gradient, and the variance introduced by approximately com- puting weight gradient, respectively. Our algorithm adaptively tunes & and v during train to control the last two variance terms to be fractional comparing to the first variance term.
Controlling Es (ar, [g(0; B, ((s))]] We adopt zeroth order method to adapt the hyperparam- eter s to keep Var, (9(0; B, p(s))] Tast Var [g( B3), where Teet < small constant. That is, the additional variance raised by approximately computing activation gradient is only Tact times the SG variance itself. Since larger increases the keep ratio and decreases the variance, we adopt the update:
(5) where sign(x) = +1 when 2 0 and sign(x) -1 when 0, and is a step size. We approximate the expectation and variance with empirical ones with M Monte Carlo repetitions. Therefore, each update requires ((M²) FP+BPs, and we run the update every F SGD iterations, where > M
8t 8t a sign [Va[[Vr, [g(0; B, p()))((s)]]] - Teot Var [g(0; B)]),
Controlling Es, [Var [g(0; B, e, &, r]]: can further decompose the variance as
As the variance sums up for each parameter (00,
L
[Var [g(0; B, e, E.
=
Σ 1=1
Es,c Varg (, C,,
(6)
where (()) is the gradient of the 1-th layer (i.e., V (oo)). ). We control the variance of each layer sepa-
rately to keep Es, [aar [g (0; E, vi]] Tw Var (0: B3). Again, this is achieved by a zeroth-order algorithm:
un (se.
Be., (0. (o..))))
where Varg [9] can be computed analytically by Eq. and a multiplier.
Now we are fully prepared to present the whole picture of VCAS in Alg. Please refer to Ap- pendix. for more details about the algorithm.
EXPERIMENTS
6.1 TRAINING FLOPS REDUCTION
We assessed VCAS on multiple fine-tuning and pre-training tasks in both vision and natural language domains. We compare our algorithm with the exact training and two previous works in BP sampling: a loss based method SB(selective backprop) in Johnson & Guestrin 20188 and gradient norm based method UB(upper bound) in Katharopoulos & Fleuret (2018). We choose these two methods since they are entirely online and need little modification to the original training pipeline like us. The results are shown in Tab. ! All results are the average of 3 different seeds except for BERT-base pretraining and VIT finetuning on ImageNet- 1k which we use 1.

---

Published as a conference paper at ICLR 2024
Algorithm 1 Variance controlled adaptive SamplinggVCASS for backpropagation
Require: update frequency F, Monte-Carlo repetition number M, variance tolerant ratio for activation Tacts for weight Tw, step size weight ratio multiplier :3
s - activation sample ratio schedule {pr f- weight sample ratios {ui
t 4 0
while not converge do
ift mod F = 0 then
for in 1, ... MM do
(X,, v6) <- batch selected randomly
SGD gradient G.,' 4- exact backward using (Xx, M6) for j in ,. M do activation gradient Gact.i.) + backward using (X,,, Us) with SampleA only calculate weight variance Vena, analytically with Eq, 3 end for end for
and Eq. 6
SGD variance V. + M-- E/
Gas
1
Em Gaill
MI
activation variance Vast + Md Li=l
M Z |Gaatt Ga,ll?))
weight variance V Mr
Σ -i=1 EM Vunsss)
update s with Vace and Vs according to Eq.
update (pttte with new according to Eq.
update (m) with V., and V. according Eq.
end if
backward with SampleA and SampleW t <- t ++1
end while
Note that to avoid falling into the pitfall of unfair comparison with baseline which is not tuned under efficient settings as is pointed out by DDegghan et al. (202T) and Kaddour et al. (2023). for all these experiments
use the same conservative setting of Tact Tw 0.025, = 0.01, 8 0.95, M 2. We preset all these values heuristically without any tuning or prior knowledge. The only hyperpamater we modified among different tasks is the variance calculation frequency F, which can be defined easily according to the total training steps.
In fact, all the hyperparameters introduced by VCAS have explicit meanings and are insensitive. We show experimentally that though extra tuming may achieve a slightly better result, overall VCAS robust to these hyperparameters with reasonable values. Please refer to Appendix. A for details about ablation studies on these insensitive hyperparameters.
For SB and UB, both adopt sample ratio of 1/3, since it's the recommended setting in the
original papers and it can achieve a FLOPs reduction of (1 + 2 * 1/3)/3 = 44.44% which is close to the results we get in most tasks. An exception is BERT-base pretraining task where we find the FLOPs reduction achievable is low so we manually set the sample ratio of SB and UB to get the same FLOPs reduction as VCAS, so that they can still give decent result. Nevertheless we are indeed favoring these methods by helping them to define a reasonable sample ratio, which can not be done themselves.
From the table we can see that overall VCAS is better than SB and UB with the least impact on final train loss and final evaluation accuracy. With FLOPs reduction of up to 49.58% VCAS can still achieve nearly the same results with the exact counterpart.
6.2 WALL-CLOCK TIME REDUCTION
We record the wall-clock time of BERT-large finetuning ImageNet- 1k with NVIDIA 3090Ti, the results are depicted in Tab. Lab.22an and Tab.
MNLI and ViT-large finetuning
From these tables, find that VCAS can translate FLOPs reduction into wall-clock time duction as effectively as simpler online batch sampling methods like UB and SB that drop part of
7

---

Published as a conference paper at ICLR 2024
Table Comparison of VCAS with other methods. Data format is Final Train Loss Final Eval Acc. for exact, and UB, Final Train
Final Eval (cc. / FLOPs reduction ratio(% for VCAS. FLOPs reduction of SB and UB 21 21.58% BERT pretraining and 44.44% for other tasks. VCAS's FL.OPs take account of the adaptation overhcad. For BERT pretraining, accuracy=average performance on GLUE. Bold indicates the best result of each metric except for exact. Underline means Eval Acc less than 0.1% off the exact training.
Task
Dataset
exact
SB
VCAS
BERT-base pretraining
C4
2.099 78.37
2.133 / 77.53
2.106 77.96
2.134 78.36 21.58
MNLI-m
0.2372 84.33
0.3833 83.71
0.2957 / 83.82
0.2428
84.23 41.56
BERT-base
QQP
0.1143 91.00
0.1441
90.76
0.1964 / 89.53
0.1189
90.92 47.10
finetuning
QNLI
0.1014 99.67
0.2017
90.58
0.1441 /91.23
0.1056 91.29 44.45
SST-2
0.0559 192.59
0.0727 92.63
0.0743 92.82
0.0600 93.04
48.28
MNLI-m
0.1439 / 86.58
0.2492
85,18
0.2266 / 86.09
0.1619
86.63 44.17
BERT-large
QQP
0.0885 99.64
0.1308 91.20
0.1751 19..51
0.0962
91.57 49.50
finetuning
QNLI
0.0877 /9.02
0. 0.1436/ 91.50
0.1325 99.98
0.0640
92.15 46.19
SST-2
0.0537 93.60
0.1136 91.81
0.0838 93.40
0.0593 93.67
49.24
CIFARIO
0.1868 98.92
0.2367 98.82
0.1923
98.94
0.1873
98. 98.90/ 45.90
ViT-base
CIFAR100
0.8760 191.19
2.248 / 89.60
1.175 89.68
0.8811 99.08 29.32
finetuning
ImageNet-1k 1k
0.6032 82.27
0.6533 82.09
0.6109 82.28
0.6089 82.27 45.29
CIFARIO
0.1359 199.24 .24224442
0. 0.1439/ 99.21
0.1378 19.17 17
0.1393
99.28 48.37
ViT-large
CIFAR100
0.4590 93.56
0.5983 93.07
0.5170 / 93.36
0.4649
93.66
38.67
finetuning
ImageNet-1k
0.4135/ 82.04
0.4637 82.21
0.4242 / 82.21
0.4228
82.27 49.58
Table 2: Wall-clock time of BERT-large finetuning on MNLI.
Method
Train Loss
Eval Acc.(%)
Wall-clock Time(h)
FLOPss(()
Time
exact
0.1439
86.58
5.478
SB
0.2492
85.18
4.320
44.44
21.14
UB
0.2266
86.09
4.266
44.44
22.12
VCAS
0.1619
86.63
4.437
44.17
19.00
Table 3: Wall-clock time of ViT-large finetuning on ImageNet-Ik.
Method
Train Loss
Eval Acc. (%)
Wall-clock Time(h)
FL.OPs.(%)
Time
exact
0.4135
82.04
52.29
SB
0.4637
82.21
42.56
44.44
18.61
UB
0,4242
82.21
41.92
44.44
19.83
VCAS
0.4228
82.27
41.28
49.58
21.06
data one-time in whole, while enjoying mirrored performance with the exact training under theo- retical guarantee.
The success of VCAS comes in ways. One is the fine-grained sampling strategy that samples activation and weight jointly, which enables us to achieve much lower FLOPs given the variance budget. The other is the variance controlled framework combined with the self-adaptation algorithm, with which we are able to learn the proper sample ratios of different training phases. In the following two subsections, will experimentally show the effectiveness of these two folds.
6.3 EFFECTIVENESS OF FINE-GRAINED SAMPLING
We compare VCAS that samples activation and weight jointly with strategies that solely sampling activation weight. Specifically, MNLI. We set Fast = T. 0.025 for VCAS, Tact = 0.05 for activation sampling only and т, = 0.05 for weight sampling only. We find that under the preliminary that Tacts Tuo 1, the results
keep equal extra variance for BERT-base finetuning
8

---

Published as a conference paper at ICLR 2024
10a
FLOPs reduction ratio
VCAS
0.4156
activation only
10a
weight only
0.3
0,22199
variance
0.2
10
0.3
0.0
10-1
10000
20000 steps
30000
40000
UB
SGD var threshold
VCAS weight var
58 var
10-2
5000 10000 15000 20000 25000 30000 35000 steps
Figure 4: FLOPs reduction ratio of sampling activation or weight solely with equal variance.
VCAS vs.
Figure 5: Gradient variance of different methods.
of these sampling strategies show significant difference due controlled variance. While as is shown in Fig. 4 VCAS can achieve a much greater FLOOPs reduction with the same total variance introduced. It's reasonable since can utilize more sparsity in both data and token dimensions with fine-grained sampling strategy of VCAS.
6.4 EFFECTIVENESS OF VARIANCE CONTROL AND SELF-ADAPTATION
In Fig. we plot the variance of different methods during training process of BERT-base finetuning on MNLI. We can find that VCAS is able to control the extra sampling variance introduced to preset threshold, while for other variance-unaware algorithms like UB and SB, the extra variance is out of control with similar FLOPs reduction.
With carefully controlled variance, a similar convergence with exact training is guaranteed as we mentioned in the introduction. As depicted in Fig. and Fig. 6 for BERT-base finetuning MNLI, VCAS shares nearly the same convergence trajectory with the exact training with reduced FLOPs, while UB converges slightly slower due uncontrolled variance, and SB converges in an entirely different trajectory with variance introduced far larger than exact.
1.0
Exact
Flops per iteration
Evall Loss
0.8
UB
SB
Acc
0.82
0.80
0.6
0.78
0.6
0.76
0.4
0.4
0
10000 20000 30000
0
10000 20000 30000
10000 20000 30000 steps
steps
steps
(a) Validation loss
(b) Validation accuracy
(c) FLOPs per iteration
Figure Convergence comparison of different sampling methods. FLOPs normalized by exact training.
CONCLUSION
We propose VCAS, robust sampling method for back propagation with controlled variance and self-adaptive sample ratios. VCAS computes approximate stochastic gradient by applying fine- grained sampling to gradually remove samples and tokens during backpropagation. VCAS enjoys similar variance, convergence trajectory. and final accuracy with exact back propagation, while duces the training cost by up to 49.58%.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
The authors would like to thank Bingrui Li and Weiyu Huang for their valuable discussions and help on algorithm design and implementation details. This work was supported by the National Key Re- search and Development Program of China (No. 2021ZD01 (2212D0110502), NSFC Projects (Nos. 62376131, 62061 136001, 62106123 12, 62076147 U19A2081, 61972224), Tsinghua Institute for Guo Qiang. and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
REFERENCES
Guillaume Alain, Alex Lamb, Chinnadhurai Sankar, Aaron Courville, and Yoshua Bengio. Variance reduction in sgd by distributed importance sampling. arXiv preprint arXiv:1511.00 1.06481, 2015.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM Review, 60(2):223-311, 2018.
Haw-Shiuan Chang, Erik Learned-Miller, and Andrew McCallum. Active bias: Training more accurate neural networks by emphasizing high variance samples. Advances Neural Information Processing Systems, 30, 2017.
Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Acton: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, pp. 1813-1813. PMLR, 2021.
Kai Lai Chung. course in probability theory. Academic press, 2001.
Cody Coleman, Christopher Yeh, Stephen Mussmann, Baharan Mirzasoleiman, Peter Bailis, Percy Liang, Jure Leskovec, and Matei Zaharia. Selection via proxy: Efficient data selection for deep learning. arXiv preprint arXiv: 17829, 1829, 2019.
Mostafa Dehghani, Anurag Amab, Lucas Beyer, Ashish Vaswani, and Yi Tay. The efficiency mis- nomer. arXiv preprint 0rXiv:2110.12810 2021.
Chaosheng Dong, Xiaojie Jin, Weihao Gao, Yijia Wang, Hongyi Zhang, Xiang Wu, Jianchao Yang, and Xiaobing Liu, One backward from ten forward, subsampling for large-scale deep learning. arXiv preprint arXiv:2104. 13114, 2021.
Petros Drineas and Michael W Mahoney. Lectures on randomized numerical linear algebra. The Mathematics of Data, 25(1). 2018.
Yang Fan, Fei Tian, Tao Qin, Jiang Bian, and Tie-Yan Liu. Learning what data to learn. arXiv preprint arxiv:1702 08635, 2017.
Mohsen Fayyaz, Ehsan Aghazadeh, Ali Modarressi, Mohammad Taher Pilehvar, Yadollah Yaghoobzadeh, and Samira Ebrahimi Kahou. Bert on data diet: Finding important examples by gradient-based pruning. arXiv preprint
2022.
Jonas Geiping and Tom Goldstein. Cramming: Training a language model on a single gpu one day. arXiv preprint arXiv:2212. 140344 2022.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. The Journal of Machine Learning Research, (011::0882211000 2021.
Le Hou, Richard Yuanzhe Pang, Tianyi Zhou, Yuexin Wu, Xinying Song, Xiaodan Song, and Denny Zhou. Token dropping for efficient bert pretraining. arXiv preprint arXiv:2203.13244 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Pision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 646-661. Springer, 2016.
10

---

Published as a conference paper at ICLR 2024
Angela H. Jiang, Daniel L-K Wong, Giulio Zhou, David G Andersen, Jeffrey Dean, Gregory R Ganger, Gauri Joshi, Michael Kaminksy, Michael Kozuch, Zachary Lipton, et al. Accelerating deep learning by focusing the biggest losers. arXiv preprint arXiv:1910.00762, 2019.
Tyler B Johnson and Carlos Guestrin. Training deep models faster with robust, approximate impor- tance sampling. Advances Neural Information Processing Systems, 31, 2018.
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt Kusner. No train gain: Revisiting efficient training algorithms for transformer- language models. arXiv preprint arXiv:2307.064400 2023.
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In International conference on machine learning, pp. 2525-2534, PMLR, 2018.
Conglong Li, Zhewei Yao, Xiaoxia Wu, Minjia Zhang, and Yuxiong He. Deepspeed data efficiency: Improving deep learning model quality and training efficiency via efficient data sampling and routing. Xrviv preprint 2rX1v:2212.03597, 2022.
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang. Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for generic network architectures. In International Conference on Machine Learning, pp. 14139-14152. PMLR, 2022.
Ilya Loshchilov and Frank Hutter. Online batch selection for faster training of neural networks. arXiv preprint arXiv:1511.06344 2015.
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training, 2018.
Sören Mindermann, Jan M Brauner, Muhammed Razzak, Mrinank Sharma, Andreas Kirsch, Win- nie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Confer- on Machine Learning, pp. 15630-15641 PMLR, 2022.
Deanna Needell, Rachel Ward, and Nati Srebro. Stochastic gradient descent, weighted sampling, and the randomized kaczmarz algorithm. Advances in neural information processing systems, 27, 2014.
Mahdi Nikdan, Tommaso Pegolotti, Eugenia lofinova, Eldar Kurtic, and Dan Alistarh. Sparse- prop: Efficient sparse backpropagation for faster training of neural networks. arXiv preprint arXiv:2302.04832 2023.
NVIDIA. Accelerating inference with sparsity using the nvidia ampere archi- tecture and nvidia tensorrt
aaceleraatig inference with https::/developer..nut nvidia. cww/blog/ 2021.
ampere and tensorrt
Xu Ouyang, Shahina Mohd Azam Ansari, Felix Xiaozhu Lin, and Yangfeng Efficient model finetuning for text classification via data filtering. arXiv preprint 1rXiv:2207.14386 2022.
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Find- ing important examples early in training. Advances in Neural Information Processing Systems, 34:20596-20677, 2021.
Md Aamir Raihan and Tor Aamodt. Sparse weight activation training. Advances in Neural Infor- mation Processing Systems, 33: 16625-15688, 2020.
David E Rumelhart, Geoffrey Hinton, and Ronald J Williams. Learning representations by back- propagating errors. nature, 323(6088) 1986.

---

Published as a conference paper at ICLR 2024
Vatsal Shah, Xiaoxia Wu, and Sujay Sanghavi. Choosing the sample with lowest loss makes sgd bust. In International Conference on Artificial Intelligence and Statistics, pp. 2120-2130. PMLR, 2020.
Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, and Yuxiong He. Random-ltd: Random and layerwise token dropping brings efficient training for large-scale transforeers. arXiv preprint arXiv:2211.11588 2022.
Minghe Zhang, Chaosheng Dong, Jinmiao Fu, Tianchen Zhou, Jia Liang, Jia Liu, Bo Liu, Michinari Momma, Bryan Wang. Yan Gao, al. Adaselection: Accelerating deep learning training through data subsampling. arXiv preprint arXiv:2306. 10728, 2023.
Minjia Zhang and Yuxiong He. Accelerating training of transformer-based language models with progressive layer dropping. Advances in Neural Information Processing Systems, 33:14011- 14023, 2020.
Peilin Zhao and Tong Zhang. Stochastic optimization with importance sampling for regularized loss minimization. In international conference on machine learning, pp. 1-9. PMLR, 2015.
A ABLATION ON HYPERPARAMETERS
There are few hyperparameters in self-adaptation algorithm, but all of them have explicit meaning. In this section we show that though extra tuning of these hyperparameters may achieve slightly better result, overall VCAS is robust to these hyperparameters with reasonable values. We conduct ablation experiments on two tasks: BERT-ber bERT-base finetuning on SST-2 and MNLI. All the results are averaged over 3 different seeds.
A.1 ACTIVATION AND WRIGHT VARIANCE THRESHOLDS Tact Tar
The main hyperparameters in VCAS is the variance thresholds of activation Toet and weight Tw- For these two thresholds, how to split total variance among them big problem with optimal solution differing across models and tasks. So without prior knowledge introduced, compromise by keeping Tact To < 1.
We further conduct ablation MNLL From the results we can find that a satisfactory outcome is assured regardless of the specific value of provided that & 1, which proves the robustness of VCAS.
from 0.01 0.5 as is shown in Tab. for SST-2 and Tab. for
Table Ablation on different variance thresholds 7 of BERT-base finetuning on SST-2
0(exact)
0.01
0.025
0.05
0.1
0.25
0.5
Final Train Loss
0.0559 92.59
0.0586
0.0600
0.0625
0.0642
0.0705
0.0761
Final Eval Acc(%)
93.07
93.04
93.25
92.81
92.79
92.18
FLOPPs reduction (%)
45.92
48.28
49.82
50.05
51.57
52.71
Table 5: Ablation on different variance thresholds т of BERT-base finetuning on MNLI
O(exact)
0.01
0.025
0.05
0.1
0.25
0.5
Final Train Loss
0.2372 84.33
0.2388
0.2428
0.2459
0.2552
0.2684
0.2805
Final Eval Acc(%)
84.31
84.23
84.33
84.07
84.13
84.08
FLOPs reduction(
38.59
41.56
43,49
45.37
47.53
48.92
A.2 MONTE-CARLA REPETITIONS M
To calculate variances, VCAS introduces an overhead of extra iterations quadratic with Monte-Carlo repetitions M.
12

---

Published as a conference paper at ICLR 2024
V.
M- 10
Vea
M 10
Vy
M-
M-
-5
Ve
M-
Ve
Mi
V
M=
Vae
Vw
M
M-2
M
104
10a
201
variance
variance
100
10
10
10
2000
4000
6000
10000
20000 steps
30000
steps
Figure 7: Variance calculated with different Monte-Carlo samples M of BERT-base finetuning on SST-2.
Figure 8: Variance calculated with different Monte-Carlo samples M of BERT-base finetuning on MNLI.
Obviously bigger M will bring more precise empirical variance, yet the cost is prohibitive.
We experiment on different M from to 10 and find significant difference in the empirical variance as is shown in Fig. for SST-2 and Fig. S for MNLI. Therefore, we adopted the setting of M with which only need to perform 6 extra iterations that is negligible if the variance calculation frequency is large enough like 100 in SST-2 and 500 in MNLI.
A.3 VARIANCE CALCULATION FREQUENCY
Similar M, the variance calculation frequency is also trade-off between better empirical approximation and less overhead introduced. We experimented on F 50, 100, 200, 500, 1000 in Tab. for SST-2 and Tab. 7 for MNLI. We can see that although as grows larger the overhead of VCAS is gradually relieved, with too large F, like F = 1000 in SST-2 that leads to only times of self-adaptation update, the sample ratio schedule not fully explored and the final FLOPs reduction is even smaller. Therefore, for all these tasks we set F to be least 1/50 of total training steps and no more than 500 due to slight marginal gains.
Table Ablation on different adaptation frequency F of BERT-base finetuning on SST-2, the num-
ber of training steps is 6315.
F Final Train Loss Final Eval Acc(%)
0(exact)
50
100
200
500
1000
0.0559 92.59
0.0589
0.0600
0.0587
0.0577
0.0562
92.71
93.04
92.56
93.15
93.19
FLOPs reduction(%)
47.33
48.28
46.06
39.43
31.03
Table Ablation on different adaptation frequency of BERT-base finetuning on MNLI, the num- ber of training steps is 36816.
F Final Train Loss
O(exact)
50
100
200
500
1000
0.2372
0.2460
0.2461
0.2440
0.2428
0.2428
Final Eval Acc(%) %)
84.33
84.20
84.23
84.12
84.23
84.21
FLOPSs reduction(%)
35.16
39.58
41.31
41.56
39.43
A.4 UPDATE STEP a AND WEIGHT RATIO MULTIPLIER
A simple grid search is conducted for f0.005, 0.01, 0.02) and {0.95, 0.9, 0.8} in Fig. for SST-2 and Fig. [0 for MNLI. From the figures, we can find that we are able to trade convergence for efficiency with more aggressive setting of larger and smaller B, yet all results here are decent
13

---

Published as a conference paper at ICLR 2024
with final accuracy drop of more than 0.3% for both tasks. Thus, VCAS is robust different and B.
93.2
0.0610
0.0597
0.0601
0.02
92.93
92.60
92.82
0.02
50
0.01
0.0600
0.0592
0.0599
0,0600
0.01
93.04
92.89
92.48
92.8
92.6
0.005
0.0591
0.0995
0.005
93.12
93.04
92.74
46.27
47.25
48.75
46
0.95
(a) Train Loss
(b) Eval Aco(%)
(c) FLOPs reduction(
Figure 9: Grid search of update step and weight ratio multiplier of BERT-base finetuning SST-2. The darker color the better.
0.02
0.2462
0.2451
0.2456
0.02
84.12
84.21
84.19
84.25
0.02
42.05
42.74
43.98
43
0.01
0.2428
0.2426
0.2457
0.01
84.23
84.29
84.17
0.01
42.89
43.51
84.15
0.005
844.10
42.38
0.242
0.95
0.9
0.8
0.95
0.9
0.95
0.9
0.8
(a) Train Loss
(b) Eval Acc(%)
(c) FLOPss reduction %
Figure 10: Grid search of s update step a and weight ratio multiplier 3 of BERT-base finetuning on MNLI. The darker color the better.
From all the ablation results above, we can see that VCAS robust to these hyperparameters with reasonable values, proving the insensitiveness.
B INSIGHTS ON UPDATE OF 8, {pt} AND {n}
In this section, we will show how the gradient norm preserving ratio as well as all the sample ratios {pi) and {vr} update across the training.
We record the update process of BERT-base finetuning on MNLI with different variance tolerance
thresholds in Appendix. A. All results are averaged three different seeds.
Fig. Ia depicts the update of s. For non-decreasing {pr). J. we plot the update of the first and the last values ,1, Pr. Fig. lbb with other values lying between. For {M1). we show the update of the first three ones V/, 1/ in Fig. II and observe similar behavior of other weights.
is seen Fig. [I that during training of BERT-base on MNLI, the gradient norm preserving ratio first decreases and then shows slight downward trend. The activation sample ratios {pr} gradually decrease with an abrupt change between epochs due the rapid decline of train loss caused by the lowered learning rate in the linear learning rate scheduler. The weight sample ratios {vr} first decrease and then fluctuate match the change of activation sample ratios.
C PERFORMANCE ON CNN
In Sec. 6 mainly experiment with Transformer-based models and Adam optimizers. But the variance controlled adaptation depicted Sec. holds universally for any DNNs with SGD-based optimizers, since it just provides an approximated stochastic gradient with controlled variance to estimate the full gradient. this section, we employ VCAS on other architectures and other opti- mizers to prove its versatility.
14

---

Published as a conference paper at ICLR 2024
r=022
T-
2.
8 62022
0.025
1-00
0.8
0.90
0.6-
0.80
0.2
0
10000
20000 Staps
30000
10000
20000
20000
10000
20000
50000
steps
steps
(a) Update of 8
(b) Update of (pt}
(c) Update of (n)
Figure I1: VCAS update process with different for BERT-base finetuning MNLI.
For CNN, it noted that the weight sampler SampleW in Sec. designed for linear layers is not us- able for convolution layers. Thus employ VCAS with degraded version of activation sampling only.
We experiment with WideResNet- 18 with widen factor 4 pretraining ImageNet. We use eight NVIDIA 3090Ti to parallel the training with Distributed Data Para)lel(DDP). We employ SGDM optimizer with momentum m 0.9. The results in Tab.
Table 8: Training results of WideResNet-188 pretraining on ImageNet with NVIDIA 3090Ti.
Method
Train Loss
Eval Acc(%)
Train Time(h)
FLOPs((%)
TimeJ(%)
exact VCAS
1.474 1.479
75.96 75.86
21.31 20.20
17.47
5.21
From the table we can VCAS is also capable of accelerating the training of CNN. Besides, the parallel setting also proves the parallelizability of VCAS. The relatively low but still decent time reduction can be explained with Amdahl's Law since VCAS only accelerate the calculation part and is not able to accelerate other parts like communication cost during parallel training.
D DETAILS ABOUT ALGORITHM.
should be noted that some parts of Alg. are simplified for clarity and we list the implementation details below:
In the algorithm table, we put the calculation of empirical variances out of the two Monte-Cart loops for simplicity. Yet practically we can calculate Vact and V., inside the loops and average the variance scalars outside. Therefore, we only need to store three tensors additionally regardless of M: SGD gradient Gs.; to calculate Vac.. and its running mean and running square mean to calculate V.. By sampling only part of parameters to keep gradients, like 1 in our experiments, the memory overhead can be neglected.
Besides, since weight sample ratios {n) are updated parameter-win according to Eq, 7 the empir- ical weight variances and SGD variances also stored parameter-wist wisee when implemented.
Update of activation sample ratios {pi} requires finding out gradient sparsity {pt} with the new s according to Eq. In implementation, this is achieved by calculating possible new {pr} with both + a and s a inside the Monte-Carlo loops and averaging them outside. Then just choose the proper one with
15

---

Published as a conference paper at ICLR 2024
E PROOF
E.
PROOF TO UNBIASEDNESS OF VCAS
Let's first consider a L-layer MLP. (Note: for simplicity we mildly abuse the term "layer" here, representing a single operation like matrix multiplication and ReLU)
For the last layer L, the output gradient V g(t) is caleulated from the loss directly, the same as the Exact BP. Since activation sampler V 24) SampleA eipo. (Vza)) unbiased, we have:
z(L)
When back propagation proceeds, we may encounter types of layers: linear and non-linear, For the linear layer, we have:
V 2c = ((2:((L)
Thus unbiasedness preserved with the output gradient of the (L - 1)-th layer:
Vz
E
4 2 g
Exact BP result
While for the non- non-linear layer like ReLU, we have:
V zıt-n Vz( Jz(
where is the Hadamard product and Jzuz is the Jacobbi matrix determined by ( which is saved in forward pass and is exact. Thus again we derive the the output gradient of the (L - 1)-th layer being unbiased:
E ZVzu-1] E Z 2 Jzun = V 2(2) 2 J2(1) Exact BP result
Thus by induction, VCAS assures activation gradients V z 1.... being unbiased.
Then for weight gradients, since weight sampler V z0 SampleW smm (V20n,Z(-1) is unbi- ased, we have:
E
E
V
2
2
Finally, we derive all weight gradients being unbiased:
E p0))
Z
(
V (
Exact result
For more complicated neural networks like CNN and Transformer, since operations like convolu- tions and layernorm are all linear transforms, by similar reasoning the unbiasedness still holds.
EXPERIMENT DETAILS
BERT-BASE PRETRAINING
For BERT-base pretraining same as the original settings of day training on a single NVIDIA 2080Ti. The full results are as follows in Tab. g
crammed BERT in Geiping & Goldstein (2022) with the recipe
From the table we can find that although VCAS achieves relatively high train loss, the downstream task performance is still competent with exact training. While SB and UB both perform worse on CoLA, which is a vulnerable task, reflecting that they have changed the original convergence trajectory of SGD.
16

---

Published as a conference paper at ICLR 2024
Table Full results BERT-base pretraining
Methods
Loss
MNLI-m
MNLI-mm
QQP
QNLI
SST2
CoLA
STSB
MRPC
RTE
Avg.
exact
2.099
82.28
82.68
87.08
88.85
91.28
48.07
83.26
86.98
54.87
78.37
SB
2.133
82.34
82.86
87.27
88.63
91.28
41.82
82.86
85.53
55.23
77.53
UB
2.106
82.95
83.46
87.27
88.66
91.05
42.80
83.68
85.90
55.95
77.96
VCAS
2.134
82.03
82.82
86.92
89.23
91.62
48.36
83.02
86.03
55.23
78.36
F.2
RECIPE OF OTHER TASKS
For BERT finetuning, we use AdamW optimizer with Ir = 26-5 and wd 0.01. The learning rate scheduler is linear one with warmup-ratio 0.1. We set epoch numbers N and batch size of batch_size 32.
For VIT finetuning, we use Adam optimizer with Ir 2e A linear Ir scheduler with warmup employed. We run = 5 epochs with batch size batch_size = 32
G LIMITATIONS
VCAS is designed for adaptively learning the proper sample ratios of large model training on large datasets. It is not suitable for small models with low gradient variances resulting increased nu- merical errors, or small datasets with few training steps that is insufficient for the update process in VCAS.
The weight sampler SampleW in VCAS is specially designed for linear layers and is not usable for other operations like convolution. But the activation sampler SampleA can be applied all mainstream architectures with deep layers. So for CNN or RNN, we need to employ a degraded version of VCAS with activation sampling only, as shown in Appendix. C
VCAS focuses on mirroring the exact training with theoretical guarantee and is lack of exploration of other possible convergence trajectories that may bring better result. Thus not recommended when the original training recipe is under-optimized.
17


Reviews:
[{"review_id": "bcgPtXnW8T", "reviewer": "Reviewer_MWbo", "summary": "The authors propose a sampling method for back propagation with controlled variance and self-adaptive sample ratios, named VCAS. It computes an approximate stochastic gradient by applying finegrained sampling to gradually remove samples and tokens during backpropagation. VCAS have similar variance as well accuracy with exact back propagation, while seems to reduce the training cost significantly.", "questions": "1. Although the proposed VCAS algorithm seems promising compared with SB and UB, I'd like to know the actual reduction in training time as these sampling methods might introduce overhead in computation. It would be helpful if the authors can report a time table for training on these datasets.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. I like it very much that the authors test the proposed algorithm on multiple fine-tuning and pre-training tasks in both vision and\nnatural language domains. Typically, papers in this domain would use small scale datasets such as MNIST or CIFAR 10/100.\n\n2. The ablation studies on a few hyperparameters are very important. I see only very few of the papers in this domain have done this before.", "weaknesses": "I think the authors can improve the paper in the following ways:\n\n1. I believe adding an algorithm table with detailed steps would make the paper more clear. \n\n2. The authors report the Final Train Loss / Final Eval Acc.(%) / FLOPs reduction ratio(%). However, I'd like to know the actual reduction in training time as these sampling methods might introduce overhead in computation. It would be helpful if the authors can report a time table for training on these datasets.\n\nTo be frank, I feel not many papers actually do this but it can be interesting to see that the actual training time might not be reduced at all, or at least not much as expected given a certain sampling ratio.\n\n3. The paper lacks discussions of related paper. For example, https://arxiv.org/pdf/2104.13114.pdf also considers the importance sampling problem by sampling data points proportionally to the loss, instead of norm of gradient.\n\nFor another example, https://arxiv.org/pdf/2306.10728.pdf also proposes adaptively sampling methods for dynamically selecting data points for mini-batch. I'd love to see the authors discussed more about these papers.\n\n4. Can the authors be more specific in terms of the notations? Adding a table of notation would be very helpful. For example, what is $h^{(l)}$ below:\n\n$$\\nabla_{Z^{(l-1)}}=h^{(l)}\\left(\\nabla_{Z^{(l)}} ; Z^{(l-1)}, \\theta^{(l)}\\right)$$", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8AW0IQm6oN", "reviewer": "Reviewer_d6vC", "summary": "This paper proposes Variance-Controlled Adaptive Sampling (VCAS), which performs an approximated stochastic gradient with an adaptive sampling rate. Based on the insight that gradients are sparse after learning has progressed to some extent, the authors improve the efficiency of learning by computing only a few selected gradients through adaptive sampling. The proposed method approximates the exact backpropagation values well in BERT and ViT training.", "questions": "I would like to see more detail in Table 2 or 3. What is the relationship between Final Eval Accuracy and FLOPs reduction? For example, is the recommending FLOPs reduction ratio for VCAS around 40%?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "VCAS performs backpropagation 20-50% faster, while following the loss curve of true full backpropagation with low variance. VCAS has an adaptive sampling rate, which allows for efficient sample selection based on learning loss and per layer. The idea is simple and highly applicable.", "weaknesses": "Comparing accuracy under the same amount of FLOPs reduction makes it difficult to understand its effectiveness compared to a metric like time to reach target accuracy[1]. As a result, it is unknown how VCAS will perform under a 50% or greater reduction.\n\n\n[1] Mindermann, S., Brauner, J. M., Razzak, M. T., Sharma, M., Kirsch, A., Xu, W., ... & Gal, Y. (2022, June). Prioritized training on points that are learnable, worth learning, and not yet learnt. In International Conference on Machine Learning (pp. 15630-15649). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0Fum3cRgiS", "reviewer": "Reviewer_Dd1x", "summary": "This work introduces a variance-controlled adaptive sampling (VCAS) method for accelerating the back-propagation of deep neural network training. VCAS computes unbiased, variance controlled gradients for both activations and network weights. By sampling both data samples and tokens in each datum in a layer-wise, fine-grained manner, VCAS can drastically reduce the computation in the back-propagation process without introducing too much variance overhead. With the similar FLOPs reduction, VCAS better optimizes the target model compared with prior loss-based and gradient-based sampling methods.", "questions": "- It is not directly clear to me whether the weight gradient sampling of VCAS is applicable to convolution neural networks (CNN). In principle, convolution is yet another linear operator, but I’m not sure how to perform this sampling in a convolutional layer. Similarly, can VCAS be applied when optimizing a recurrent neural network (RNN)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This work introduces a fine-grained strategy that 1) increasingly removes data samples when back-propagating to the input layer, and 2) samples tokens in each data sample when computing gradients for network weights. This fine-grained strategy allows high FLOPs reduction with controlled variance.\n\n- The sampling ratios are adaptively adjusted according to the estimated variance. As training progresses, the network may require changing sampling ratios, and the adaptive sampling ratios can better satisfy this need.\n\n- Compared with prior methods, the proposed method, VCAS, can better simulate exact back-propagation, leading to better optimized loss and better evaluation performance.", "weaknesses": "- Training time reduction: When comparing with baselines, this work uses FLOPs as the efficiency indicator. However, FLOP reduction may not directly translate into wall-clock time reduction due to various factors like parallel computation efficiency. It is suggested to also list the wall-clock training time of each method for a more straightforward comparison.\n\n- Insights on sampling ratio updates: In Section 7 this work has discussed the design choices that determine the sampling ratio updates. For better comprehension, it may be useful to include a figure that shows how $s$ and $\\nu_l$ changes as training progresses.\n\n- Figure/Table clarity: Figure 2 seems to lack some more detailed explanation. In Table 1, it is not clear which numbers should be bolded. For example, for ViT-base fine-tuning on CIFAR-100, UB seems to be highlighted for the highest eval accuracy, but for ViT-large fine-tuning on CIFAR-100, UB seems to be highlighted for the lowest train loss? Also for Table 1, how significant is the relative improvement over baselines?\n\n- Limitations: It is suggested to include some detailed discussion on the limitations (e.g., applicable model architectures, base optimizer, dataset) of the proposed method. In this paper, only Transformer-based architectures and Adam-like optimization algorithms are tested. It is not clear whether we can extrapolate the conclusion to other settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NLCBbtuxhP", "reviewer": "Reviewer_PsQT", "summary": "This manuscript proposes a sampling algorithm to eliminate computation during forward and/or BP. More specifically, a variance-controlled adaptive sampling (VCAS) method is designed to accelerate BP by computing unbiased stochastic gradients.\n\nThe effectiveness of the VCAS is justified by pre-training and fine-tuning tasks in both vision and NLP domains. Some ablation studies are also included to discuss the effects of hyper-parameters.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* This manuscript is well-structured, with a clearly explained methodology section.\n* The manuscript evaluates the effectiveness of the VCAS on pre-training and fine-tuning tasks in both vision and NLP domains.\n* An ablation study is also included.", "weaknesses": "1. Limited literature review. This manuscript did not carefully examine the relevant papers, and some closely related papers were omitted from the discussion, e.g., [1, 2, 3] and the related work in [1].\n2. Limited baseline evaluations. Close baseline methods e.g., [2, 3] should be considered.\n3. Additional hyper-parameters and insufficient ablation studies. The ablation study on one task, i.e., fine-tuning BERT-base on MNLI, is insufficient to justify the insensitivity of these hyper-parameters.\n4. Clarity. Some design choices in Section 5 are just given and have no explanation. For example, how to derive the provided equation from the mentioned zeroth-order algorithm? \n5. The claim on the unbiased stochastic gradient needs to be more careful and some theoretical justifications should be provided.\n\n### Reference\n1. SparseProp: Efficient Sparse Backpropagation for Faster Training of Neural Networks at the Edge\n2. Back Razor: Memory-Efficient Transfer Learning by Self-Sparsified Backpropagation\n3. Sparse weight activation training", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)