BP-Modified Local Loss for Efficient Training of Deep Neural Networks

Dear Reviewer UZvZ,

Thank you for reviewing the paper. We appreciate the time and effort dedicated to reviewing this draft and value your insightful comments. Following is a point-by-point response to the reviewer’s comments.

1. motivation of local training algorithms is to train large datasets or large models in resource-constrained environments. So I think authors could scale to some datasets larger than CIFAR-100, such as the ImageNet, to better illustrate the applicability of the proposed method and meet the design motivation.

   Thank you for pointing this out. Due to the resources limitations, we have added additional experiments on the ImageNet-Tiny dataset with 10 epochs, which is usually enough for fine-tuning. The results are shown in Table 2 which is consistent with the results in CIFAR100 dataset. We found that the BP-modified local loss method can achieve better performance than the original local loss methods in most cases, especially when the original method cannot perform well. We observed that the LoKr tuning in original methods have comparable or better test accuracy on the ImageNet-Tiny dataset compared to the CIFAR100 dataset, while in BP tuning and other two model settings, the performance is much lower than the CIFAR100 dataset. We do not know what causes this phenomena and this maybe the reason why the BP-modified local loss method has lower performance than the original local loss methods in LoKr tuning case.

2. The optimization of the Loss function tends to significantly affect the training convergence performance. So I think it is necessary to provide model training time comparison results between the BP-Modified Local Loss method and other Local Loss baseline methods, such as total epochs needed for convergence.

   you for pointing this out, we have added Table 5 to show the time overhead of the BP-modified local loss method and the original methods. It shows that the BP-modified local loss method only has a slight additional time overhead compared to the original methods.

   The local loss method may not be able to converge to the same loss as the BP algorithm since the local gradient is biased. We report the average epoch that the BP-modified local loss method first achieved the same test accuracy as the final test accuracy of the original methods in CIFAR100 dataset in Table 8.

   We find that the BP-modified local loss method requires much less epochs to reach the same test accuracy as the final test accuracy of the original methods, demonstrates the efficiency of the proposed method.

3. In addition, I suggest adding some figure descriptions if possible so that the readers could more clearly understand the authors’ intention.

   Thank you for pointing this out. We have added Figure 1 on Page 3 which describes the idea of the BP-modified local loss method.

   As the figure demonstrates, to modify the local gradient in batch $B$ which usually has large bias with the true gradient, the BP-modified local loss method will use additional information of local gradient and true gradient in batch $B'$. It will compute the offset $\Delta \tilde g_{h}$ in batch $B'$ to adjust the local gradient in batch $B$. Finally, the modified local gradient will be used by the optimizers to update the model parameters.

   The figure suggests the idea of the BP-modify local loss method and the lazy update technique used in the method. The bias-variance trade-off is not presented since it is haed to combine the gradient and its distribution property in one figure.

Thank you for your reponse, which resolved all of my doubts. I am willing to raise my rating and recommend to accept this paper.

