DropBP: Accelerating Fine-Tuning of Large Language Models by Dropping Backward Propagation

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q3.1. Pervious work like LayerDrop and others omit the layer computation in the forward pass. Then the computation could be removed in the subsequent backward computation with essential engineering efforts. Thus it lacks a clear distinction in terms of technical innovation compared to these previous works.

Q3.2. It's unclear why the forward pass computation remains unchanged.

A3.1-2. We appreciate the reviewer's insightful comments. We understand the reviewer's questions as follows:

1. Motivation: What is the reason for focusing on reducing computations in the backward propagation rather than the entire computation process?
2. Methodology: Why does the approach omit only backward propagation computations while keeping forward propagation unchanged?
3. Related Works: How does this differ from algorithms like LayerDrop [2], which drop computations during forward propagation as well?

First, the goal of LayerDrop is to ultimately perform layer-wise pruning to reduce inference time by dropping layers during forward propagation. In contrast, the aim of our DropBP is to enhance training efficiency directly by reducing FLOPs and training memory usage. Therefore, our approach does not require strictly dropping the forward path. Of course, while it is possible to drop the forward propagation, we chose to avoid this because the training process is highly sensitive to dropping the forward propagation, as shown in Fig. E.

• Figure E in the attached PDF.

In Fig. E, we compared the loss curves of DropBP and Progressive Layer Dropping (PLD) [3], a representative layer dropping algorithm, when the same drop rate was applied to training computations. The results demonstrate that DropBP achieves much more stable training compared to PLD. This is because dropping the forward path can cause output deviations, which can negatively impact loss and all gradients.

Furthermore, in Appendix D of the manuscript, we compared our DropBP with Layer Dropping algorithms for fine-tuning LLMs. As shown in Table 7 of the manuscript, our DropBP achieved higher accuracy compared to traditional Layer Dropping algorithms.

In summary, while Layer Dropping algorithms aim to improve inference efficiency by randomly dropping the forward path and pruning it during inference, our DropBP focuses on selectively dropping backward paths, that are relatively less sensitive than the forward path, to accelerate training and reduce memory usage while maintaining high accuracy.

Following the reviewer's advice, we will incorporate these clarifications into the paper to reduce confusion and strengthen the contribution.

Q3.3. The faster convergence observed in Figure 5 with DropBP compared to the vanilla model is counterintuitive. The observation here quite confuses me since the backward pass optimizes a partial computation graph, concerns regarding overfitting arise. The paper would benefit from a discussion on potential regularization techniques employed to address this, and a comparison with related work (e.g., [1]) that utilizes sublayer dropping for regularization in training a deep Transformer model.

A3.3. We agree with the reviewer's observation that Fig. 5 of the manuscript appears to show rapid convergence, which may suggest overfitting. However, this is because the x-axis is plotted against training time. When plotted against training steps, a completely different pattern emerges, as shown below:

• Figure B in the attached PDF.

As shown in Fig. B, when the drop rate is 0.5, the convergence of loss per step is almost identical to the baseline. However, with drop rates of 0.75 and 0.875, the convergence speed per step is slower. Nonetheless, DropBP significantly reduces the time consumed per training step because it skips the backward propagation computations for the dropped layers. Consequently, the convergence speed per training time is actually faster for DropBP compared to the baseline.

Moreover, our DropBP does not optimize a partial computation graph but instead randomly drops layers according to the drop rate, meaning the layers being trained change with each step. This can be interpreted as structured dropout, leading to an ensemble effect where multiple models are effectively trained simultaneously. As a result, it can serve as a regularization technique to mitigate overfitting issues. In fact, when applying DropBP, we did not observe overfitting, where training loss decreases but validation loss worsens.

Additionally, through [1], our DropBP can also be interpreted that DropBP reduces co-adaptation among layers. Specifically, [1] claims that Layer Dropping (LD) [2] successfully achieves regularization by reducing co-adaptation among layers through dropping layers during training. This effect is also applicable to DropBP, where some backward paths are randomly dropped instead of strictly training all layers.

Thanks to the reviewer's insight, we have confirmed that DropBP can also be interpreted as a regularization technique. In the revised manuscript, we will incorporate this analysis, and in future research, we plan to specifically analyze the regularization effects of applying DropBP and work on improving it.

