Paper: Efficient Backpropagation with Variance Controlled Adaptive Sampling (ICLR-2024)

Reviewer comment (Novelty):
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].

Author response:
**Q1. Limited literature review on sparsity-based training methods**

Thanks for pointing out relevant works on sparse training. However, we'd like to place our work on another closely related field of ***online batch selection*** that discards less informative data, but not prunes weights or neurons as is typical in sparse training. The main difference comes in two ways:
- Wall-clock time reduction: Online batch selection methods can always translate FLOPs reduction into time reduction since they neatly wipe the entire datum out. While few sparse training works report time reduction usually for the lack of supportive hardware, especially for unstructured sparsity.
- Biasedness: In online batch selection, importance sampling is usually adopted which assures unbiasedness. While in sparse training, pruning methods typically cannot guarantee unbiased gradients. Using unbiased gradients for SGD is critical for the algorithm to converge properly.

Moreover, sparse training and online batching selection should be orthogonal since they reduce computation along different dimensions (data vs. model) and the speedup could be multiplicative. 

Nevertheless we agree with the reviewer that the paper can be improved with more discussion on sparse training methods. We have added the discussion to Related Work.

Reviewer comment (Experiments):
Limited baseline evaluations. Close baseline methods e.g., [2, 3] should be considered.

Author response:
**Q2. Limited baseline evaluations**

VCAS can translate FLOPs reduction into time reduction which is rare for sparse training methods. We compare VCAS with the official codebases of methods mentioned by the reviewer:

For Back Razor, we tested on BERT-large finetuning on MNLI. We find though Back Razor can reduce up to 90\% FLOPs, it cannot reduce training time with its unstructured sparsity, and the accuracy drops a lot.

**Tab. Comparison of VCAS with Back Razor**
|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs Reduction(%)|Time Reduction(%)|
|-|-|-|-|-|-|
|Exact|0.1439|86.58|5.48|-|-|
|Back Razor|0.3827|84.26|5.19|90.00|5.29|
|VCAS|0.1619|86.63|4.44|44.17|19.03|

SWAT seems close to our work since it conducts a fine-grained pruning on weight and activation and implements a structured sparsity version that uniformly prune channels.

Since the official codebase of SWAT only supports CNNs, we compare SWAT with VCAS on WideResNet18 with widen factor $k=4$ on CIFAR100. We find VCAS can reduce training time while preserving accuracy, while SWAT is even slower than exact (dense) training due to its high overhead. 

**Tab. Comparison of VCAS with SWAT**
|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs Reduction(%)|Time Reduction(%)|
|-|-|-|-|-|-|
|Exact|0.0070|73.25|4.32|-|-|
|SWAT|0.0115|72.01|4.81|50.00|-11.49|
|VCAS|0.0070|74.16|3.81|19.05|11.72|

Reviewer comment (Experiments):
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.

Author response:
**Q3.Additional hyperparameters and insufficient ablations**

First we'd like to emphasize that we employ ***the same set of hyperparameters in all our experiments*** across pretraining and finetuning, CV and NLP, and all get decent FLOPs reduction and accuracy. We believe it can already reflect the insensitivity. As said by Reviewer MWbo in "Strength", "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". 

From the experimental perspective, we agree with the reviewer that more experiments will help to justify our statement. So we have done ***all ablations again for BERT-base finetuning on SST-2***. Please refer to Appendix A for the results.

Reviewer comment (Theory):
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?

Author response:
**Q4. Clarity of design choices in Section 5**

Sorry for the confusion. To make our algorithm clearer we've added an algorithm table at the end of Section 5. Hope it can relieve your puzzle.
> How to derive the provided equation from the mentioned zeroth-order algorithm?

The zeroth-order algorithm is embarassingly simple: since the variance increase monotonically with the drop ratio, our algorithm simply decrease the drop ratio by a constant if the variance is too large, and increase the drop ratio if the variance is too small. The only learning signal utilized is whether the variance is higher than the desired threshold. For a more intuitive view of the update, please refer to figures in Appendix C

