Paper: Scaling Supervised Local Learning with Augmented Auxiliary Networks (ICLR-2024)

Reviewer comment (Experiments):
The author presents comparison results between AugLocal, DGL, and InfoPro. However, looking at Table 2 in the Appendix, AugLocal has higher FLOPs than the other methods. Typically, networks with a larger number of FLOPs tend to achieve higher accuracy. Therefore, it seems appropriate to compare the accuracy among AugLocal, DGL, and InfoPro with the same FLOPs, that is, auxiliary networks with equivalent FLOPs.

Author response:
Thank you for your suggestion. In response to your comment, we have performed additional experiments to scale up the auxiliary networks of DGL by using 3x3 convolutions with the same auxiliary network depth and a multiplier to scale up the channel numbers of the convolutional layers to ensure similar FLOPs as AugLocal. We further implement PredSim and InfoPro, which incorporate additional local losses and auxiliary networks, resulting in higher FLOPs than DGL. Our results in the table below consistently demonstrate that AugLocal outperforms the other methods with similar FLOPs. We have incorporated this experiment in Appendix A.5 of the updated paper.  

| Method | d=2 | d=3 | d=4 | d=5 | d=6 |
|:---:|:---:|:---:|:---:|:---:|:---:|
| AugLocal | 90.98±0.05 | 92.62±0.22 | 93.22±0.17 | 93.75±0.20 | 93.96±0.15 |
| DGL | 83.03±0.24 | 85.82±0.08 | 87.84±0.46 | 89.19±0.20 | 90.01±0.05 |
| PredSim | 72.06±0.63 | 80.16±0.47 | 86.00±0.56 | 88.27±0.72 | 88.34±0.34 |
| InfoPro | 83.71±0.20 | 89.14±0.17 | 90.75±0.22 | 91.45±0.05 | 92.10±0.17 |

Reviewer comment (Experiments):
AugLocal utilizes the subsequent layers of the hidden layers. As a result, this method has to wait for the subsequent layers to be trained for training the current layer with the next data (update locking). In contrast, traditional local learning methods like DGL and InfoPro completely solve this update locking issue, enabling asynchronous training.

Author response:
We would like to clarify that the construction of augmented auxiliary networks happens before the training starts, and the selected auxiliary layers only share the same structures with their counterparts in the primary network. This particular design ensures hidden layers in the primary network can be updated independently and asynchronously, thereby solving the update locking problem. However, as already mentioned in the paper, the update locking problem still exists in the training of the hidden layers along with their auxiliary networks and will become exacerbated when the depth of auxiliary networks grows. Therefore, we constrain the depth of auxiliary networks to control the degree of update locking. To alleviate the update locking problem in auxiliary networks, one potential solution could involve gradually freezing some auxiliary layers during training. This strategy will be investigated as part of our future work. Based on your suggestions, we have updated to mention this point in the paper's conclusion. Thanks.

Reviewer comment (Experiments):
It might be challenging to compare AugLocal at the same FLOPs since the method of determining the auxiliary network in DGL and InfoPro is already fixed. Nevertheless, if there is a comparison result with AugLocal at the same FLOPs, it would enhance the quality of the paper if it is possible.

Author response:
Thank you for your suggestion. In response to your comment, we have performed additional experiments to scale up the auxiliary networks of DGL by using 3x3 convolutions with the same auxiliary network depth and a multiplier to scale up the channel numbers of the convolutional layers to ensure similar FLOPs as AugLocal. We further implement PredSim and InfoPro, which incorporate additional local losses and auxiliary networks, resulting in higher FLOPs than DGL. Our results in the table below consistently demonstrate that AugLocal outperforms the other methods with similar FLOPs. We have incorporated this experiment in Appendix A.5 of the updated paper.  

| Method | d=2 | d=3 | d=4 | d=5 | d=6 |
|:---:|:---:|:---:|:---:|:---:|:---:|
| AugLocal | 90.98±0.05 | 92.62±0.22 | 93.22±0.17 | 93.75±0.20 | 93.96±0.15 |
| DGL | 83.03±0.24 | 85.82±0.08 | 87.84±0.46 | 89.19±0.20 | 90.01±0.05 |
| PredSim | 72.06±0.63 | 80.16±0.47 | 86.00±0.56 | 88.27±0.72 | 88.34±0.34 |
| InfoPro | 83.71±0.20 | 89.14±0.17 | 90.75±0.22 | 91.45±0.05 | 92.10±0.17 |