Q3.4. some typos exist.

A3.4. hank you for catching the typo. We will make sure to correct it in the revised version.

[1] Li et al., "Learning Light-Weight Translation Models from Deep Transformer." [2] Fan et al., "Reducing Transformer Depth on Demand with Structured Dropout." [3] Zheng and He, "Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping"

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q2.1. Can you provide more details on the sensitivity calculation process? Specifically, how is the sensitivity of each layer computed, and what is the computational overhead associated with this step?

A2.1. Please referto the global response GA2 above.

Q2.2. The paper could benefit from a more detailed theoretical analysis of why DropBP works as effectively as it does. This would strengthen the paper by providing a deeper understanding of the underlying principles.

A2.2.

• Figure C in the attached PDF.

That's a good point. We interpret transformer models as a collection of numerous blocks, each composed of various modules with residual connections. Our hypothesis is that we can fine-tune LLMs well by training only certain shallow submodules. To theoretically analyze this hypothesis, we measured the impact of submodules based on their path lengths in LLaMA2-7B, as shown in Fig. C. Specifically, we followed these steps as suggested in [1]:

1. We first perform a forward pass through the entire network.
2. During the backward pass, we randomly sample $k$ residual blocks, which are back-propagated without passing through skip connections, while the remaining $n-k$ blocks are bypassed through the skip connections.
3. We then measure the norm of the gradient at the input.

We take 100 measurements for each path length $k$. Subsequently, we multiply by the distribution of all possible path lengths, which follows a Binomial distribution, to quantify the gradient contribution from paths of a specific length.

In Fig. C(b), we observed that the gradient per path length decreases as the path length increases. Consequently, Fig. C(c) demonstrates that shorter path lengths have a greater impact on the gradient. These observations are consistent with the findings in [1], which attributed this phenomenon to vanishing gradients. We confirmed that this also occurs in transformers, where the paths that significantly influence training in LLMs are relatively short. Therefore, DropBP enables effective training by focusing on these short submodules.

Thanks to the reviewer's advice, this analysis will be included in the final version, and we plan to conduct a more theoretical analysis of this phenomenon.

Q2.3. Are there guidelines or heuristics that practitioners can follow to optimize performance for their specific use cases?

A2.3. We are very pleased that the reviewer has shown interest in the use case of DropBP. Empirically, using the identical settings as the baseline is sufficient to achieve good convergence of loss and high accuracy when applying DropBP. At higher drop rates such as p=0.75 and p=0.875, however, increasing the learning rate by about 1.5 times can slightly improve accuracy. Thanks to the reviewer, we will incorporate these guidelines into the code we release.

Q2.4. How well does DropBP integrate with other recent advancements in efficient training techniques, such as mixed precision training or distributed training frameworks?

A2.4. We have developed a DropBP library that can be easily integrated into PyTorch, allowing it to be readily combined with most efficient training techniques that can be applied on a single GPU, such as parameter-efficient fine-tuning and mixed precision training. As shown in Fig. 4 and Table 2 of the manuscript, we incorporated these combinations into our experiments.

Additionally, we have confirmed that our library works well with recent distributed training frameworks based on PyTorch, such as FSDP [2]. However, we have encountered some errors when integrating with the DeepSpeed [3] framework. We are currently debugging these issues and plan to resolve them. In the future, we intend to analyze the experimental results once these issues are addressed.

[1] Veit and al., "Residual Networks Behave Like Ensembles of Relatively Shallow Networks." [2] Zho and al., "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel". [3] Rajbhandari and al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models".

We thank the reviewer for carefully reviewing our submission and providing valuable feedback. Please see below for our response to the questions and comments.

Q1.1. The idea of optimizing NNs with sparse gradient is not new.

A1.1. We acknowledge that the idea of optimizing neural networks with sparse gradients is not novel. However, our method differs significantly from sparse gradient methods like meProp, FISH Mask, and Freezeout in terms of purpose, methodology, and effect. Specifically:

• meProp: While meProp accelerates training by masking the output gradient of individual layers, our DropBP skips entire transformer blocks by leveraging skip connections. As shown in Table G.1.2 in global response, our DropBP achieves high accuracy compared to meProp. Furthermore, meProp requires computing top-K gradient masks at each iteration, whereas DropBP only needs to calculate the sensitivities once for the entire training process, reducing the overhead on training time. Additionally, meProp must store activations for all layers during training, while DropBP saves activation memory by not storing activations for dropped layers.
• FISH Mask: FISH Mask, which is a parameter efficient fine-tuning such as LoRA, reduces communication costs by updating sparse parameters without decreasing FLOPs. In contrast, our DropBP reduces computational costs directly by dropping backward operations. Furthermore, while FISH Mask have to store all activations for backward propagation, our DropBP eliminates the need to store activations for dropped layers, reducing activation memory. DropBP and FISH Mask employ complementary methods, which makes it possible to apply them simultaneously, just as DropBP and LoRA can be. We are working on this implementation but facing delays due to the complexity of FISH Mask's code. We will include related experiments in the paper after completion.
• Freezeout: Freezeout accelerates training by gradually freezing earlier layers, while DropBP randomly drops layers from the start, regardless of their order. Consequently, Freezeout requires storing all activations initially, which complicates increasing sequence length and parallel processing. In contrast, DropBP maintains low and consistent memory allocation as shown in Table G.1.2, facilitating easier management, longer sequences, and better parallel processing.

We will include these distinctions and a more comprehensive comparison to suggested related works in our revised paper.

Q1.2. It is unclear if the method works well for generation tasks and domain-specific transfer learning.

A1.2. Please refer to the global response GA1 above.

Q1.3. It is unclear which algorithm is used to solve the constraint minimization problem, i.e., to determine the layer-specific rates based on sensitivity, and its extra computational cost.

A1.3. Please refer to the global response GA2 above.

Q1.4. (Minor) The sensitivity of a set of layers may not be accurately represented by the direct summation of the sensitivities of individual layers in the set.

A1.4. We agree with the reviewer's opinion that the total sensitivities of the dropped network do not strictly equal the sum of each individual layer's sensitivity. However, given the practical constraints of calculating the sensitivities for the vast number of possible combinations of dropped networks in a deep neural network, we have made the assumption that the total sensitivities of a dropped network can be approximated by the sum of the sensitivities of its individual layers. In our future research, we will explore more accurate methods to calculate the network's sensitivity and determine optimal drop rates.

Q1.5. What is the long context modeling performance after applying DropBP?

A1.5.

• Table R.1.1. Processing Time Analysis for Calculating Sensitivities When Fine-Tuning LLMs on the Alpaca Dataset.

• Figure A in the attached PDF.

In response to the reviewer's request, we trained the LLaMA2-7B-chat model on the LongAlpaca dataset using sequence lengths of 16K and 32K. The settings followed those of LongLoRA [1]. As shown in Fig. A, our experiments demonstrated that the model successfully converged on loss with long sequence data of 16K or more. Additionally, when evaluating the fine-tuned model on a subset of the PG16 test set with a sequence length of 16K, we achieved lower perplexity (PPL) compared to the non-fine-tuned model as shown in Table R.1.1, confirming that our DropBP method enables effective long sequence modeling.

Due to limited GPU resources and review time, we conducted experiments on smaller training and test sets. Once we secure sufficient resources and time, we plan to obtain more robust experimental results for the revised version and integrate DropBP into LongLoRA for future work..

Q1.6 Could the authors present Figure 5 with # of steps as the x-axis to demonstrate faster convergence?

A1.6 Please refer to the global response GA3 above.

Q1.7 I wonder if the sensitivities would evolve, and the drop rate needs to be re-allocated through training.

A1.7

• Figure D in the attached PDF.

As shown in Fig. D, our experiments demonstrated that the sensitivity of each layer converges as training progresses. Accordingly, we calculates sensitivity just once at the 10% of the training process, minimizing overhead from sensitivity calculations. This approach proved effective in most experiments. However, we believe the reviewer's concern is also valid and we plan to add a feature that periodically recalculates sensitivity to adjust the drop rate to the DropBP library.

[1] Chen at al., "LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models"