Reviewer comment (Theory):
The claim on the unbiased stochastic gradient needs to be more careful and some theoretical justifications should be provided.

Author response:
**Q5. Lack of proof on the unbiasedness of VCAS**

Thank you for posing the vital problem of the unbiasedness of VCAS! We've added our proof in Appendix E.

Proof sketch: since our sampling is unbiased, we just need to prove the unbiasedness preserves in BP. As BP proceeds, linear operations like matmul, convolution and layernorm can obviously preserve unbiasedness. Nonlinear operations like ReLU incur Hadamard product with Jacobian matrix that is determined in forward pass, and thus also preserve unbiasedness. By induction we derive VCAS being unbiased.

Reviewer comment (Evaluation):
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.

Author response:
**Q1. Training time reduction.**

Thanks for the suggestion. To compare wall-clock training time, we choose two tasks with the longest training time: BERT-large finetuning on MNLI and ViT-large finetuning on ImageNet. The results are as follows:

**Tab. BERT-large finetuning on MNLI**
|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs Reduction(%)|Time Reduction(%)|
|-|-|-|-|-|-|
|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|

**Tab. ViT-large finetuning on ImageNet**
|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs Reduction(%)|Time Reduction(%)|
|-|-|-|-|-|-|
|exact|04135|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|

We find ***VCAS can translate FLOPs reduction into wall-clock time reduction as effectively as simpler online-batch sampling methods*** like UB and SB that drop part of data one-time in a whole, while enjoying mirrored performance with the exact training under theoretical guarantee.

Our current implementation of VCAS is purely based on Python and PyTorch. The gap between FLOPs reduction and time reduction can be further narrowed by optimizing the implementation with CUDA or Triton.

We have added the wall-clock time results to Section 6.2, thanks again for the suggestion!

Reviewer comment (Presentation):
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.

Author response:
**Q2. Insights on sampling ratio updates.**

Thanks for the constructive suggestion. We've added a few figures in Appendix B recording the update process of our sample ratios. For all variance thresholds $\tau$ from 0.01 to 0.5, gradient norm preserve ratio $s$ first drops rapidly and then slowly. Activation ratios $\rho_l$ gradually decrease with training. Weight ratios $\nu_l$ first drop and then fluctuate around a certain value that is different across $l$.

Reviewer comment (Presentation):
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?

Author response:
**Q3. Figure/Table clarity.**

Sorry for the confusion. Now we've added more explanations in Fig.2. The bold formatting in Tab.1 originally highlights the highest *Eval Acc / (1-FLOPs reduction)* which is a balance between performance and efficiency. Now we've changed it to indicate the best of each metric for better clarity. In Tab.1 we also added underlines to mark Eval Acc that is less than 0.1% off the exact training for a more intuitive look.

**Q4. For Table 1, how significant is the relative improvement over baselines?**

As in Tab.1, the FLOPs improvement is relatively marginal compared with UB and SB in our *conservative settings* of variance tolerance ratio $\tau=0.025$. But the ***accuracy improvement is significant***, where the train loss and eval accuracy of VCAS are much closer to exact training.

However, it should be noted that the usage of VCAS is different from existing sampling baselines like UB and SB. For UB and SB, we need to preset a sample ratio $r$ carefully and pray for a decent outcome. While for VCAS, we just need to set a variance tolerance ratio $\tau$ far less than 1, and then it can adaptive learn the proper fine-grained sample ratios which can keep the training curve closely to that of exact training. Therefore, two more important improvements are that ***VCAS is tuning-free for sample ratios*** and ***the FLOPs reduction achieved by VCAS is "safer" with theoretical guarantee*** compared to other baselines.

Reviewer comment (Experiments):
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.

Author response:
**Q5. Limitations should be discussed.**

Thanks for the constructive advice. We've added a detailed discussion of the limitations of VCAS in Appendix G. 