4. Although the authors adopted some strategies like lazy update and split mini-batch, computing the offset $\Delta \tilde g_{h}$ involves additional operations. I wonder what’s the overhead in terms of runtime and memory usage as the depth of the network increases? Can authors quantify the impact of these additional operations on the overall training efficiency? Such as total training time, per-epoch time as well as GPU memory usage as the depth of the network increases?

   The additional relative computation overhead of the BP-modified local loss method will be about $\frac{(1+r)B'}{BK}$ where $r$ is the ratio of computation overhead between BP algorithm and local loss algorithm on one sample. The memory overhead will be on the scale of $O(Bh,b'T)$ where $b'$ is the batch size after splitting, and the additional memory usage is the size of trainable parameters.

   Since we use the lazy update technique, the additional computation will be the acquisition of the local gradient $\tilde{g}_h$ and the true gradient $\tilde{g}_T$ in the batch $B'$ every $K$ steps. Therefore, the additional relative computation overhead is $\frac{(1+r)B'}{BK}$ where $r$ is the ratio of computation overhead between the BP algorithm and the local loss algorithm on one sample. The additional computation required for the scaling factors $\lambda_h$ and $a_h$ like $\|g_h\|_2,tr(Cov(g_h))$ mentioned in the previous response is of the size of the trainable parameters and is usually negligible.

   Furthermore, since we use the mini-batch splitting technique to split $B'$ into smaller mini-batches, it will not change the total computation but will slightly increase the computation time than no split. Then the additional time overhead will be slightly larger than $\frac{(1+r)B'}{BK}$. We have added Table 5 in Appendix D.1 to show the time overhead of the BP-modified local loss method and the original methods.

   As the model becomes deeper, the only parameter involved in computation that will change is $r$, the ratio of computation overhead between the BP algorithm and local loss algorithm on one sample. However, how $r$ changes with the depth of the network depends on the specific model and the design of the local loss method.

   On the other hand, our method needs additional memory to store the offset $\Delta \tilde g_{h}$ which is the size of trainable parameters and is usually negligible, especially for CNN and low-rank tuning cases such as LoRA. The scaling factors are some float or double numbers which are also negligible.

   As for the memory usage, the memory usage will increase linearly with batch size $B$ and it is reported that the memory usage will grow linearly with respect to horizon $h$, and these two factors are independent. When computing the offset $\Delta \tilde g_{h}$, the memory usage in getting local gradient $\tilde g_{h}$ and true gradient $\tilde{g}_{T}$ are $O(B'h),O(B'T)$ respectively. Since we use the mini-batch splitting technique to split $B'$ into smaller mini-batches with size $b'$, the memory usage will be $O(b'T)$. Further, acquiring different gradients is asynchronous so we only need to consider the maximum for different gradients, i.e. O(Bh,b'T). As long as $b'$ is small enough like $b'T\leq Bh$, the memory usage in computing gradients in batch $B'$ will be comparable to or less than computing local gradient in batch $B$.

   As the depth of the model increases, $T$ will become larger, however, as long as $b'T\leq Bh$, the memory usage in computing gradients in batch $B'$ will be comparable to or less than computing local gradient in batch $B$. In the extreme case, when $T\ll Bh$, the memory overhead $O(b'T)$ will be larger than $O(Bh)$ for the original method.

   For the numerical results, the ViT-b16 model and ResNet-50 model have different depths (12 and 17), and the results in Table 4 in Appendix D.1 show that the additional memory overhead is mainly the trainable parameters and negligible. There is also no significant additional computation overhead difference between the two models as Table 5 shows.

   Thank you for pointing this out. We have revised the last paragraph of Section 3 to indicate the additional computation and memory overhead and Section 5 to emphasize the additional computation and memory overhead of the BP-modified local loss method.

Dear Reviewer EzNV,

Thank you for reviewing the paper. We appreciate the time and effort dedicated to reviewing this draft and value your insightful comments. Following is a point-by-point response to the reviewer’s comments.

1. Generalisability is one of the major concerns, as the scope of the experiments is a bit limited. The results are primarily conducted on CIFAR-100, a relatively small dataset. While the improvements are evident, a more robust evaluation that includes larger datasets like ImageNet, as well as tasks other than image classification, would make the proposed method more convincing.

   Thank you for pointing this out. Due to the resources limitations, we have added additional experiments on the ImageNet-Tiny dataset with 10 epochs, which is usually enough for fine-tuning. The results are shown in Table 2 which is consistent with the results in CIFAR100 dataset. We found that the BP-modified local loss method can achieve better performance than the original local loss methods in most cases, especially when the original method cannot perform well. We observed that the LoKr tuning in original methods have comparable or better test accuracy on the ImageNet-Tiny dataset compared to the CIFAR100 dataset, while in BP tuning and other two model settings, the performance is much lower than the CIFAR100 dataset. We do not know what causes this phenomena and this maybe the reason why the BP-modified local loss method has lower performance than the original local loss methods in LoKr tuning case.

2. The authors leveraged the Ornstein-Uhlenbck process to model the bias-variance dynamic, however, the OU process is a simplified linear model, and deep neural network training is inherently non-linear and non-convex. The validity of the approximation discussed in the paper would require some empirical evidence to support, as well as some discussions in more depth. For example, deep neural networks exhibit different training dynamics across layers, with shallow and deep layers experiencing different gradient behaviours. Can authors provide how gradients evolve for different layers and test if the OU process holds similarly for both shallow and deep layers.

   Yes, the OU process is a much simplified linear model and there is no evidence that the training process of any neural network will be similar to an OU process.

   However, the OU process is just a toy example to introduce the bias-variance trade-off analysis. The result of the OU process gives us some insight into the bias-variance trade-off and the sub-optimality of the unbiased modification. These are also consistent with the one-step loss analysis in Section 4.2. The one-step loss analysis is the main analysis in our paper and it can apply to any deep neural network and task with Lipschitz continuous gradient and the variances of the local gradient and true gradient being bounded.

   Thank you for pointing this out. We have revised the introduction paragraph in Section 4.1. The different training dynamics across layers is a good point and we will consider this in our future work.

3. The paper mentions hyperparameters such as offset batch size B' and sampling period K, but does not discuss the selection of these parameters in detail. It is recommended that the authors provide more analysis on how these hyperparameters affect model performance and memory usage.

   For the influence of period $K$ and batch size $B'$ on the performance, please refer to the Global Response 2. As discussed in Appendix C, a larger period $K$, or equivalently, a lazier update of the $\Delta \tilde{g}_{h}$ will lead to larger variance since the noise of the same offset will accumulate, and enhance the delayed bias, thus reducing the performance. On the other hand, as Eq. (7) suggests, a smaller $B'$ will also increase the variance. However, a smaller $K$ and larger $B'$ mean more computation.

   We have added some numerical results on the sensitivity of the period $K$ and mini-batch size $B'$ and the results are shown in Appendix D.2. The results align with the previous discussion that larger $K$ and smaller $B'$ save some time at the expense of lower performance and vice versa.

   Additionally, if the mini-batch size $B'$ is too small, the variance of the modified gradient will be large and the performance might be worse, while it seems that it has more tolerance for larger $K$. If the period $K$ is too small or $B'$ is too large, the computation overhead will be large with diminishing performance improvement. Our experience suggests that $B'$ should be large enough to reduce the variance of the modified gradient while $K$ can be relatively large such that the computation overhead is small. $B'$ being $1\%$ of $BK$ are usually good choices.

   Note that there is no memory usage difference between different choices of $K$ and $B'$. Since we use the mini-batch splitting technique to split $B'$ into smaller mini-batches, the memory overhead of computing offset $\Delta \tilde g_h$ in batch $B'$ will be comparable to or less than computing local gradient in batch $B$. We do not change the smaller batch size so larger $B'$ means only more computation but not more memory usage. The additional memory usage of our method is to store the offset $\Delta \tilde{g}_{h}$ which is the size of trainable parameters, and scaling factors which are some double numbers. They are independent with the choice of $K$ and $B'$.

   The theoretical analysis of the influence of period $K$ and mini-batch size $B'$ is a challenging task and beyond the scope of this paper. We will consider this in our future work.

Dear Reviewer z11R,

Thank you for reviewing our paper. We greatly appreciate the time and effort you have dedicated and value your insightful comments. Following is a point-by-point response to the reviewer’s comments.

1. Impact of delayed gradient adjustment: The paper discusses the noise accumulation and bias adjustment problems that may be caused by delayed gradient adjustment, but does not provide a detailed analysis. It is recommended that the authors further study the impact of these factors on training dynamics and explore possible solutions.

   Thank you for pointing this out. The detailed analysis of the delayed gradient adjustment is much more difficult and complex since it involves multiple steps of training. The estimate in the one-step analysis may not apply to the multi-step analysis. This is beyond the scope of this paper. We will consider this in our future work.

   However, we have added some numerical results on the sensitivity of the period $K$ in Appendix D.2. The results show that larger $K$ will lead to lower performance but still can achieve better performance than the original local loss methods. This observation may indicate that the bias adjustment problem is not severe in the experiment period range and that noise accumulation is the majority issue. Especially the use of momentum in the optimizer will enlarge the noise accumulation because the same accumulated noise will be considered as the accurate direction thus harder to be averaged out in momentum.

2. More extensive experimental validation: Although the paper conducts experiments on multiple models and tasks, it is still recommended to conduct further experimental validation on larger models and datasets to demonstrate the universality and scalability of the method.

   Thank you for pointing this out. Due to the resources limitations, we have added additional experiments on the ImageNet-Tiny dataset with 10 epochs, which is usually enough for fine-tuning. The results are shown in Table 2 which is consistent with the results in CIFAR100 dataset. We found that the BP-modified local loss method can achieve better performance than the original local loss methods in most cases, especially when the original method cannot perform well. We observed that the LoKr tuning in original methods have comparable or better test accuracy on the ImageNet-Tiny dataset compared to the CIFAR100 dataset, while in BP tuning and other two model settings, the performance is much lower than the CIFAR100 dataset. We do not know what causes this phenomena and this maybe the reason why the BP-modified local loss method has lower performance than the original local loss methods in LoKr tuning case.

Dear Reviewer Zuvr,

Thank you for reviewing our paper. We appreciate the time and effort you have dedicated and value your insightful comments. Below is our point-by-point response to each of your comments:

1. Misc: Paraphrase lines 31, 32, and 33 for better readability.

   Thank you for pointing this out. We have revised the sentence for better readability:

   "This issue becomes more severe as the model becomes deeper and requires larger batch sizes, like SimCLR (Chen et al., 2020). Consequently, reducing the memory footprint of training and fine-tuning remains an active area of research (Krizhevsky et al., 2017; Rhu et al., 2016; Malladi et al., 2023)."

2. Is there any overhead in latency or throughput for BP-modified loss compared to other local loss-based training methods?

   Yes, the additional relative computation overhead of the BP-modified local loss method will be about $\frac{(1+r)B'}{BK}$ where $r$ is the ratio of computation overhead between the BP algorithm and the local loss algorithm on one sample. The memory overhead will be on the scale of $O(Bh,b'T)$ where $b'$ is the batch size after splitting, and the additional memory usage is the size of trainable parameters.

   The additional computation and memory requirement of the method is mainly on the offset $\Delta \tilde g_{h}$ on batch $B'$. Since we use the lazy update technique, the additional computation will be the acquisition of the local gradient $\tilde{g}_h$ and the true gradient $\tilde{g}_T$ in the batch $B'$ every $K$ steps. Therefore, the additional relative computation overhead is $\frac{(1+r)B'}{BK}$ where $r$ is the ratio of computation overhead between the BP algorithm and the local loss algorithm on one sample.

   However, since we use the mini-batch splitting technique to split $B'$ into smaller mini-batches, it will not change the total computation but will slightly further increase the computation time. Then the additional time overhead will be slightly larger than $\frac{(1+r)B'}{BK}$. We have added Table 5 in Appendix D.1 to show the time overhead of the BP-modified local loss method and the original methods.

   On the other hand, our method needs additional memory to store the offset $\Delta \tilde{g}_{h}$ which is the size of trainable parameters and is usually negligible, especially for CNN and low-rank tuning cases such as LoRA.

   As for the memory usage, the memory usage will increase linearly with batch size $B$ and it is reported that the memory usage will grow linearly with respect to horizon $h$, and these two factors are independent. When computing the offset $\Delta \tilde g_{h}$, the memory usage in getting local gradient $\tilde g_{h}$ and true gradient $\tilde{g}_{T}$ are $O(B'h),O(B'T)$ respectively. Since we use the mini-batch splitting technique to split $B'$ into smaller mini-batches with size $b'$, the memory usage will be $O(b'T)$. Further, acquiring different gradients is asynchronous so we only need to consider the maximum for different gradients, i.e. $O(Bh,b'T)$. As long as $b'$ is small enough like $b'T\leq Bh$, the memory usage in computing gradients in batch $B'$ will be comparable to or less than computing local gradient in batch $B$.

   Thank you for pointing this out. We have revised the last paragraph of Section 3 to indicate the additional computation and memory overhead and Section 5 to emphasize the additional computation and memory overhead of the BP-modified local loss method.

3. Is the bias-variance trade-off analysis generalizable across various model architectures and tasks, or are there conditions under which it may not apply?

   Yes, the bias-variance trade-off analysis is generalizable across various model architectures and tasks since it is based on the stochastic modified equation Eq. (6) (7) and the one-step loss analysis in Section 4.2.

   The Stochastic Modified Equation Eq. (6) was previously proposed by Li et al and Mandt et al and it is a general model that will be satisfied if the model has uniformly bounded gradient and Hessian which is satisfied by most of the modern models and tasks. Further, Eq. (7) is its natural extension in the BP-modified local loss method and it further requires the local gradient to hold the above properties.

   However, the OU process in Section 4.1 is a simplified linear model to introduce the bias-variance trade-off analysis. It is just a toy example to provide insight and capture the key phenomena of bias-variance balance and the sub-optimality of the unbiased modification. There is no evidence that the training process of any deep neural network will satisfy the OU process.

   Nonetheless, our main analysis is based on the one-step loss analysis in Section 4.2 where we just need the gradient of the loss function to be Lipschitz continuous, and the variance of both the true gradient and local gradient is bounded. This is a general condition that is satisfied by most of the modern models and tasks