Module-wise Training of Neural Networks via the Minimizing Movement Scheme

Summary

This paper proposes a new training method for greedy layer-wise or module-wise training of neural networks, which is compelling in constrained and on-device settings where memory is limited and suffers from a stagnation problem. Experimental results show that their method improves the accuracy of module-wise training of various architectures. -------After Rebuttal------- I thank the authors for the answers to my comments and questions. Most of my concerns are properly addressed. Therefore I raise my score. While I still think the word "early overfitting" in the paper should be carefully used. I agree with the phenomenon that vanilla module-wise training performs very well in the early layers but stagnates and gets overtaken later. While it can be due to the insufficient mutual information between the learned features and the inputs, instead of the early layer overfitting. Generally, I think the early modules containing fewer learnable parameters are difficult to overfit the large-scale dataset, such as ImageNet.

Strengths

1. Overall, the idea is new and the proposed method is theoretically sound. 2. The paper is well-written and easy to follow.

Weaknesses

On Line 34, the authors claim that the early modules can overfit and learn more discriminative features than end-to-end training, destroying task-relevant information, and deeper modules don’t improve the test accuracy significantly, or even degrade it. However, the experimental results in this paper are not sufficient enough to support this assumption, especially for the overfitting problem. As the listed experiments are all conducted on small or medium-scale datasets, the overfitting problem might happen for the early layers. While, for the large-scale dataset, such as ImageNet, the early layers are more likely to be under-fitting due to the insufficient model capacity of the first few layers.

Questions

1. The wall-time training cost is also an important evaluation metric for training algorithms. I suggest the authors provide a comparison of the training time between their proposed TRGL and other related methods.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See weaknesses and questions.

Summary

The paper explores the module-wise training of neural networks via the Minimizing Movement Scheme. This approach aims to overcome the stagnation problem often encountered in layer-wise training, leading to improved accuracy and reduced memory usage. The authors compare their results with those of other methods and demonstrate the effectiveness of their approach on various architectures, including ResNets, Transformers, and VGG.

Strengths

1. The introduced methodology effectively addresses the frequently observed stagnation issue associated with layer-wise training. As a result, there is a noteworthy enhancement in model accuracy along with a commendable reduction in memory usage. 2. Establishing a connection between neural network training and optimal transport theory is not only intellectually stimulating, but it also presents intriguing theoretical insights. 3. Generally, the manuscript is well-structured and clearly articulated. However, some convoluted notations used in Section 2.2 could potentially be simplified for improved reader comprehension.

Weaknesses

1. The paper is currently missing a detailed convergence analysis. It's vital to elucidate whether the overall training process will indeed converge in order to reinforce the reliability and robustness of the proposed approach. 2. The study could benefit from a more comprehensive ablation analysis, specifically, an exploration on the impact of variations in the regularization penalty. This would further validate the robustness of the model and provide additional insights into its behavior under different conditions.

Questions

N/A

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

The paper proposed a new regularization for module-wise training via the distance of the input and output of the module.

Strengths

1. The proposed regularization is quite straightforward and easy to apply. 2. The paper connects the proposed regularization with optimal transport via theoretical analysis. 3. Some experimental results are good, like adding more modules on STL-10.

Weaknesses

1. Although the author shows the connection between the proposed regularization and optimal transport, it does not provide insights regarding convergence. Whether the proposed regularization helps the convergence of the multi-module training might be more interesting for larger audiences. 2. Authors argue that 'residual connections are already biased towards small displacements.' However, residual connections are not toward small values for early layers, as shown in [28]. As a result, how does the proposed method control each module's regularization strength? If they are set to the same value, then the regularization may be harmful to early layers. In addition, the selection of the regularization strength becomes more complex when increasing the number of modules K. 3. The improvements over the previous method are only obvious on STL-10, and STL-10 is not a large dataset. In addition, larger datasets and models are omitted. I think memory saving will be more meaningful for larger datasets and models. As a result, the scalability of the proposed method is not examined.

Questions

1. Why is the memory saving of the proposed method often less than other methods? Where this overhead comes from? 2. Some notations are not introduced properly. For example, the `#' notation, and what is the difference of $T_k$ and $T_k^{\tau}$.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Limitations are provided.

Summary

TRGL offers a promising approach to module-wise training that addresses the stagnation problem, improves accuracy, and reduces memory usage in constrained and on-device settings. The method introduces a module-wise regularization inspired by the minimizing movement scheme for gradient flows in distribution space. By minimizing the kinetic energy of modules along with the training loss, TRGL encourages modules to change their inputs as little as possible, preserving task-relevant information.

Strengths

(1) TRGL significantly reduces memory usage compared to end-to-end training, ranging from 10% to 60% less memory. This makes it particularly suitable for constrained settings and on-device training where memory resources are limited. (2) The authors provide theoretical analysis, proving that TRGL leads to more regular and stable greedy modules that progressively minimize the loss. (3) TRGL can be applied to various network architectures, especially those using residual connections such as ResNets and vision transformers.

Weaknesses

(1) TRGL introduces an additional regularization term to the training objective, which increases the computational complexity of the training process. The calculation of the kinetic energy and its incorporation into the loss function may require additional computational resources, leading to longer training times. How about the actual training time compared to other methods? (2) what is VanGL? is it just a module-wise training without the regularization terms? (3) In my opinion, the performance gain is marginal compared to vanilla VanGL. (4) The principle of finding Tau is vague.

Questions

See weakness

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

This increased complexity may result in longer training times and higher computational requirements, which could be a limitation in resource-constrained environments or when training large-scale models.

Dear reviewers

Dear reviewers, thank you once again for your valuable comments and insightful recommendations. As the deadline for the discussion approaches, we would like to gently remind you of our eagerness to receive your response and engage in further discourse regarding the points raised in the rebuttal. We would like to draw your attention to the fact that one of the primary weaknesses identified was the lack of comprehensive large-scale evaluation. In addressing this concern, we have conducted additional experiments and are pleased to present new quantitative results obtained from large datasets (imageNet). These results serve to reinforce and substantiate the conclusions we initially put forth. Thanks again, Authors