Regarding to the limitations on applicable model architecture, base optimizer, dataset, theoretically, there will not be such limitations. But the performance of VCAS is closely related with the "redundancy" of architectures and datasets. Nevertheless, in such scenario, VCAS can learn to keep all data to control the variance (i.e., fall back to exact training), and the accuracy will not be harmed. On the other hand, other online batch sampling methods are likely to fail.

Reviewer comment (Reproducibility):
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)?

Author response:
**Q6. Whether can VCAS apply to other architectures?**

Yes. We apply VCAS to train WideResNet-18 on ImageNet with Momentum SGD, i.e., we generalize it beyond Transformers and beyond Adam. For CNNs, the weight gradient sampling part is not applicable, since the "tokens" (i.e., "pixels") in CNNs are not individually processed, but tightly coupled. Therefore, we only apply activation gradient sampling and adapataion of VCAS.

We've conducted an experiment using 8\*3090Ti to pretrain WideResNet-18 with widen factor $k=4$ on ImageNet. We use SGDM with $momentum=0.9$. The results are as follows:

**Table. VCAS performance for training WideResNet-18 on ImageNet**
|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs Reduction(%)|Time Reduction(%)|
|-|-|-|-|-|-|
|exact|1.474|75.96|21.31|-|-|
|VCAS |1.479|75.86|20.20|17.47|5.21|

From the table we can see that VCAS is also capable of ***CNN acceleration***, and also usable for ***other optimizers like SGDM*** here. Besides, the parallel setting proves the ***parallelizability*** of VCAS.

Results above are added to Appendix C.

Reviewer comment (Presentation):
I believe adding an algorithm table with detailed steps would make the paper more clear.

Author response:
**Q1. Add an algorithm table.**

We appreciate it a lot for your constructive advice. We've now added an algorithm table in Algorithm 1 at the end of Section 5. Thank you!

Reviewer comment (Writing):
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:

$$\nabla_{Z^{(l-1)}}=h^{(l)}\left(\nabla_{Z^{(l)}} ; Z^{(l-1)}, \theta^{(l)}\right)$$

Author response:
**Q4. Clarity of notations.**

Sorry for the confusion. $h^{(l)}$ means the function of calculating input gradient with output gradient, input and weight. For example for linear layers, we have input gradient equals to output gradient multiplied by weight. The same with $g^{(l)}$ for weight gradient.

We have now added more explanations on every notations in Section 4 and 5. Hope it can relieve your puzzle!

Reviewer comment (Evaluation):
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.

Author response:
**Q2. Report wall-clock time reduction.**

Thanks for your insightful advice. As you can see, not many papers report the wall-clock time reduction, mainly because the wall-clock time is highly dependent on hardware and implementation. FLOPs reduction, however, is a high-level metric directly reflecting the real power of algorithm. Each metric has its own pros and cons.

Though, we agree that wall-clock time is indispensable since that's what we really want to accelerate and FLOPs reduction can not always translate into equal wall-clock time reduction, not only because of the overhead introduced, but also due to parallelizability, hardware friendliness, and compatibility to many other minor but crucial parts in the training procedure.

For this metric we choose two tasks with the longest training time, which are most worthy of acceleration: BERT-large finetuning on MNLI and ViT-large finetuning on ImageNet. The results are as follows:

**Tab. BERT-large finetuning on MNLI:**

|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs reduction(%)|Time reduction(%)|
|-|-|-|-|-|-|
|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|

**Tab. ViT-large finetuning on ImageNet:**

|Method|Train Loss|Eval Acc(%)|Train Time(h)|FLOPs reduction(%)|Time reduction(%)|
|-|-|-|-|-|-|
|exact|04135|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|

From these tables, we can find that VCAS can translate FLOPs reduction into wall-clock time reduction ***as effectively as simpler online-batch sampling methods*** like UB and SB that drop part of data one-time in a whole, while enjoying mirrored performance with the exact training under theoretical guarantee.

Our current implementation of VCAS is purely based on Python and PyTorch. The gap between FLOPs reduction and time reduction can be further narrowed by optimizing the implementation with CUDA or Triton.

We have added the wall-clock time results to Section 6.2, thanks again for the suggestion!