Reviewer comment (Presentation):
Due to the nature of the algorithm, the current layer cannot be trained with the next data until the subsequent layer of the hidden layer is trained (update locking). It would be good to mention this limitation in the paper's conclusion or Appendix and suggest directions for future work. If possible, briefly mentioning a possible solution would also be beneficial.

Author response:
We would like to clarify that the construction of augmented auxiliary networks happens before the training starts, and the selected auxiliary layers only share the same structures with their counterparts in the primary network. This particular design ensures hidden layers in the primary network can be updated independently and asynchronously, thereby solving the update locking problem. However, as already mentioned in the paper, the update locking problem still exists in the training of the hidden layers along with their auxiliary networks and will become exacerbated when the depth of auxiliary networks grows. Therefore, we constrain the depth of auxiliary networks to control the degree of update locking. To alleviate the update locking problem in auxiliary networks, one potential solution could involve gradually freezing some auxiliary layers during training. This strategy will be investigated as part of our future work. Based on your suggestions, we have updated to mention this point in the paper's conclusion. Thanks.

Reviewer comment (Experiments):
It would be even better if results for simple models like GCN2 or BERT-Base could be added.

Author response:
Thank you for raising the potential applicability of our proposed method on graph neural networks and the transformer architecture model. We believe they are interesting topics to explore independently as the local learning research on these architectures is not yet available. It would require more time to build the baselines for a comprehensive comparison and, therefore, we would like to leave it as future works.

Reviewer comment (Reproducibility):
It is not clear when the authors say: "when more layers of the primary network are selected into the auxiliary networks". Does that mean the weights are shared between layers in the primary network and the auxiliary network? Or just the initialization is shared, or just the architecture is shared?

Author response:
**Q1:** It is not clear when the authors say: "when more layers of the primary network are selected into the auxiliary networks". Does that mean the weights are shared between layers in the primary network and the auxiliary network? Or just the initialization is shared, or just the architecture is shared?

**A1:** Thank you for bringing this up. We would like to clarify that only the architecture is shared in our method. To improve the clarity on this point, we have revised Section 3.2 of our paper, emphasizing that "the selection of the auxiliary network takes place prior to training, and only the architectural structure of the selected layers are employed in the construction of the auxiliary network.".

Reviewer comment (Experiments):
Based on Point 2, the guess is that other layers are not shared based on the last paragraph of Sec 3.2 and the last sentence in A.2 in the supplementary. In this case, it is not clear why Unif., Seq., Repe. can show a large difference even for d = 2 in Table 5. Does that mean downsampling in the auxiliary network is critical? If the authors want to claim the importance of architectural bias, the missing results are C1×1/C3x3 with downsampling, which should not be considered architectural bias. The authors should also show that C3x3 performs equally well with the VGG network. Additionally, the authors should ablate to what circumstances, when altering the auxiliary architecture presented in Table 1 in the supplementary, the AugLocal performance degraded and performed similarly to other baselines. Simply attributing the performance gain to architecture bias is not convincing as the learning is free-form if the auxiliary network weights are not shared with the primary network.

Author response:
**Q2.1:**	It is not clear why Unif., Seq., Repe. can show a large difference even for d = 2 in Table 5. 

**A2.1:** The large accuracy difference for d=2 can be attributed to the different channel numbers used in their residual blocks. The sequential (Seq.) and repetitive (Repe.) selection strategies will select layers that are closer to the primary network layer that we intend to train, which have smaller channel numbers compared to those chosen by the uniform (Unif.) sampling method. This leads to a noticeable decrease in accuracy. We have revised  Section 4.5 of our paper to provide a detailed explanation on this point. 

**Q2.2:**	If the authors want to claim the importance of architectural bias, I believe the missing results are C1×1/C3x3 with downsampling, which should not be considered architectural bias from my perspective. The authors should also show C3x3 performs equally well with the VGG network. 

**A2.2:** Our experiments involving C1×1/C3×3 are specifically designed to demonstrate the advantage of our proposed AugLocal strategy compared to other handcrafted ways to construct auxiliary networks. Following your suggestion, we conduct an experiment of C3×3 with downsampling at comparable FLOPs to AugLocal. Our experimental results in the table below demonstrate improved accuracies compared to those models without downsampling. Furthermore, it is evident that AugLocal consistently outperforms C3×3 with downsampling, suggesting the effectiveness of the proposed auxiliary network construction strategy in AugLocal (i.e., perform uniform sampling from the subsequent layers in the primary network). 

For the VGG networks used in this paper, our augmented auxiliary networks were constructed exactly following the way you have suggested (i.e., uniformly sample the C3×3 layers from the subsequent layers in the primary network). This approach delivers superior performance for VGG networks, as illustrated by the results presented in Table 2 of our paper.  

| |d=2|d=3|d=4|d=5|d=6|
|:-:|:-:|:-:|:-:|:-:|:-:|
|AugLocal|90.98±0.05|92.62±0.22|93.22±0.17|93.75±0.20|93.96±0.15|
|C3×3 **without** downsampling|83.03±0.24|85.82±0.08|87.84±0.46|89.19±0.20|90.01±0.05|
|C3×3 **with** downsampling|85.18±0.26 |89.90±0.13|91.36±0.23|91.72±0.11|91.97±0.18|

Reviewer comment (Evaluation):
Table 4 shows unusual results; it is unclear how the authors obtained the 157.12 GB GPU memory result as there is no single GPU that can handle it. The ResNet-110 result looks like an outlier. The authors should elaborate on why it shows significantly better memory saving while the architecture is similar to other cases (ResNet-32/34/101).

Author response:
**Q3:** Table 4 shows weird results; how did the authors get the 157.12 GB GPU memory results as there is no single GPU that can handle it AFAIK? The ResNet-110 result looks like an outlier. Can the authors elaborate on why it shows significantly better memory saving while the architecture is similar to other cases (ResNet-32/34/101)?

**A3:** Thank you for the feedback. First of all, we want to clarify that the reported GPU memory usage represents the aggregate utilization across all GPUs within the experimental machine, rather than the capacity of a single GPU card. This has been explicitly detailed in Appendix A.2 of our paper. Specifically, our experiments were conducted on a system equipped with 10 $\times$ NVIDIA RTX 3090 GPU cards with 24GB GPU memory each, meeting the memory requirements for our evaluations. To avoid any confusion, we have revised Section 4.4 of our paper to explicitly state this point for clarity.

Regarding the observed memory efficiency of ResNet-110, it is essential to consider the underlying mechanics. In AugLocal, the memory allocated to a local layer (defined as a hidden layer along with its associated auxiliary network) is freed upon the completion of the local layer's update. As a result, the peak GPU memory consumption is determined by the local layer with the largest memory footprint, relative to other local layers. Therefore, the memory-saving characteristic of AugLocal stems from the ratio between the maximum memory utilization among local layers and the overall memory footprint of the entire network. As compared with other network architectures, ResNet-110 has a much smaller memory footprint of its local layers when contrasted with the total memory consumption of the network. Therefore, ResNet-110 shows substantial GPU memory saving in comparison to other network architectures.

Reviewer comment (Writing):
It would be helpful to show FLOPs comparison between local learning and end-to-end learning instead of just showing FLOPs reduction in Figure 5. Also, gamma is introduced in equation 2 but never used in the following text, making it confusing how it guides the rule design.

Author response:
**Q4:** It would be great to show FLOPs comparison between local learning and end-to-end learning instead of just showing FLOPs reduction in Figure 5. Also, gamma is introduced in eq2 but never used in the following text, making it confusing about how it guides the rule design.

**A4:** Thanks for your suggestions. We have updated our paper to provide explicit values of FLOPs for local learning and end-to-end learning in Appendix Table 6. In addition, the value of gamma is set to 0.5 in our experiments and we analyzed the influence of gamma in Figure 5. We have updated our paper to make this point clearer.

Reviewer comment (Evaluation):
In Table 1, the authors can provide a more complete picture by including some complexity measure (similar to what is reported in Table 2 in Appendix). What is the wall-clock time of all of these methods? I suspect that AugLocal will incur some significant overhead due to having deeper auxiliary networks.

Author response:
**A1:** Thank you for your suggestion. As you have rightly pointed out, the wall-clock time for the proposed AugLocal depends on the specific implementation. In the case of the sequential implementation, where each hidden layer is trained sequentially one after another, AugLocal does incur increased computational overhead compared to other local learning methods. This is due to the presence of additional auxiliary networks. To provide a more complete picture, we have compared the wall-clock time across all local learning methods under the sequential implementation using ResNet-110, as shown in the table below. We have also included the results in Table 2 in Appendix. These results reveal that AugLocal (d=2) can achieve higher accuracy at the trade-off of slightly increased computational overhead in comparison to other local learning rules, and the accuracy of AugLocal can be further enhanced by leveraging deeper auxiliary nets, albeit at the expense of increased wall-clock time. Nevertheless, we would like to highlight that when having a dedicated parallel implementation, which we are still working on, the AugLocal can yield significant training speed-ups compared to end-to-end BP training, as theoretically analyzed in Appendix A.1. 

|  |BP|Gradient Checkpoint|PredSim|DGL|InfoPro|AugLocal(d=2)|AugLocal(d=3)|AugLocal(d=4)|AugLocal(d=5)|AugLocal(d=6)|
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|GPU Memory (GB)|9.27|3.03|1.54|1.61|3.98|1.71|1.62|1.70|1.71|1.72|
|Computational Overhead (Wall-clock Time)|-|34.1%|65.9%|76.2%|292.3%|87.2%|115.9%|135.6%|180.8%|214.6%|
|Acc.|94.61±0.18|94.61±0.18|74.95±0.36|85.69±0.32|86.95±0.46|90.98±0.05|92.62±0.22|93.22±0.17|93.75±0.20|93.96±0.15|

Reviewer comment (Theory):
The theoretical speedup of (d+1)/(L+1) is not very informative, as it assumes an ideal setting where all local modules can be run in parallel, which is not trivial to implement. The point of local learning is save GPU memory, which I suspect comes at the expense of extra time. It is also worth comparing AugLocal with back-prop with gradient checkpointing.

Author response:
**A2:** We fully agree with you on this point and thank you for suggesting adding the comparison with gradient checkpointing.  As already replied in **A1**, we have updated our paper to include the wall-clock time of our method under the sequential implementation in Appendix Table 2. Moreover, to fully realize the potential of AugLocal, we are working on a GPU-based parallel implementation and we hope to provide an update in future work. 

Additionally, in response to the suggestion to compare AugLocal with gradient checkpointing, we have revised our paper to include this comparison in Table 4 and Appendix Table 2. Our results demonstrate that AugLocal can achieve a much lower GPU memory footprint, albeit accompanied by a moderate increase in wall-clock time.

Reviewer comment (Writing):
eq(2) is a bit misleading, it implies that only the auxiliary network parameters $\Phi^l$ impact the local losses, and the final loss is a function of all network parameters.

Author response:
**A5:** Thank you for pointing this out. We have fixed it in our revised paper and please refer to it for the update.

Reviewer comment (Experiments):
In the ImageNet experiments, the DGL numbers correspond to a much waker end-to-end baseline (66.6 vs 71.59 in this paper), thus making DGL seem weaker than it actually is. The authors already reproduced the DGL numbers (as well as other methods) for the other datasets to ensure a fair comparison. Why haven't the authors done the same for ImageNet?

Author response:
**A3:** Thank you for your suggestions. We have conducted the DGL experiment on ImageNet using consistent training setups. The accuracy achieved by DGL is 67.32%, which is much better than the original paper's reported accuracy of 64.4%. Our AugLocal still outperforms DGL by a large margin of 3.6% (70.92% vs. 67.32%). We have updated our paper to reflect the enhanced accuracy of DGL.

Reviewer comment (Experiments):
What would the GPU memory savings look like with varying batch sizes, does it require a large BS of 1024 to show significant improvement?

Author response:
**A4:** Thank you for pointing this out. We have conducted experiments to analyze the GPU memory footprint of AugLocal with varying batch sizes (ranging from 128 to 2048) using ResNet-110 on CIFAR-10. As the results are given below, even with a small batch size of 128, AugLocal can achieve significant memory savings. Moreover, the memory savings further improve as the batch size increases. 

|Batch Size|128|256|512|1024|2048|
|:-:|:-:|:-:|:-:|:-:|:-:|
|GPU Memory of BP (GB)|1.18|2.34|4.65|9.27|18.51|
|GPU Memory of AugLocal (GB)|0.38 |0.54|0.93|1.72|3.36|
|GPU Memory Saving by AugLocal|67.8%|76.9%|80.0%|81.4%|81.8%|