Question:
You are a highly experienced, conscientious, and fair academic reviewer, please help me review this paper. The review should be organized into nine sections:
1. Summary: A summary of the paper in 100-150 words.
2. Strengths/Weaknesses/Questions: The Strengths/Weaknesses/Questions of paper, which should be listed in bullet points, with each point supported by specific examples from the article where possible.
3. Soundness/Contribution/Presentation: Rate the paper's Soundness/Contribution/Presentation, and match this score to the corresponding description from the list below and provide the result. The possible scores and their descriptions are:
    1 poor
    2 fair
    3 good
    4 excellent
4. Rating: Give this paper an appropriate rating, match this rating to the corresponding description from the list below and provide the result. The possible Ratings and their descriptions are:
    1 strong reject
    2 reject, significant issues present
    3 reject, not good enough
    4 possibly reject, but has redeeming facets
    5 marginally below the acceptance threshold
    6 marginally above the acceptance threshold
    7 accept, but needs minor improvements
    8 accept, good paper
    9 strong accept, excellent work
    10 strong accept, should be highlighted at the conference
5. Paper Decision: It must include the Decision itself(Accept or Reject) and the reasons for this decision, based on the criteria of originality, methodological soundness, significance of results, and clarity and logic of presentation.

Here is the template for a review format, you must follow this format to output your review result:
**Summary:**
Summary content

**Strengths:**
- Strength 1
- Strength 2
- ...

**Weaknesses:**
- Weakness 1
- Weakness 2
- ...

**Questions:**
- Question 1
- Question 2
- ...

**Soundness:**
Soundness result

**Presentation:**
Presentation result

**Contribution:**
Contribution result

**Rating:**
Rating result

**Paper Decision:**
- Decision: Accept/Reject
- Reasons: reasons content

Please ensure your feedback is objective and constructive. The paper is as follows:

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

Skander Karkar

Criteo, Sorbonne Universite

&Ibrahim Ayed

Sorbonne Universite, Thales

Emmanuel de Bezenac

ETH Zurich

&Patrick Gallinari

Criteo, Sorbonne Universite

###### Abstract

Greedy layer-wise or module-wise training of neural networks is compelling in constrained and on-device settings where memory is limited, as it circumvents a number of problems of end-to-end back-propagation. However, it suffers from a stagnation problem, whereby early layers overfit and deeper layers stop increasing the test accuracy after a certain depth. We propose to solve this issue by introducing a module-wise regularization inspired by the minimizing movement scheme for gradient flows in distribution space. We call the method TRGL for Transport Regularized Greedy Learning and study it theoretically, proving that it leads to greedy modules that are regular and that progressively solve the task. Experimentally, we show improved accuracy of module-wise training of various architectures such as ResNets, Transformers and VGG, when our regularization is added, superior to that of other module-wise training methods and often to end-to-end training, with as much as \(60\%\) less memory usage.

## 1 Introduction

End-to-end backpropagation is the standard training method of neural networks. However, it requires storing the whole model and computational graph during training, which requires large memory consumption. It also prohibits training the layers in parallel. Dividing the network into modules, a module being made up of one or more layers, accompanied by auxiliary classifiers, and greedily solving module-wise optimization problems sequentially (i.e. one after the other fully) or in parallel (i.e. at the same time batch-wise), consumes much less memory than end-to-end training as it does not need to store as many activations, and when done sequentially, only requires loading and training one module (so possibly one layer) at a time. Module-wise training has therefore been used in constrained settings in which end-to-end training can be impossible such as training on mobile devices [58, 57] and dealing with very large whole slide images [65]. When combined with batch buffers, parallel module-wise training also allows for parallel training of the modules [8]. Despite its simplicity, module-wise training has been recently shown to scale well [8, 47, 60, 45], outperforming more complicated alternatives to end-to-end training such as synthetic [33, 14] and delayed [32, 31] gradients, while having superior memory savings.

In a classification task, module-wise training splits the network into successive modules, a module being made up of one or more layers. Each module takes as input the output of the previous module, and each module has an auxiliary classifier so that a local loss can be computed, with backpropagation happening only inside the modules and not between them (see Figure 1 below).

The main drawback of module-wise training is the well-documented _stagnation problem_ observed in [43, 7, 60, 47], whereby early modules overfit and learn more discriminative features than end-to-endtraining, destroying task-relevant information, and deeper modules don't improve the test accuracy significantly, or even degrade it, which limits the deployment of module-wise training. We further highlight this phenomenon in Figures 2 and 3 in Section 4.4. To tackle this issue, InfoPro [60] propose to maximize the mutual information that each module keeps with the input, in addition to minimizing the loss. [7] make the auxiliary classifier deeper and Sedona [47] make the first module deeper. These last two methods lack a theoretical grounding, while InfoPro requires a second auxiliary network for each module besides the classifier. We propose a different perspective, leveraging the analogy between residual connections and the Euler scheme for ODEs [61]. To preserve input information, we minimize the kinetic energy of the modules along with the training loss. Intuitively, this forces the modules to change their input as little as possible. We leverage connections with the theories of gradient flows in distribution space and optimal transport to analyze our method theoretically.

Our approach is particularly well-adapted to networks that use residual connections such as ResNets [27; 28], their variants (e.g. ResNeXt [62], Wide ResNet [63], EfficientNet [56] and MobileNetV2 [48]) and vision transformers that are made up essentially of residual connections [39; 17], but is immediately usable on any network where many layers have the same input and output dimension such as VGG [52]. Our contributions are the following:

* We propose a new method for module-wise training. Being a regularization, it is lighter than many recent state-of-the-art methods (PredSim [45], InfoPro [60]) that train another auxiliary network besides the auxiliary classifier for each module.
* We theoretically justify our method, proving that it is a transport regularization that forces the module to be an optimal transport map making it more regular and stable. We also show that it amounts to a discretization of the gradient flow of the loss in probability space, which means that the modules progressively minimize the loss and explains why the method avoids the accuracy collapse observed in module-wise training.
* Experimentally, we consistently improve the test accuracy of module-wise trained networks (ResNets, VGG and Swin-Transformer) beating 8 other methods, in sequential and parallel module-wise training, and also in _multi-lap sequential_ training, a variant of sequential module-wise training that we introduce and that performs better in many cases. In particular, our regularization makes parallel module-wise training superior or comparable in accuracy to end-to-end training, while consuming \(10\%\) to \(60\%\) less memory.

## 2 Transport-regularized module-wise training

The typical setting of (sequential) module-wise training for minimizing a loss \(L\), is, given a dataset \(\mathcal{D}\), to solve one after the other, for \(1{\leq}k{\leq}K\), Problems

\[(T_{k},F_{k})\in\arg\min_{T,F}\sum_{x\in\mathcal{D}}L(F,T\circ G_{k-1}(x)) \tag{1}\]

Figure 1: Module-wise training.

where \(G_{k}=T_{k}\circ\ldots\circ T_{1}\) for \(1{\leq}k{\leq}K\), \(G_{0}{=}\mathsf{id}\), \(T_{k}\) is the module (one or many layers) and \(F_{k}\) is an auxiliary classifier. Module \(T_{k}\) receives the output of module \(T_{k-1}\), and auxiliary classifier \(F_{k}\) computes the prediction from the output of \(T_{k}\) so the loss can be computed. The inputs are \(x\) and \(L\) has access to their labels \(y\) to calculate the loss. i.e. \(L(F,T\circ G_{k-1}(x))=l(F\circ T\circ G_{k-1}(x),y)\) where \(l\) is a machine learning loss such as cross-entropy. See Figure 1. The final network trained this way is \(F_{K}\circ G_{K}\). But, at inference, we can stop at any depth \(k\) and use \(F_{k}\circ G_{k}\) if it performs better. Indeed, an intermediate module often performs as well or better than the last module because of the early overfitting and subsequent stagnation or collapse problem of module-wise training [43; 7; 60; 47].

We propose below in (2) a regularization that avoids the destruction of task-relevant information by the early modules by forcing them to minimally modify their input. Proposition 2.2 proves that by using our regularization (2), we are indeed making the modules build upon each other to solve the task, which is the property we desire in module-wise training, as the modules now act as successive proximal optimization steps in the _minimizing movement scheme_ optimization algorithm for maximizing the separability of the data representation. The background on optimal transport (OT), gradient flows and the minimizing movement scheme is in Appendices A and B.

### Method statement

To keep greedily-trained modules from overfitting and destroying information needed later, we penalize their kinetic energy to force them to preserve the geometry of the problem as much as possible. If each module is a single residual block (that is a function \(T{=}\mathsf{id}{+}r\), which includes many transformer architectures [39; 17]), its kinetic energy is simply the squared norm of its residue \(r{=}T{-}\mathsf{id}\), which we add to the loss \(L\) in the target of the greedy problems (1). All layers that have the same input and output dimension can be rewritten as residual blocks and the analysis applies to a large variety of architectures such as VGG [52]. Given \(\tau{>}0\), we now solve, for \(1{\leq}k{\leq}K\), Problems

\[(T^{\tau}_{k},F^{\tau}_{k})\in\arg\min_{T,F}\sum_{x\in\mathcal{D}}L(F,T\circ G^ {\tau}_{k-1}(x))+\frac{1}{2\tau}\|T\circ G^{\tau}_{k-1}(x){-}G^{\tau}_{k-1}(x) \|^{2} \tag{2}\]

where \(G^{\tau}_{k}{=}T^{\tau}_{k}{\circ}\ldots\circ T^{\tau}_{1}\) for \(1{\leq}k{\leq}K\) and \(G^{\tau}_{0}{=}\mathsf{id}\). The final network is \(F^{\tau}_{K}{\circ}G^{\tau}_{K}\). Intuitively, this biases the modules towards moving the points as little as possible, thus at least keeping the performance of the previous module. Residual connections are already biased towards small displacements and this bias is desirable and should be encouraged [35; 64; 26; 15; 36]. But the method can be applied to any module where \(T(x)\) and \(x\) have the same dimension so that \(T(x){-}x\) can be computed.

To facilitate the theoretical analysis, we rewrite the method in a more general formulation using data distribution \(\rho\), which can be discrete or continuous, and the distribution-wide loss \(\mathcal{L}\) that arises from the point-wise loss \(L\). Then Problem (2) is equivalent to Problem

\[(T^{\tau}_{k},F^{\tau}_{k})\in\arg\min_{T,F}\ \mathcal{L}(F,T_{\sharp}\rho^{ \tau}_{k})+\frac{1}{2\tau}\int_{\Omega}\|T(x)-x\|^{2}\,\mathrm{d}\rho^{\tau}_ {k}(x) \tag{3}\]

with \(\rho^{\tau}_{k+1}{=}(T^{\tau}_{k})_{\sharp}\rho^{\tau}_{k}\), \(\rho^{\tau}_{1}{=}\rho\) and \(\mathcal{L}(F,T_{\sharp}\rho^{\tau}_{k})=\int L(F,T(x))\,\mathrm{d}\rho^{\tau} _{k}(x)=\int L(F,z)\,\mathrm{d}T_{\sharp}\rho^{\tau}_{k}(x)\).

### Link with the minimizing movement scheme

We now formulate our main result: solving Problems (3) is equivalent to following a _minimizing movement scheme (MMS)_[50] in distribution space for minimizing \(\mathcal{Z}(\mu):=\min_{F}\mathcal{L}(F,\mu)\), which is the loss of the best classifier. If we are limited to linear classifiers, \(\mathcal{Z}(\rho^{\tau}_{k})\) is the linear separability of the representation \(\rho^{\tau}_{k}\) at module \(k\) of the data distribution \(\rho\). The MMS, introduced in [24; 23], is a metric counterpart to Euclidean gradient descent for minimizing functionals over distributions. In our case, \(\mathcal{Z}\) is the functional we want to minimize. We define the MMS below in Definition 2.1

The distribution space we work in is the metric Wasserstein space \(\mathbb{W}_{2}(\Omega)=(\mathcal{P}(\Omega),W_{2})\), where \(\Omega\subset\mathbb{R}^{d}\) is a convex compact set, \(\mathcal{P}(\Omega)\) is the set of probability distributions over \(\Omega\) and \(W_{2}\) is the Wasserstein distance over \(\mathcal{P}(\Omega)\) derived from the optimal transport problem with Euclidean cost:

\[W_{2}^{2}(\alpha,\beta)=\min_{T\,\mathrm{s.t.}\ T_{\sharp}\alpha=\beta}\int_{ \Omega}\|T(x)-x\|^{2}\,\mathrm{d}\alpha(x) \tag{4}\]

where we assume that \(\partial\Omega\) is negligible and that the distributions are absolutely continous.

**Definition 2.1**.: Given \(\mathcal{Z}:\mathbb{W}_{2}(\Omega)\rightarrow\mathbb{R}\), and starting from \(\rho_{1}^{\tau}\in\mathcal{P}(\Omega)\), the Minimizing Movement Scheme (MMS) takes proximal steps for minimizing \(\mathcal{Z}\). It is s given by

\[\rho_{k+1}^{\tau}\in\arg\min_{\rho\in\mathcal{P}(\Omega)}\ \mathcal{Z}(\rho)+ \frac{1}{2\tau}W_{2}^{2}(\rho,\rho_{k}^{\tau}) \tag{5}\]

The MMS (5) can be seen as a non-Euclidean implicit Euler step for following the gradient flow of \(\mathcal{Z}\), and \(\rho_{k}^{\tau}\) converges to a minimizer of \(\mathcal{Z}\) under some conditions (see the end of this section).

So under the mentioned assumptions on \(\Omega\) and absolute continuity of the distributions, we have that Problems (3) are equivalent to the minimizing movement scheme (5):

**Proposition 2.2**.: _The distributions \(\rho_{k+1}^{\tau}=(T_{k}^{\tau})_{t}\rho_{k}^{\tau}\), where the functions \(T_{k}^{\tau}\) are found by solving (3) and \(\rho_{1}^{\tau}=\rho\) is the data distribution, coincide with the MMS (5) for \(\mathcal{Z}=\min_{F}\mathcal{L}(F,.)\)._

Proof.: The minimizing movement scheme (5) is equivalent to taking \(\rho_{k+1}^{\tau}=(T_{k}^{\tau})_{\tau}\rho_{k}^{\tau}\) where

\[T_{k}^{\tau}\in\arg\min_{T:\Omega\rightarrow\Omega}\mathcal{Z}(T_{\sharp}\rho _{k}^{\tau})+\frac{1}{2\tau}W_{2}^{2}(T_{\sharp}\rho_{k}^{\tau},\rho_{k}^{\tau }) \tag{6}\]

under conditions that guarantee the existence of a transport map between \(\rho_{k}^{\tau}\) and any other measure, and absolute continuity of \(\rho_{k}^{\tau}\) suffices, and the loss can ensure that \(\rho_{k+1}^{\tau}\) is also absolutely continuous. Among the functions \(T_{k}^{\tau}\) that solve problem (6), is the optimal transport map from \(\rho_{k}^{\tau}\) to \(\rho_{k+1}^{\tau}\). To solve specifically for this optimal transport map, we have to solve the equivalent Problem

\[T_{k}^{\tau}\in\arg\min_{T}\mathcal{Z}(T_{\sharp}\rho_{k}^{\tau})+\frac{1}{2 \tau}\int_{\Omega}\|T(x)-x\|^{2}\,\mathrm{d}\rho_{k}^{\tau}(x) \tag{7}\]

Problems (6) and (7) have the same minimum value, but the minimizer of (7) is now an optimal transport map between \(\rho_{k}^{\tau}\) and \(\rho_{k+1}^{\tau}\). This is immediate from the definition (4) of the \(W_{2}\) distance. Equivalently minimizing first over \(F\) and then over \(T\) in (3), it follows from the definition of \(\mathcal{Z}\) that Problems (3) and (7) are equivalent, which concludes.

Since we solve Problems (3) over neural networks, their representation power shown by universal approximation theorems [13; 29] is important to get close to equivalence between (5) and (3), as we need to approximate an optimal transport map. We also know that the training of each module, if it is shallow, converges [5; 6; 34; 22; 18].

If \(\mathcal{Z}\) is lower-semi continuous then Problems (5) always admit a solution because \(\mathcal{P}(\Omega)\) is compact. If \(\mathcal{Z}\) is also \(\lambda\)-geodesically convex for \(\lambda{>}0\), we have convergence of \(\rho_{k}^{\tau}\) as \(k{\rightarrow}\infty\) and \(\tau{\rightarrow}0\) to a minimizer of \(\mathcal{Z}\), potentially under more technical conditions (see Appendix B). Even though a machine learning loss will usually not satisfy these conditions, this analysis offers hints as to why our method avoids in practice the problem of stagnation or collapse in performance of module-wise training as \(k\) increases, as we are making proximal local steps in Wasserstein space to minimize the loss. This convergence discussion also suggests taking \(\tau\) as small as possible and many modules.

### Regularity result

As a secondary result, we show that Problem (3) has a solution and that the solution module \(T_{k}^{\tau}\) is an optimal transport map between its input and output distributions, which means that it comes with some regularity. [36] show that these networks generalize better and overfit less in practice. We assume that the minimization in \(F\) is over a compact set \(\mathcal{F}\), that \(\rho_{k}^{\tau}\) is absolutely continuous, that \(\mathcal{L}\) is continuous and non-negative, that \(\Omega\) is convex and compact and that \(\partial\Omega\) is negligible.

**Proposition 2.3**.: _Problem (3) has a minimizer \((T_{k}^{\tau},F_{k}^{\tau})\) such that \(T_{k}^{\tau}\) is an optimal transport map. And for any minimizer \((T_{k}^{\tau},F_{k}^{\tau})\), \(T_{k}^{\tau}\) is an optimal transport map._

The proof is in Appendix C. OT maps have regularity properties under some boundedness assumptions. Given Theorem A.1 in Appendix A taken from [20], \(T_{k}^{\tau}\) is \(\eta\)-Holder continuous almost everywhere and if the optimization algorithm we use to solve the discretized problem (2) returns an approximate solution pair \((\tilde{F}_{k}^{\tau},\tilde{T}_{k}^{\tau})\) such that \(\tilde{T}_{k}^{\tau}\) is an \(\epsilon\)-optimal transport map, i.e. \(\|\tilde{T}_{k}^{\tau}-T_{k}^{\tau}\|_{\infty}\leq\epsilon\), then we have (using the triangle inequality) the following stability property of the module \(\tilde{T}_{k}^{\tau}\):

\[\|\tilde{T}_{k}^{\tau}(x)-\tilde{T}_{k}^{\tau}(y)\|\leq 2\epsilon+C\|x-y\|^{\eta} \tag{8}\]for almost every \(x,y\in\text{supp}(\rho_{k}^{\tau})\) and \(C{>}0\). Composing these stability bounds on \(T_{k}^{\tau}\) and \(\tilde{T}_{k}^{\tau}\) allows to get bounds for the composition networks \(G_{k}^{\tau}\) and \(\tilde{G}_{k}^{\tau}{=}\tilde{T}_{k}^{\tau}\circ\dots\circ\tilde{T}_{1}^{\tau}\).

To summarize Section 2, the transport regularization makes each module more regular and it allows the modules to build on each other as \(k\) increases to solve the task, which is the property we desire.

## 3 Practical implementation

### Multi-block modules

For simplicity, we presented in (2) the case where each module is a single residual block. However, in practice, we often split the network into modules that are made-up of many residual blocks each. We show here that regularizing the kinetic energy of such modules still amounts to a transport regularization, which means that the theoretical results in Propositions 2.2 and 2.3 still apply.

If each module \(T_{k}\) is made up of \(M\) residual blocks, i.e. applies \(x_{m+1}{=}x_{m}{+}r_{m}(x_{m})\) for \(0{\leq}m{<}M\), then its total discrete kinetic energy for a single data point \(x_{0}\) is the sum of its squared residue norms \(\sum\|r_{m}(x_{m})\|^{2}\), since a residual network can be seen as a discrete Euler scheme for an ordinary differential equation [61] with velocity field \(r\):

\[x_{m+1}=x_{m}+r_{m}(x_{m})\ \longleftrightarrow\ \partial_{t}x_{t}=r_{t}(x_{t}) \tag{9}\]

and \(\sum\limits_{m}\|r_{m}(x_{m})\|^{2}\) is then the discretization of the total kinetic energy \(\int_{0}^{1}\|r_{t}(x)\|^{2}\,\mathrm{d}t\) of the ODE. If \(\psi_{m}^{x}\) denotes the position of a point \(x\) after \(m\) residual blocks, then regularizing the kinetic energy of multi-block modules now means solving

\[(T_{k}^{\tau},F_{k}^{\tau})\in\arg\min_{T,F}\sum_{x\in\mathcal{D} }(L(F,T(G_{k-1}^{\tau}(x))+\frac{1}{2\tau}\sum_{m=0}^{M-1}\|r_{m}(\psi_{m}^{x} )\|^{2}) \tag{10}\] \[\text{s.t. }T=(\mathtt{id}+r_{M-1})\circ...\circ(\mathtt{id}+r_{0 }),\ \psi_{0}^{x}=G_{k-1}^{\tau}(x),\psi_{m+1}^{x}=\psi_{m}^{x}+r_{m}(\psi_{m}^{x})\]

where \(G_{k}^{\tau}{=}T_{k}^{\tau}\circ\dots\circ T_{1}^{\tau}\) for \(1{\leq}k{\leq}K\) and \(G_{0}^{\tau}{=}\mathtt{id}\). We also minimize this sum of squared residue norms instead of \(\|T(x)-x\|^{2}\) (the two no longer coincide) as it works better in practice, which we assume is because it offers a more localized control of the transport. As expressed in (9), a residual network can be seen as an Euler scheme of an ODE and Problem (10) is then the discretization of

\[(T_{k}^{\tau},F_{k}^{\tau})\in\arg\min_{T,F}\ \mathcal{L}(F,T_{ \sharp}\rho_{k}^{\tau})+\frac{1}{2\tau}\int_{0}^{1}\|v_{t}\|_{L^{2}((\phi_{t}) \rho_{k}^{\tau})}^{2}\,\mathrm{d}t \tag{11}\] \[\text{s.t. }T=\phi_{1},\ \partial_{t}\phi_{t}^{x}=v_{t}(\phi_{t}^ {x}),\ \phi_{0}^{\cdot}=\mathtt{id}\]

where \(\rho_{k+1}^{\tau}=(T_{k}^{\tau})_{\sharp}\rho_{k}^{\tau}\) and \(r_{m}\) is the discretization of vector field \(v_{t}\) at time \(t=m/M\). Here, distributions \(\rho_{k}^{\tau}\) are pushed forward through the maps \(T_{k}^{\tau}\) which correspond to the flow \(\phi\) at time \(t=1\) of the kinetically-regularized velocity field \(v_{t}\). We recognize in the second term in the target of (11) the optimal transport problem in its dynamic formulation (15) from [9], and given the equivalence between the Monge OT problem (4) and the dynamic OT problem (15) in Appendix A, Problem (11) is in fact equivalent to the original continuous formulation (3), and the theoretical results in Section 2 follow immediately (see also the proof of Proposition 2.3 in Appendix C).

### Solving the module-wise problems

The module-wise problems can be solved in two ways. One can completely train each module with its auxiliary classifier for \(N\) epochs before training the next module, which receives as input the output of the previous trained module. We call this _sequential_ module-wise training. But we can also do this batch-wise, i.e. do a complete forward pass on each batch but without a full backward pass, rather a backward pass that only updates the current module \(T_{k}^{\tau}\) and its auxiliary classifier \(F_{k}^{\tau}\), meaning that \(T_{k}^{\tau}\) forwards its output to \(T_{k+1}^{\tau}\) immediately after it computes it. We call this _parallel_ module-wise training. It is called _decoupled_ greedy training in [8], which shows that combining it with batch buffers solves all three locking problems and allows a linear training parallelization in the depth of the network. We propose a variant of sequential module-wise training that we call _multi-lap sequential_ module-wise training, in which instead of training each module for \(N\) epochs, we train each module from the first to the last sequentially for \(N/R\) epochs, then go back and train from the first module to the last for \(N/R\) epochs again, and we do this for \(R\) laps. For the same total number of epochs and training time, and the same advantages (loading and training one module at a time) this provides a non-negligible improvement in accuracy over normal sequential module-wise training in most cases, as shown in Section 4. Despite our theoretical framework being that of sequential module-wise training, our method improves the test accuracy of all three module-wise training regimes.

### Varying the regularization weight

The discussion in Section 2.2 suggests taking a fixed weight \(\tau\) for the transport cost that is as small as possible. However, instead of using a fixed \(\tau\), we might want to vary it along the depth \(k\) to further constrain with a smaller \(\tau_{k}\) the earlier modules to avoid that they overfit or the later modules to maintain the accuracy of earlier modules. We might also want to regularize the network further in earlier epochs when the data is more entangled. We propose in Appendix D to formalize this varying weight \(\tau_{k,i}\) across modules \(k\) and SGD iterations \(i\) by using a scheme inspired by the method of multipliers to solve Problems (2) and (10). However, it works best in only one experiment in practice. The observed dynamics of \(\tau_{k,i}\) suggest simply finding a fixed value of \(\tau\) that is multiplied by 2 for the second half of the network, which works best in all the other experiments (see Appendix E).

## 4 Experiments

We call our method TRGL for Transport-Regularized Greedy Learning. For the auxiliary classifiers, we use the architecture from DGL [7; 8], that is a convolution followed by an average pooling and a fully connected layer, which is very similar to that used by InfoPro [60], except for the Swin Transformer where we use a linear layer. We call vanilla greedy module-wise training with the same architecture but without our regularization VanGL, and we include its results in all tables for ablation study purposes. The code is available at github.com/block-wise/module-wise and implementation details are in Appendix E.

### Parallel module-wise training

To compare with other methods, we focus first on parallel training, as it performs better than sequential training and has been more explored recently. The first experiment is training in parallel 3 residual architectures and a VGG-19 [52] divided into 4 modules of equal depth on TinyImageNet. We compare in Table 1 our results in this setup to three of the best recent parallel module-wise training methods: DGL [8], PredSim [45] and Sedona [47], from Table 2 in [47]. We find that our TRGL has a much better test accuracy than the three other methods, especially on the smaller architectures. It also performs better than end-to-end training on the three ResNets. Parallel TRGL in this case with 4 modules consumes \(10\) to \(21\%\) less memory than end-to-end training (with a batch size of 256).

The second experiment is training in parallel two ResNets divided into 2 modules on CIFAR100 [37]. We compare in Table 2 our results in this setup to the two delayed gradient methods DDG [32] and FR [31] from Table 2 in [31]. Here again, parallel TRGL has a better accuracy than the other two methods and than end-to-end training. With only two modules, the memory gains from less backpropagation are neutralized by the weight of the extra classifier and there are negligible memory savings compared to end-to-end training. However, parallel TRGL has a better test accuracy by up to almost 2 percentage points.

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline Architecture & Parallel VanGL & Parallel TRGL (ours) & PredSim & DGL & Sedona & E2E \\ \hline VGG-19 & 56.17 \(\pm\) 0.29 (\(\downarrow 27\%\)) & **57.28 \(\pm\)** 0.20 (\(\downarrow 21\%\)) & 44.70 & 51.40 & 56.56 & 58.74 \\ ResNet-50 & 58.43 \(\pm\) 0.45 (\(\downarrow 26\%\)) & **60.30 \(\pm\)** 0.58 (\(\downarrow 20\%\)) & 47.48 & 53.96 & 54.40 & 58.10 \\ ResNet-101 & 63.64 \(\pm\) 0.30 (\(\downarrow 24\%\)) & **63.71 \(\pm\)** 0.40 (\(\downarrow 11\%\)) & 53.92 & 53.80 & 59.12 & 62.01 \\ ResNet-152 & 63.87 \(\pm\) 0.16 (\(\downarrow 21\%\)) & **64.23 \(\pm\)** 0.14 (\(\downarrow 10\%\)) & 51.76 & 57.64 & 64.10 & 62.32 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Test accuracy of parallel TRGL with 4 modules (average and 95\(\%\) confidence interval over 5 runs) on TinyImageNet, compared to DGL, PredSim, Sedona and E2E from Table 2 in [47], with memory saved compared to E2E as a percentage of E2E memory consumption in red.

The third experiment is training in parallel a ResNet-110 divided into two, four, eight and sixteen modules on STL10 [12]. We compare in Table 3 our results in this setup to the recent methods InfoPro [60] and DGL [8] from Table 2 in [60]. TRGL largely outperforms the other methods. It also outperforms end-to-end training in all but one case (that with 16 modules). With a batch size of 64, memory savings of parallel TRGL compared to end-to-end training reach \(48\%\) and \(58.5\%\) with 8 and 16 modules respectively, with comparable test accuracy. With 4 modules, TRGL training weighs \(24\%\) less than end-to-end-training, and has a test accuracy that is better by \(2\) percentage points (see Section 4.2 and Table 6 for a detailed memory usage comparison with InfoPro).

The fourth experiment is training (from scratch) in parallel a Swin-Tiny Transformer [39] divided into 4 modules on three datasets. We compare in Table 4 our results with those of InfoPro [60] and InfoProL, a variant of InfoPro proposed in [46]. TRGL outperforms the other module-wise training methods. It does not outperform end-to-end training in this case, but consumes \(29\%\) less memory on CIFAR10 and CIFAR100 and \(50\%\) less on STL10, compared to \(38\%\) for InfoPro and \(45\%\) for InfoProL in [46].

Finally, we compare our method to InfoPro, DGL and Sedona in Table 5 below on a large scale experiment on ImageNet.

### Memory savings and training time

As seen above, parallel TRGL is lighter than end-to-end training by up to almost \(60\%\). The extra memory consumed by our regularization compared to parallel VanGL is between 2 and \(13\%\) of

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline Dataset & Parallel VanGL & Parallel TRGL (ours) & InfoPro & InfoProL & E2E \\ \hline STL10 & 67.00 \(\pm\) 1.36 (\(\downarrow 55\%\)) & **67.92**\(\pm\) 1.12 (\(\downarrow 50\%\)) & 64.61 (\(\downarrow 38\%\)) & 66.89 (\(\downarrow 45\%\)) & 72.19 \\ CIFAR10 & 83.94 \(\pm\) 0.42 (\(\downarrow 33\%\)) & **86.48**\(\pm\) 0.54 (\(\downarrow 29\%\)) & 83.38 (\(\downarrow 38\%\)) & 86.28 (\(\downarrow 45\%\)) & 91.37 \\ CIFAR100 & 69.34 \(\pm\) 0.91 (\(\downarrow 33\%\)) & **74.11**\(\pm\) 0.31 (\(\downarrow 29\%\)) & 68.36 (\(\downarrow 38\%\)) & 73.00 (\(\downarrow 45\%\)) & 75.03 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Test accuracy of parallel TRGL with 4 modules (average and 95\(\%\) confidence interval over 5 runs) on a Swin-Tiny Transformer, compared to InfoPro, InfoProL and E2E from Table 3 in [46], with memory saved compared to E2E as a percentage of E2E memory consumption in red.

\begin{table}
\begin{tabular}{c c c c c c} \hline \hline Architecture & Parallel VanGL & Parallel TRGL (ours) & DDG & FR & E2E \\ \hline ResNet-101 & 77.31 \(\pm\) 0.27 & **77.87**\(\pm\) 0.44 & 75.75 & 76.90 & 76.52 \\ ResNet-152 & 75.40 \(\pm\) 0.75 & **76.55**\(\pm\) 1.90 & 73.61 & 76.39 & 74.80 \\ \hline \hline \end{tabular}
\end{table}
Table 2: Test accuracy of parallel TRGL with 2 modules (average and 95\(\%\) confidence interval over 3 runs) on CIFAR100, compared to DDG, FR and E2E from Table 2 in [31].

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline \(K\) & Par VanGL & Par TRGL (ours) & DGL & InfoPro S & InfoPro C & E2E \\ \hline
2 & 79.85 \(\pm\) 0.93 & **80.04**\(\pm\) 0.85 & 75.03 \(\pm\) 1.18 & 78.98 \(\pm\) 0.51 & 79.01 \(\pm\) 0.64 & 77.73 \(\pm\) 1.61 \\
4 & 77.11 \(\pm\) 2.31 & **79.72**\(\pm\) 0.81 & 73.23 \(\pm\) 0.64 & 78.72 \(\pm\) 0.27 & 77.27 \(\pm\) 0.40 & 77.73 \(\pm\) 1.61 \\
8 & 75.71 \(\pm\) 0.55 & **77.82**\(\pm\) 0.73 & 72.67 \(\pm\) 0.24 & 76.40 \(\pm\) 0.49 & 74.85 \(\pm\) 0.52 & 77.73 \(\pm\) 1.61 \\
16 & 73.57 \(\pm\) 0.95 & **77.22**\(\pm\) 1.20 & 72.27 \(\pm\) 0.58 & 73.95 \(\pm\) 0.71 & 73.73 \(\pm\) 0.48 & 77.73 \(\pm\) 1.61 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Test accuracy of Parallel (Par) TRGL with \(K\) modules (average and 95\(\%\) confidence interval over 5 runs) using a ResNet-110 on STL10, compared to DGL, two variants of InfoPro and E2E from Table 2 in [60].

end-to-end memory. Memory savings depend then mainly on the size of the auxiliary classifier, which can easily be adjusted. Note that delayed gradients method DDG and FR increase memory usage [31], and Sedona does not claim to save memory, but rather to speed up training [47]. DGL is architecture-wise essentially identical to VanGL and consumes the same memory.

We compare in Table 6 the memory consumption of our method to that of InfoPro [60] on a ResNet-110 on STL10 with a batch size of 64 (so the same setting as in Table 3). InfoPro [60] also propose to split the network into modules that have the same weight but not necessarily the same number of layers. They only implement this for \(K{\leq}4\) modules. When the modules are even in weight and not in depth, we call the training methods VanGL*, TRGL* and InfoPro*. In practice, this leads to shallower early modules which slightly hurts performance according to [47], and as seen below. However, TRGL* still outperforms InfoPro and end-to-end training, and it leads to even bigger memory savings than InfoPro*. We see in Table 6 below that TRGL saves more memory than InfoPro in two out of three cases (4 and 8 modules), and about the same in the third case (16 modules), with much better test accuracy in all cases. Likewise, TRGL* is lighter than InfoPro*, with better accuracy.

However, parallel module-wise training does slightly slow down training. Epoch time increases by \(6\%\) with 2 modules and by \(16\%\) with 16 modules. TRGL is only slower than VanGL by \(2\%\) for all number of modules due to the additional regularization term. This is comparable to InfoPro which reports a time overhead between 1 and \(27\%\) compared to end-to-end training. See Appendix F for details.

### Sequential full block-wise training

Block-wise sequential training, meaning that each module is a single residual block and that the blocks are trained sequentially, therefore requiring only enough memory to train one block and its classifier. Even though it has been less explored in recent module-wise training methods, it has been used in practice in very constrained settings such as on-device training [58; 57]. We therefore test our regularization in this section in this setting, with more details in Appendix G.

We propose here to use shallower ResNets that are initially wider. These architectures are well-adapted to layer-wise training as seen in [7]. We check first in Table 10 in Appendix G that this architecture works well with parallel module-wise training with 2 modules by comparing it favorably on CIFAR10 [37] with methods DGL [8], InfoPro [60] and DDG [32] that use a ResNet-110 with the same number of parameters.

We then train a 10-block ResNet block-wise on CIFAR100. In Tables 11 and 12 in Appendix G, we see that MLS training improves the accuracy of sequential training by 0.8 percentage points when the trainset is full, but works less well on small train sets. Of the two, the regularization mainly improves the test accuracy of MLS training. The improvement increases as the training set gets smaller and reaches 1 percentage point. While parallel module-wise training performs quite close to end-to-end training in the full data regime and much better in the small data regime, sequential and multi-lap sequential training are competitive with end-to-end training in the small data regime. Combining the multi-lap trick and the regularization improves the accuracy of sequential training by 1.2 percentage points when using the entire trainset. We report further results for full block-wise training on MNIST [38] and CIFAR10 [37] in Tables 13 and 14 in Appendix G.

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline  & \multicolumn{4}{c}{Equally deep modules} & \multicolumn{4}{c}{Equally heavy modules} \\ \cline{2-7} \(K\) & Par VanGL & Par TRGL (ours) & InfoPro & Par VanGL* & Par TRGL* (ours) & InfoPro* \\ \hline
4 & 27\% (77.11) & 24\% (**79.72**) & 18\% (78.72) & 41\% (77.14) & 39\% (**78.94**) & 33\% (78.78) \\
8 & 50\% (75.71) & 48\% (**77.82**) & 37\% (76.40) & & & \\
16 & 61\% (73.57) & 58\% (**77.22**) & 59\% (73.95) & & & \\ \hline \hline \end{tabular}
\end{table}
Table 6: Memory savings using a ResNet-110 on STL10 split into \(K\) modules trained in parallel with a batch size of 64, as a percentage of the weight of end-to-end training. Average test accuracy over 5 runs is between brackets. Test accuracy of end-to-end training is 77.73%.

The \(88\%\) accuracy of sequential training on CIFAR10 in Table 13 is the same as in Table 2 of [7], which is the best method for layer-wise sequential training available, with VGG networks of comparable depth and width.

### Accuracy after each module

Finally, we verify that our method avoids the stagnation or collapse in accuracy with depth. In Figure 2 below, we show the accuracy of each module with and without the regularization.

On the left, from parallel module-wise training experiments from Table 3, TRGL performs worse than vanilla greedy learning early, but surpasses it in later modules, indicating that it does avoid early overfitting. On the right, from sequential block-wise training experiments from Table 13, we see a large decline in performance that the regularization avoids. We see similar patterns in Figure 3 in Appendix G with parallel and MLS block-wise training.

## 5 Limitations

The results in Appendix G show a few limitations of our method, as the improvements from the regularization are sometimes minimal on sequential training. However, the results show that our approach works in all settings (parallel and sequential with many or few modules), whereas other papers don't test their methods in all settings, and some show problems in other settings than the original one in subsequent papers (e.g. delayed gradients methods when the number of modules increases [31] and PredSim in [47]). Also, for parallel training in Section 4.1, the improvement from the regularization compared to VanGL is larger and increases with the number of modules (so with the memory savings) and reaches almost 5 percentage points. We show in Appendix H that our method is not very sensitive to the choice of hyperparameter \(\tau\) over a large scale.

## 6 Related work

Layer-wise training was initially considered as a pre-training and initialization method [10; 43] and was recently shown to be competitive with end-to-end training [7; 45]. Many papers consider using a different auxiliary loss, instead of or in addition to the classification loss: kernel similarity [42], information-theory-inspired losses [53; 44; 41; 60] and biologically plausible losses [53; 45; 25; 11]. Methods [7], PredSim [45], DGL [8], Sedona [47] and InfoPro [60] report the best module-wise training results. [7; 8] do it simply through the architecture choice of the auxiliary networks. Sedona applies architecture search to decide on where to split the network into modules and what auxiliary classifier to use before module-wise training. Only BoostResNet [30] also proposes a block-wise training idea geared for ResNets. However, their results only show better early performance and end-to-end fine-tuning is required to be competitive. A method called ResIST [19] that is similar to block-wise training of ResNets randomly assigns ResBlocks to one of up to 16 modules that

Figure 2: Test accuracy after each module averaged over 10 runs with \(95\%\) confidence intervals. Left: parallel vanilla (VanGL, in blue) and regularized (TRGL, in red) module-wise training of a ResNet-110 with 16 modules on STL10 (Table 3). Right: sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training of a 10-block ResNet on \(2\%\) of CIFAR10 (Table 13).

are trained independently and reassembled before another random partition. More of a distributed training method, it is only compared to local SGD [54]. These methods can all be combined with our regularization, and we do use the auxiliary classifier from [7, 8].

Besides module-wise training, methods such as DNI [33, 14], DDG [32] and FR [31], solve the update and backward locking problems with an eye towards parallelization by using delayed or predicted gradients, or even predicted inputs to address forward locking, which is what [55] do. But they observe training issues with more than 5 modules [31]. This makes them compare unfavorably to module-wise training [8]. The high dimension of the predicted gradient, which scales with the size of the network, makes [33, 14] challenging in practice. Therefore, despite its simplicity, greedy module-wise training is more appealing when working in a constrained setting.

Viewing ResNets as dynamic transport systems [16, 36] followed from their view as a discretization of ODEs [61]. Transport regularization of ResNets in particular is motivated by the observation that they are naturally biased towards minimally modifying their input [35, 26]. We further linked this transport viewpoint with gradient flows in the Wasserstein space to apply it in a principled way to module-wise training. Gradient flows on the data distribution appeared recently in deep learning. In [1], the focus is on functionals of measures whose first variations are known in closed form and used, through their gradients, in the algorithm. This limits the scope of their applications to transfer learning and similar tasks. Likewise, [21, 40, 4, 3] use the explicit gradient flow of \(f\)-divergences and other distances between measures for generation and generator refinement. In contrast, we use the discrete minimizing movement scheme which does not require computation of the first variation and allows to consider classification.

## 7 Conclusion

We introduced a transport regularization for module-wise training that theoretically links module-wise training to gradient flows of the loss in probability space. Our method provably leads to more regular modules and experimentally improves the test accuracy of module-wise parallel, sequential and multi-lap sequential (a variant of sequential training that we introduce) training. Through this simple method that does not complexify the architecture, we make module-wise training competitive with end-to-end training while benefiting from its lower memory usage. Being a regularization, the method can easily be combined with other layer-wise training methods. Future work can experiment with working in Wasserstein space \(W_{p}\) for \(p{\neq}2\), i.e. regularizing with a norm \(\|\cdot\|_{p}\) with \(p{\neq}2\).

## References

* [1]Alvarez-Melis, D., and Fusi, N. Dataset dynamics via gradient flows in probability space. _ICML_ (2021).
* [2]Ambrosio, L., Gigli, N., and Savare, G. _Gradient Flows in Metric Spaces and in the Space of Probability Measures_. Birkhauser Basel, 2005.
* [3]Ansari, A. F., Ang, M. L., and Soh, H. Refining deep generative models via discriminator gradient flow. In _ICLR_ (2021).
* [4]Arbel, M., Korba, A., Salim, A., and Gretton, A. Maximum mean discrepancy gradient flow. In _NeurIPS_ (2019).
* [5]Arora, R., Basu, A., Mianjy, P., and Mukherjee, A. Understanding deep neural networks with rectified linear units. In _International Conference on Learning Representations_ (2018).
* [6]Bach, F. Breaking the curse of dimensionality with convex neural networks. _Journal of Machine Learning Research 18_, 19 (2017), 1-53.
* [7]Bellovsky, E., E., Eickenberg, M., and Oyallon, E. Greedy layerwise learning can scale to imagenet. In _ICML_ (2019).
* [8]Bellovsky, E., Eickenberg, M., and Oyallon, E. Decoupled greedy learning of cnns. In _ICML_ (2020).
* [9]Benamou, J., and Brenier, Y. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem. _Numerische Mathematik_ (2000).

* [10]Bengio, Y., Lamblin, P., Popovici, D., and Larochelle, H. Greedy layer-wise training of deep networks. In _NeurIPS_ (2006).
* [11]Bernd Illing, Wulfram Gerstner, G. B. Towards truly local gradients with clapp: Contrastive, local and predictive plasticity. _arXiv_ (2020).
* [12]Coate, A., Lee, H., and Ng, A. Y. An analysis of single layer networks in unsupervised feature learning. In _AISTATS_ (2011).
* [13]Cybenko, G. Approximation by superpositions of a sigmoidal function. _Mathematics of Control, Signals and Systems 2_ (1989), 303-314.
* [14]Czarnecki, W. M., Swirszcz, G., Jaderberg, M., Osindero, S., Vinyals, O., and Kavukcuoglu, K. Understanding synthetic gradients and decoupled neural interfaces. In _ICML_ (2017).
* [15]De, S., and Smith, S. L. Batch normalization biases residual blocks towards the identity function in deep networks. In _NeurIPS_ (2020).
* [16]de Bezenac, E., Ayed, I., and Gallinari, P. Optimal unsupervised domain translation. _arXiv_ (2019).
* [17]Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In _International Conference on Learning Representations_ (2021).
* [18]Du, S. S., and Goel, S. Improved learning of one-hidden-layer convolutional neural networks with overlaps. _arXiv_ (2018).
* [19]Dun, C., Wolfe, C. R., Jermaine, C. M., and Kyrillidis, A. Resist: Layer-wise decomposition of resnets for distributed training. _arXiv_ (2021).
* [20]Figalli, A. _The Monge-Ampere Equation and Its Applications_. Zurich lectures in advanced mathematics. European Mathematical Society, 2017.
* [21]Gao, Y., Jiao, Y., Wang, Y., Wang, Y., Yang, C., and Zhang, S. Deep generative learning via variational gradient flow. In _ICML_ (2019).
* [22]Ge, R., Lee, J. D., and Ma, T. Learning one-hidden-layer neural networks with landscape design. In _International Conference on Learning Representations_ (2018).
* [23]Gianazza, U., Gobbino, M., and Savare, G. Evolution problems and minimizing movements. _Atti della Accademia Nazionale dei Lincei. Classe di Scienze Fisiche, Matematiche e Naturali. Rendiconti Lincei. Matematica e Applicazioni 5_ (1994), 289-296.
* [24]Giorgi, E. D., Marino, A., and Tosques, M. Problemi di evoluzione in spazi metrici e curve di massima pendenza. _Atti della Accademia Nazionale dei Lincei. Classe di Scienze Fisiche, Matematiche e Naturali. Rendiconti 68_, 3 (3 1980), 180-187.
* [25]Gupta, S. K. A more biologically plausible local learning rule for anns. _arXiv_ (2020).
* [26]Hauser, M. On residual networks learning a perturbation from identity. _arXiv_ (2019).
* [27]He, K., Zhan, X., Ren, S., and Sun, J. Identity mappings in deep residual networks. In _ECCV_ (2016).
* [28]He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In _CVPR_ (2016).
* [29]Hornik, K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. _Neural Networks 2_, 5 (1989), 359-366.
* [30]Huang, F., Ash, J., Langford, J., and Schapire, R. Learning deep resnet blocks sequentially using boosting theory. In _ICML_ (2018).
* [31]Huo, Z., Gu, B., and Huang, H. Training neural networks using features replay. In _NeurIPS_ (2018).
* [32]Huo, Z., Gu, B., Yang, Q., and Huang, H. Decoupled parallel backpropagation with convergence guarantee. In _ICML_ (2018).
* [33]Jaderberg, M., Czarnecki, W. M., Osindero, S., Vinyals, O., Graves, A., Silver, D., and Kavukcuoglu, K. Decoupled neural interfaces using synthetic gradients. In _ICML_ (2017).

* [34]Janzamin, M., Sedghi, H., and Anandkumar, A. Beating the perils of non-convexity: Guaranteed training of neural networks using tensor methods. _arXiv_ (2016).
* [35]Jastrzebski, S., et al. Residual connections encourage iterative inference. In _ICLR_ (2018).
* [36]Karkar, S., Ayed, I., de Bezenac, E., and Gallinari, P. A principle of least action for the training of neural networks. In _ECML-PKDD_ (2020).
* [37]Krizhevsky, A. Learning multiple layers of features from tiny images. _University of Toronto Technical Report_ (2009).
* [38]LeCun, Y., Cortes, C., and Burges, C. J. MNIST handwritten digit database. _yann.lecun.com/exdb/mnist_ (2010).
* [39]Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In _Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)_ (2021).
* [40]Liutkus, A., Ilsekli, U., Majewski, S., Durmus, A., and Stoter, F.-R. Sliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions. In _ICML_ (2019).
* [41]Ma, W.-D. K., Lewis, J., and Kleijn, W. B. The hsic bottleneck: Deep learning without back-propagation. In _AAAI_ (2020).
* [42]Mandar Kulkarni, S. K. Layer-wise training of deep networks using kernel similarity. In _DLPR workshop, ICPR_ (2016).
* [43]Marquez, E. S., Hare, J. S., and Niranjan, M. Deep cascade learning. _IEEE Transactions on Neural Networks and Learning Systems_ (2018).
* [44]Nguyen, T. T., and Choi, J. Layer-wise learning of stochastic neural networks with information bottleneck. _Entropy 21_ (2019).
* [45]Nokland, A., and Eidnes, L. H. Training neural networks with local error signals. In _ICML_ (2019).
* [46]Pathak, P., Zhang, J., and Samaras, D. Local learning on transformers via feature reconstruction, 2022.
* [47]Pyeon, M., Moon, J., Hahn, T., and Kim, G. Sedona: Search for decoupled neural networks toward greedy block-wise learning. In _ICLR_ (2021).
* [48]Sandler, M., Howard, A. G., Zhu, M., Zhmoginov, A., and Chen, L. Mobilenetv2: Inverted residuals and linear bottlenecks. In _2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018_ (2018), Computer Vision Foundation / IEEE Computer Society, pp. 4510-4520.
* [49]Santambrogio, F. _Optimal Transport for Applied Mathematicians_. Birkhauser, 2015.
* [50]Santambrogio, F. Euclidean, metric, and wasserstein gradient flows: an overview. _arXiv_ (2016).
* [51]Saxe, A. M., Mcclelland, J. L., and Ganguli, S. Exact solutions to the nonlinear dynamics of learning in deep linear neural network. In _ICLR_ (2014).
* [52]Simonyan, K., and Zisserman, A. Very deep convolutional networks for large-scale image recognition. In _International Conference on Learning Representations (ICLR 2014)_ (2014).
* [53]Sindy Lowe, Peter O'Connor, B. V. Putting an end to end-to-end: Gradient-isolated learning of representations. In _NeurIPS_ (2019).
* [54]Stich, S. U. Local sgd converges fast and communicates little. In _ICLR_ (2019).
* [55]Sun, Q., Dong, H., Chen, Z., Sun, J., Li, Z., and Dong, B. Layer-parallel training of residual networks with auxiliary-variable networks. _arXiv_ (2021).
* [56]Tan, M., and Le, Q. EfficientNet: Rethinking model scaling for convolutional neural networks. In _Proceedings of the 36th International Conference on Machine Learning_ (09-15 Jun 2019), K. Chaudhuri and R. Salakhutdinov, Eds., vol. 97 of _Proceedings of Machine Learning Research_, PMLR, pp. 6105-6114.
* [57]Tang, Y., Teng, Q., Zhang, L., Min, F., and He, J. Layer-wise training convolutional neural networks with smaller filters for human activity recognition using wearable sensors. _IEEE Sensors Journal_ (2021).

* [58]Teng, Q., Wang, K., Zhang, L., and He, J. The layer-wise training convolutional neural networks using local loss for sensor-based human activity recognition. _IEEE Sensors Journal_ (2020).
* [59]Villani, C. _Optimal Transport: Old and New_. Springer-Verlag, 2008.
* [60]Wang, Y., Ni, Z., Song, S., and Le Yang, G. H. Revisiting locally supervised learning: an alternative to end-to-end training. In _ICLR_ (2021).
* [61]Weinan, E. A proposal on machine learning via dynamical systems. _Commun. Math. Stat_ (2017).
* [62]Xie, S., et al. Aggregated residual transformations for deep neural networks. In _CVPR_ (2017).
* [63]Zagoruyko, S., and Komodakis, N. Wide residual networks. In _BMVC_ (2016).
* [64]Zhang, J., et al. Towards robust resnet: A small step but a giant leap. In _IJCAI_ (2019).
* MICCAI 2022_ (Cham, 2022), Springer Nature Switzerland, pp. 192-201.
Background on optimal transport

The Wasserstein space \(\mathbb{W}_{2}(\Omega)\) with \(\Omega\) a convex and compact subset of \(\mathbb{R}^{d}\) is the space \(\mathcal{P}(\Omega)\) of probability measures over \(\Omega\), equipped with the distance \(W_{2}\) given by the solution to the optimal transport problem

\[W_{2}^{2}(\alpha,\beta)=\min_{\gamma\in\Pi(\alpha,\beta)}\int_{\Omega\times \Omega}\|x-y\|^{2}\,\mathrm{d}\gamma(x,y) \tag{12}\]

where \(\Pi(\alpha,\beta)\) is the set of probability distribution over \(\Omega\times\Omega\) with first marginal \(\alpha\) and second marginal \(\beta\), i.e. \(\Pi(\alpha,\beta)=\{\gamma\in\mathcal{P}(\Omega\times\Omega)\mid\pi_{1\sharp} \gamma=\alpha,\;\pi_{2\sharp}\gamma=\beta\}\) where \(\pi_{1}(x,y)=x\) and \(\pi_{2}(x,y)=y\). The optimal transport problem can be seen as looking for a transportation plan minimizing the cost of displacing some distribution of mass from one configuration to another. This problem indeed has a solution in our setting and \(W_{2}\) can be shown to be a geodesic distance (see for example [49, 59]). If \(\alpha\) is absolutely continuous and \(\partial\Omega\) is \(\alpha\)-negligible then the problem in (12) (called the Kantorovich problem) has a unique solution and is equivalent to the Monge problem, i.e.

\[W_{2}^{2}(\alpha,\beta)=\min_{T\,s.t.\;T_{\sharp}\alpha=\beta}\int_{\Omega}\|T (x)-x\|^{2}\,\mathrm{d}\alpha(x) \tag{13}\]

and this problem has a unique solution \(T^{\star}\) linked to the solution \(\gamma^{\star}\) of (12) through \(\gamma^{\star}=(\mathrm{id},T^{\star})_{\sharp}\alpha\). Another equivalent formulation of the optimal transport problem in this setting is the dynamical formulation [9]. Here, instead of directly pushing samples of \(\alpha\) to \(\beta\) using \(T\), we can equivalently displace mass, according to a continuous flow with velocity \(v_{t}:\mathbb{R}^{d}\to\mathbb{R}^{d}\). This implies that the density \(\alpha_{t}\) at time \(t\) satisfies the _continuity equation_\(\partial_{t}\alpha_{t}+\nabla\cdot(\alpha_{t}v_{t})=0\), assuming that initial and final conditions are given by \(\alpha_{0}=\alpha\) and \(\alpha_{1}=\beta\) respectively. In this case, the optimal displacement is the one that minimizes the total action caused by \(v\) :

\[W_{2}^{2}(\alpha,\beta)= \min_{v}\int_{0}^{1}\|v_{t}\|_{L^{2}(\alpha_{t})}^{2}\,\mathrm{d}t \tag{14}\] \[\mathrm{s.t.}\;\partial_{t}\alpha_{t}+\nabla\cdot(\alpha_{t}v_{t })=0,\;\alpha_{0}=\alpha,\alpha_{1}=\beta\]

Instead of describing the density's evolution through the continuity equation, we can describe the paths \(\phi_{t}^{x}\) taken by particles at position \(x\) from \(\alpha\) when displaced along the flow \(v\). Here \(\phi_{t}^{x}\) is the position at time \(t\) of the particle that was at \(x\sim\alpha\) at time 0. The continuity equation is then equivalent to \(\partial_{t}\phi_{t}^{x}=v_{t}(\phi_{t}^{x})\). See chapters 4 and 5 of [49] for details. Rewriting the conditions as necessary, Problem (14) becomes

\[W_{2}^{2}(\alpha,\beta)= \min_{v}\int_{0}^{1}\|v_{t}\|_{L^{2}((\phi_{t})_{\sharp}\alpha)} ^{2}\,\mathrm{d}t \tag{15}\] \[\mathrm{s.t.}\;\partial_{t}\phi_{t}^{x}=v_{t}(\phi_{t}^{x}),\; \phi_{0}=\mathrm{id},(\phi_{1})_{\sharp}\alpha=\beta\]

and the optimal transport map \(T^{\star}\) that solves (13) is in fact \(T^{\star}(x)=\phi_{1}^{x}\) for \(\phi\) that solves the continuity equation together with the optimal \(v^{\star}\) from (15). We refer to [49, 59] for these results on optimal transport.

Optimal transport maps have some regularity properties under some boundedness assumptions. We mention the following result from [20]:

**Theorem A.1**.: _Let \(\alpha\) and \(\beta\) be absolutely continuous measures on \(\mathbb{R}^{d}\) and \(T\) the optimal transport map between \(\alpha\) and \(\beta\) for the Euclidean cost. Suppose there are bounded open sets \(X\) and \(Y\), such that the density of \(\alpha\) (respectively of \(\beta\)) is null on \(X^{\mathrm{c}}\) (respectively \(Y^{\mathrm{c}}\)) and bounded away from zero and infinity on \(X\) (respectively \(Y\))._

_Then there exists two relatively closed sets of null measure \(A\subset X\) and \(B\subset Y\), such that \(T\) is \(\eta\)-Holder continuous from \(X\setminus A\) to \(Y\setminus B\), i.e. \(\forall\;x,y\in X\setminus A\) we have_

\[\|T(x)-T(y)\|\leq C\|x-y\|^{\eta}\text{ for constants }\eta,C>0\]

## Appendix B Background on gradient flows

We follow [50, 2] for this background on gradient flows. Given a function \(\mathcal{L}:\mathbb{R}^{d}\to\mathbb{R}\) and an initial point \(x_{0}\in\mathbb{R}^{d}\), a _gradient flow_ is a curve \(x:[0,\infty[\to\mathbb{R}^{d}\) that solves the Cauchy problem

\[\begin{cases}x^{\prime}(t)=-\nabla\mathcal{L}(x(t))\\ x(0)=x_{0}\end{cases} \tag{16}\]A solution exists and is unique if \(\nabla\mathcal{L}\) is Lipschitz or \(\mathcal{L}\) is convex. Given \(\tau>0\) and \(x_{0}^{\tau}=x_{0}\) define a sequence \((x_{k}^{\tau})_{k}\) through the _minimizing movement scheme_:

\[x_{k+1}^{\tau}\in\arg\min_{x\in\mathbb{R}^{d}}\ \mathcal{L}(x)+\frac{1}{2\tau} \|x-x_{k}^{\tau}\|^{2} \tag{17}\]

\(\mathcal{L}\) lower semi-continous and \(\mathcal{L}(x)\geq C_{1}-C_{2}\|x\|^{2}\) guarantees existence of a solution of (17) for \(\tau\) small enough. \(\mathcal{L}\)\(\lambda\)-convex meets these conditions and also provides uniqueness of the solution because of strict convexity of the target. See [49, 50, 2].

We interpret the point \(x_{k}^{\tau}\) as the value of a curve \(x\) at time \(k\tau\). We can then construct a curve \(x^{\tau}\) as the piecewise constant interpolation of the points \(x_{k}^{\tau}\). We can also construct a curve \(\tilde{x}^{\tau}\) as the affine interpolation of the points \(x_{k}^{\tau}\).

If \(\mathcal{L}(x_{0})<\infty\) and \(\inf\mathcal{L}>-\infty\) then \((x^{\tau})\) and \((\tilde{x}^{\tau})\) converge uniformly to the same curve \(x\) as \(\tau\) goes to zero (up to extracting a subsequence). If \(\mathcal{L}\) is \(\mathcal{C}^{1}\), then the limit curve \(x\) is a solution of (16) (i.e. a gradient flow of \(\mathcal{L}\)). If \(\mathcal{L}\) is not differentiable then \(x\) is solution of the problem defined using the subdifferential of \(\mathcal{L}\), i.e. \(x\) satisfies \(x^{\prime}(t)\in-\partial\mathcal{L}(x(t))\) for almost every \(t\).

If \(\mathcal{L}\) is \(\lambda\)-convex with \(\lambda>0\), then the solution to (16) converges exponentially to the unique minimizer of \(\mathcal{L}\) (which exists by coercivity). So taking \(\tau\to 0\) and \(k\to\infty\), we tend towards the minimizer of \(\mathcal{L}\).

The advantage of the minimizing movement scheme (17) is that it can be adapted to metric spaces by replacing the Euclidean distance by the metric space's distance. In the (geodesic) metric space \(\mathbb{W}_{2}(\Omega)\) with \(\Omega\) convex and compact, for \(\mathcal{L}:\mathbb{W}_{2}(\Omega)\to\mathbb{R}\cup\{\infty\}\) lower semi-continuous for the weak convergence of measures in duality with \(\mathcal{C}(\Omega)\) (equivalent to lower semi-continuous with respect to the distance \(W_{2}\)) and \(\rho_{0}^{\tau}=\rho_{0}\in\mathcal{P}(\Omega)\), the minimizing movement scheme (17) becomes

\[\rho_{k+1}^{\tau}\in\arg\min_{\rho\in\mathcal{P}(\Omega)}\ \mathcal{L}(\rho)+ \frac{1}{2\tau}W_{2}^{2}(\rho,\rho_{k}^{\tau}) \tag{18}\]

This problem has a solution because the objective is lower semi-continuous and the minimization is over \(\mathcal{P}(\Omega)\) which is compact by Banach-Alaoglu.

We can construct a piecewise constant interpolation between the measures \(\rho_{k}^{\tau}\), or a geodesic interpolation where we travel along a geodesic between \(\rho_{k}^{\tau}\) and \(\rho_{k+1}^{\tau}\) in \(\mathbb{W}_{2}(\Omega)\), constructed using the optimal transport map between these measures. Again, if \(\mathcal{L}(x_{0})<\infty\) and \(\inf\mathcal{L}>-\infty\) then both interpolations converge uniformly to a limit curve \(\tilde{\rho}\) as \(\tau\) goes to zero. Under further conditions on \(\mathcal{L}\), mainly \(\lambda\)-geodesic convexity (i.e. \(\lambda\)-convexity along geodesics) for \(\lambda>0\), we can prove stability and convergence of \(\tilde{\rho}(t)\) to a minimizer of \(\mathcal{L}\) as \(t\to\infty\), see [49, 50, 2].

## Appendix C Proof of Proposition 2.3

Proof.: Take a minimizing sequence \((\tilde{F}_{i},\tilde{T}_{i})\), i.e. such that \(\mathcal{C}(\tilde{F}_{i},\tilde{T}_{i})\to\min\mathcal{C}\), where \(\mathcal{C}\geq 0\) is the target function in (3) and denote \(\beta_{i}=\tilde{T}_{i*}\rho_{k}^{\tau}\). Then by compacity \(\tilde{F}_{i}\to F^{\star}\) and \(\beta_{i}\rightharpoonup\beta^{\star}\) in duality with \(\mathcal{C}_{b}(\Omega)\) by Banach-Alaoglu. There exists \(T^{\star}\) an optimal transport map between \(\rho_{k}^{\tau}\) and \(\beta^{\star}\). Then \(\mathcal{C}(F^{\star},T^{\star})\leq\lim\mathcal{C}(\tilde{F}_{i},\tilde{T}_{i })=\min\mathcal{C}\) by continuity of \(\mathcal{L}\) and because

\[\int_{\Omega}\|T^{\star}(x)-x\|^{2}\,\mathrm{d}\rho_{k}^{\tau}(x) =W_{2}^{2}(\rho_{k}^{\tau},\beta^{\star})\] \[=\lim W_{2}^{2}(\rho_{k}^{\tau},\beta_{i})\] \[\leq\lim\int_{\Omega}\|\tilde{T}_{i}(x)-x\|^{2}\,\mathrm{d}\rho_{ k}^{\tau}(x)\]

as \(W_{2}\) metrizes weak convergence of measures. We take \((F_{k}^{\tau},T_{k}^{\tau})=(F^{\star},T^{\star})\). It is also immediate that for any minimizing pair, the transport map has to be optimal. Taking a minimizing sequence \((\tilde{F}_{i},\tilde{v}^{i})\) and the corresponding induced maps \(\tilde{T}_{i}\) we get the same result for Problem (11). Problems (3) and (11) are equivalent by the equivalence between Problems (13) and (15).

## Appendix D Varying the regularization weight

The discussion in Section 2.2 suggests taking a fixed weight \(\tau\) for the transport cost that is as small as possible. However, instead of using a fixed \(\tau\), we might want to vary it along the depth \(k\) to further constrain with a smaller \(\tau_{k}\) the earlier modules to avoid that they overfit or the later modules to maintain the accuracy of earlier modules. We might also want to regularize the network further in earlier epochs when the data is more entangled. To unify and formalize this varying weight \(\tau_{k,i}\) across modules \(k\) and SGD iterations \(i\), we use a scheme inspired by the method of multipliers to solve Problems (2) and (10). To simplify the notations, we will instead consider the weight \(\lambda_{k,i}{:=}2\tau_{k,i}\) given to the loss. We denote \(\theta_{k,i}\) the parameters of both \(T_{k}\) and \(F_{k}\) at SGD iteration \(i\). We also denote \(L(\theta,x)\) and \(W(\theta,x)\) respectively the loss and the transport regularization as functions of parameters \(\theta\) and data \(x\). We now increase the weight \(\lambda_{k,i}\) of the loss every \(s\) iterations of SGD by a value that is proportional to the current loss. Given increase factor \(h{>}0\), initial parameters \(\theta_{k,1}\), initial weights \(\lambda_{k,1}{\geq}0\), learning rates (\(\eta_{i}\)) and batches (\(x_{i}\)), we apply for module \(k\) and \(i{\geq}1\):

\[\begin{cases}\theta_{k,i+1}&=\theta_{k,i}-\eta_{i}\nabla_{\theta}(\lambda_{k,i }\,L(\theta_{k,i},x_{i})+W(\theta_{k,i},x_{i}))\\ \lambda_{k,i+1}&=\lambda_{k,i}+hL(\theta_{k,i+1},x_{i+1})\text{ if }i\text{ mod }s=0 \text{ else }\lambda_{k,i}\end{cases}\]

The weights \(\lambda_{k,i}\) will vary along modules \(k\) because they will evolve differently with iterations \(i\) for each \(k\). They will increase more slowly with \(i\) for larger \(k\) because deeper modules will have smaller loss. This method can be seen as a method of multipliers for the problem of minimizing the transport under the constraint of zero loss. Therefore it is immediate by slightly adapting the proof of Proposition 2.3 or from [36] that we are still solving a problem that admits a solution whose non-auxiliary part is an optimal transport map with the same regularity as stated above. We use the same initial value \(\lambda_{1}=\lambda_{k,1}\) for all modules so that this method requires choosing three hyper-parameters (\(h\), \(s\) and \(\lambda_{1}\)). In practice (see Section 4.1 and Appendix E), it works best in only one experiment. Simply manually finding a value of \(\tau\) that is multiplied by 2 for the second half of the network works best in all the other experiments.

## Appendix E Implementation details

We use standard data augmentation and standard implementations for VGG-19, ResNet-50, ResNet-101, ResNet-110, ResNet-152 and Swin-Tiny Transformer (the same as for the other methods in Section 4.1). We use NVIDIA Tesla V100 16GB GPUs for the experiments. Training a Resnet-152 on TinyImageNet in Table 1 takes about 36 hours. Training a Resnet-152 on CIFAR100 in Table 1 takes about 11 hours. Training a ResNet-110 on STL10 in Table 3 takes about 3 hours. Training a Swin-Tiny Transformer in Table 4 take between 40 minutes and 1 hour.

For sequential and multi-lap sequential training, we use SGD with a learning rate of \(0.007\). With the exception of the Swin Transformer in Table 4, we use SGD for parallel training with learning rate of \(0.003\) in all Tables but Table 3 where the learning rate is \(0.002\). For the Swin Transformer in Table 4, we use the AdamW optimizer with a learning rate of \(0.007\) and a CosineLR scheduler.

For end-to-end training we use a learning rate of 0.1 that is divided by five at epochs 120, 160 and 200. Momentum is always 0.9. For parallel and end-to-end training, we train for 300 epochs. For sequential and multi-lap sequential training, the number of epochs varies per module (see Section G).

For experiments in Section 4.1, we use a batch size of 256, orthogonal initialization [51] with a gain of 0.1, label smoothing of 0.1 and weight decay of 0.0002. The batch size changes to 64 for Table 3 and to 1024 for Table 4.

For experiments in Section 4.3, we use a batch size of 128, orthogonal initialization with a gain of 0.05, no label smoothing and weight decay of 0.0001.

In Table 1, we use \(\tau=500000\) for the first two modules and then double it for the last two modules for TRGL. In Table 2, we use \(\lambda_{k,1}=1,h=1\) and \(s=50\) for TRGL. In Table 3, we use \(\tau=0.5\) and double it at the midpoint, expect for the first line where \(\tau=50\).

## Appendix F Memory savings and training time

We compare in Table 7 the memory consumption of our method to that of InfoPro [60] on a ResNet-110 split into \(K\) modules trained in parallel on STL10 with a batch size of 64 (so the same setting as in Table 3 in Section 4.1). We report in Table 7 the memory saved as a percentage of the 6230 MiB memory required by end-to-end training with the same batch size. VanGL refers to our architecture trained without the regularization. InfoPro [60] also propose to split the network into \(K\) modules that have the same weight but not necessarily the same number of layers. They only implement this for \(K{\leq}4\) modules. When the modules are even in weight and not in depth, we call the training methods VanGL*, TRGL* and InfoPro*. In practice, this leads to shallower early modules which slightly hurts performance according to [47]. We verify this in Table 8 (to be compared with Table 3 in Section 4.1). However, TRGL* still outperforms InfoPro and end-to-end training, and it leads to even bigger memory savings. We see in Table 7 that TRGL saves more memory than InfoPro in two out of three cases (4 and 8 modules), and about the same in the third case (16 modules), with much better test accuracy in all cases. Likewise, TRGL* is lighter than InfoPro*, with better accuracy. We also see that the added memory cost of the regularization compared to vanilla greedy learning is small. However, parallel module-wise training does slightly slow down training. Epoch time increases by \(6\%\) with 2 modules and by \(16\%\) with 16 modules. TRGL is only slower than VanGL by \(2\%\) for for all number of modules due to the additional regularization term. This is comparable to InfoPro which report a time overhead between 1 and \(27\%\) compared to end-to-end training.

Note that methods DDG [32] and FR [31], being delayed gradient methods and not module-wise training methods, do no save memory (they actually increase memory usage, see FR [31]). Sedona [47] also does not claim to save memory, as their first module (the heaviest) is deeper than the others, but rather to speed up computation. Finally, DGL [8] is architecture-wise essentially identical to VanGL and consumes the same memory.

## Appendix G Sequential full block-wise training

To show that our method works well with all types of module-wise training when using few modules, we train a ResNet-101 split in 2 modules on CIFAR100, sequentially and multi-lap sequentially. Results are in Table 9. We see that our idea of multi-lap sequential training adds one percentage point of accuracy to sequential training, and that the regularization further improves the accuracy by about half a percentage point. As only one module has to be trained at a time, these two training methods require only around half the memory end-to-end training requires (the size of the heaviest module and its classifier more exactly).

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline  & \multicolumn{4}{c}{Equally deep modules} & \multicolumn{4}{c}{Equally heavy modules} \\ \cline{2-7} \(K\) & Par VanGL & Par TRGL (ours) & InfoPro & Par VanGL* & Par TRGL* (ours) & InfoPro* \\ \hline
4 & 27\% (77.11) & 24\% (**79.72**) & 18\% (78.72) & 41\% (77.14) & 39\% (**78.94**) & 33\% (78.78) \\
8 & 50\% (75.71) & 48\% (**77.82**) & 37\% (76.40) & & & \\
16 & 61\% (73.57) & 58\% (**77.22**) & 59\% (73.95) & & & \\ \hline \hline \end{tabular}
\end{table}
Table 7: Memory savings using a ResNet-110 on STL10 split into \(K\) modules trained in parallel with a batch size of 64, as a percentage of the weight of end-to-end training. Average test accuracy over 5 runs is between brackets. Test accuracy of end-to-end training is 77.73\(\%\).

\begin{table}
\begin{tabular}{c c c c} \hline \hline \(K\) & Par VanGL* & Par TRGL* (ours) & InfoPro* \\ \hline
2 & 79.05 \(\pm\) 1.33 & **79.47**\(\pm\) 1.36 & 79.05 \(\pm\) 0.57 \\
4 & 77.14 \(\pm\) 1.23 & **78.94**\(\pm\) 1.13 & 78.78 \(\pm\) 0.72 \\ \hline \hline \end{tabular}
\end{table}
Table 8: Test accuracy of parallel (Par) TRGL* with \(K\) modules (average and 95\(\%\) confidence interval over 5 runs) on a ResNet-110 trained on STL10, compared to InfoPro* and E2E training from Table 3 in [60]

\begin{table}
\begin{tabular}{c c c c} \hline \hline Seq VanGL & Seq TRGL & MLS VanGL & MLS TRGL \\ \hline
73.31 & 73.61 & 74.34 & **74.78** \\ \hline \hline \end{tabular}
\end{table}
Table 9: Test accuracy of sequential (Seq) and multi-lap sequential (MLS) TRGL and VanGL with 2 modules on CIFAR100 using ResNet-101 (average of 2 runs).

We now focus on full block-wise training, meaning that each module is a single ResBlock, mostly sequentially. We propose here to use shallower and initially wider ResNets with a downsampling and 256 filters initially and a further downsampling and doubling of the number of filters at the midpoint, no matter the depth. In these ResNets, we use the ResBlock from [27] with two convolutional layers. If such a network is divided in \(K\) modules of \(M\) ResBlocks each, we call the network a \(K\)\(-\)\(M\) ResNet. These wider shallower architectures are well-adapted to layer-wise training as seen in [7]. We check in Table 10 that this architecture works well with parallel module-wise training by comparing favorably on CIFAR10 ([37]) a 2-7 ResNet with DGL, InfoPro ([60]) and DDG [32]. The 2-7 ResNet has 45 millions parameters, which is about the same as the ResNet-110 divided in two used by the other methods, and performs better when trained in parallel.

We now train a 10-block ResNet block-wise on CIFAR100 (a 10-1 ResNet in our notations). We report even the small improvements in accuracy to show that our method works in all settings (parallel or sequential with many or few splits), which other methods don't do. For sequential training, block \(k\) is trained for \(50\)\(+\)\(10k\) epochs where \(0\)\(\leq\)\(k\)\(\leq\)\(10\), block 0 being the encoder. This idea of increasing the number of epochs along with the depth is found in [43]. For MLS training, block \(k\) is trained for \(10\)\(+\)\(2k\) epochs, and this is repeated for 5 laps. In block-wise training, the last block does not always perform the best and we report the accuracy of the best block. In Table 11, we see that MLS training improves the test accuracy of sequential training by around \(0.8\) percentage points when the training dataset is full, but works less well on small training sets. Of the two, the regularization mainly improves the test accuracy of MLS training. The improvement increases as the training set gets smaller and reaches 1 percentage point. That is also the case for parallel module-wise training in Table 12, which already performs quite close to end-to-end training in the full data regime and much better in the small data regime. Combining the multi-lap trick and the regularization improves the performance of sequential training by 1.2 percentage points.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline Parallel VanGL (ours) & Parallel TRGL (ours) & DGL & DDG & InfoPro \\ \hline
94.01 \(\pm\).17 & **94.05**\(\pm\).18 & 93.50 & 93.41 & 93.58 \\ \hline \hline \end{tabular}
\end{table}
Table 10: Average test accuracy and 95\(\%\) confidence interval of 2-7 ResNet over 10 runs on CIFAR10 with parallel TRGL and VanGL, compared to DGL and DDG from [8] and InfoPro from [60] that split a ResNet-110 in 2 module-wise-parallel-trained modules.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline Train size & \multicolumn{1}{c}{Seq VanGL} & Seq TRGL & MLS VanGL & MLS TRGL & E2E \\ \hline
50000 & 68.74 \(\pm\) 0.45 & **68.79**\(\pm\) 0.56 & 69.48 \(\pm\) 0.53 & **69.95**\(\pm\) 0.50 & 75.85 \(\pm\) 0.70 \\
25000 & 60.48 \(\pm\) 0.15 & **60.59**\(\pm\) 0.14 & 61.33 \(\pm\) 0.23 & **61.71**\(\pm\) 0.32 & 65.36 \(\pm\) 0.31 \\
12500 & 51.64 \(\pm\) 0.33 & **51.74**\(\pm\) 0.26 & 51.30 \(\pm\) 0.22 & **51.89**\(\pm\) 0.30 & 52.39 \(\pm\) 0.97 \\
5000 & 36.37 \(\pm\) 0.33 & **36.40**\(\pm\) 0.40 & 33.68 \(\pm\) 0.48 & **34.61**\(\pm\) 0.59 & 36.38 \(\pm\) 0.31 \\ \hline \hline \end{tabular}
\end{table}
Table 11: Average highest test accuracy and 95\(\%\) confidence interval of 10-1 ResNet over 10 runs on CIFAR100 with different train sizes and sequential (Seq), multi-lap sequential (MLS) and parallel (Par) TRGL and VanGL, compared to E2E.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Train size & Par VanGL & Par TRGL & E2E \\ \hline
50000 & 72.59 \(\pm\) 0.40 & **72.63**\(\pm\) 0.40 & 75.85 \(\pm\) 0.70 \\
25000 & 64.84 \(\pm\) 0.19 & **65.01**\(\pm\) 0.27 & 65.36 \(\pm\) 0.31 \\
12500 & 55.13 \(\pm\) 0.24 & **55.40**\(\pm\) 0.35 & 52.39 \(\pm\) 0.97 \\
5000 & 39.45 \(\pm\) 0.23 & **40.36**\(\pm\) 0.23 & 36.38 \(\pm\) 0.31 \\ \hline \hline \end{tabular}
\end{table}
Table 12: Average highest test accuracy and 95\(\%\) confidence interval of 10-1 ResNet over 10 runs on CIFAR100 with different train sizes and sequential (Seq), multi-lap sequential (MLS) and parallel (Par) TRGL and VanGL, compared to E2E.

We report further results of block-wise training on CIFAR10 in Table 13 and on MNIST [38] in Table 14, but now we report the accuracy of the last block. We see again greater improvement due to the regularization as the training set gets smaller, gaining up to 6 percentage points.

The \(88\%\) accuracy of sequential training on CIFAR10 in Table 13 is the same as for sequential training in table 2 of [7], which is the best method for layer-wise sequential training available, with VGG networks of comparable depth and width.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Train size & Seq VanGL & Seq TRGL & E2E \\ \hline
50000 & 88.02 \(\pm\).18 & **88.20 \(\pm\).24 & 91.88 \(\pm\).18 \\
25000 & 83.95 \(\pm\).13 & **84.28 \(\pm\).22 & 88.75 \(\pm\).27 \\
10000 & 76.00 \(\pm\).39 & **77.18 \(\pm\).34 & 82.61 \(\pm\).35 \\
5000 & 67.74 \(\pm\).49 & **69.67 \(\pm\)**.44 & 73.93 \(\pm\).67 \\
1000 & 45.67 \(\pm\).88 & **51.34 \(\pm\)**.90 & 50.63 \(\pm\).98 \\ \hline \hline \end{tabular}
\end{table}
Table 13: Average last block test accuracy and 95\(\%\) confidence interval of 10-1 ResNet over 10 runs on CIFAR10 with different train sizes and sequential TRGL and VanGL, compared to E2E.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Train size & Par VanGL & Par TRGL & E2E \\ \hline
60000 & 99.07 \(\pm\).04 & **99.08 \(\pm\)**.04 & 99.30 \(\pm\).03 \\
30000 & 98.90 \(\pm\).05 & **98.93 \(\pm\)**.06 & 99.22 \(\pm\).03 \\
12000 & 98.52 \(\pm\).06 & **98.59 \(\pm\)**.06 & 98.96 \(\pm\).06 \\
6000 & 98.05 \(\pm\).09 & **98.16 \(\pm\)**.07 & 98.62 \(\pm\).06 \\
1500 & 96.34 \(\pm\).12 & **96.91 \(\pm\)**.07 & 97.19 \(\pm\).08 \\
1200 & 95.80 \(\pm\).12 & **96.58 \(\pm\)**.09 & 96.88 \(\pm\).09 \\
600 & 91.35 \(\pm\).99 & **95.16 \(\pm\)**.15 & 95.30 \(\pm\).17 \\
300 & 89.81 \(\pm\).73 & **92.86 \(\pm\)**.24 & 92.87 \(\pm\).28 \\
150 & 81.84 \(\pm\) 1.22 & **87.48 \(\pm\)**.42 & 87.82 \(\pm\).59 \\ \hline \hline \end{tabular}
\end{table}
Table 14: Average last block test accuracy and 95\(\%\) confidence interval of 20-1 ResNet (32 filters, fixed encoder, same classifier) over 20/50 runs on MNIST with different train sizes and parallel TRGL and VanGL, compared to E2E.

Figure 3: Test accuracy after each block of 10-1 ResNet averaged over 10 runs with \(95\%\) confidence intervals. Left: multi-lap sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training on \(10\%\) of the CIFAR100 training set. Right: parallel vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training on \(10\%\) of CIFAR100 training set.

Sensitivity to hyper-parameters

We show in Figure 4 below that TRGL still performs better than VanGL (in the same setting as in Table 3 in Section 4.1) for values of \(\tau\) from 0.03 to 100 and is still roughly equivalent to it for values up to 5000.

## Appendix I Broader impact

Less memory usage has a positive environmental impact and allows organizations with less resources to use deep learning.

Figure 4: Average test accuracy over 5 runs of parallel TRGL using a ResNet110 on STL10 with 16 modules with different values of \(\tau\) (in red), and of VanGL (blue line).

Response:
## Review

### Summary

This paper presents a new training method for greedy layer-wise or module-wise training of neural networks, specifically targeting the stagnation problem often encountered in traditional training approaches. By introducing a module-wise regularization inspired by the minimizing movement scheme, the method significantly improves accuracy while reducing memory usage, making it suitable for constrained and on-device settings. The authors provide theoretical analyses and experimental results to support their claims, demonstrating the effectiveness of their approach across various architectures, including ResNets and Transformers. Overall, the submission shows promise but raises some questions regarding convergence analysis and the scalability of the method on larger datasets.

### Strengths

- The proposed method significantly reduces memory usage compared to end-to-end training, making it particularly suitable for on-device training.
- The authors provide a theoretical analysis, showing that their approach leads to more regular and stable greedy modules.
- The methodology effectively addresses the stagnation issue associated with layer-wise training, resulting in improved accuracy.
- The link between their training method and optimal transport theory provides intriguing theoretical insights.
- The paper is well-structured and generally easy to follow.

### Weaknesses

- The introduction of an additional regularization term increases the computational complexity and may lead to longer training times.
- There is a lack of detailed convergence analysis, which is crucial for establishing the reliability of the proposed approach.
- The performance gain over previous methods appears marginal, especially on smaller datasets like STL-10.
- Concerns were raised regarding the strength of the regularization and its potential negative impact on early layers.
- The experimental results do not sufficiently support claims regarding early overfitting and the effectiveness of the method on large-scale datasets.

### Questions

- What is VanGL? Is it just a module-wise training method without the regularization terms?
- Why does the proposed method demonstrate less memory savings compared to other methods?
- How does the proposed regularization control each module's strength, particularly with varying numbers of modules?
- What is the wall-time training cost of the proposed method compared to existing methods?
- Some notations in the paper are not introduced properly; for example, what is the difference between $T_k$ and $T_k^{\tau}$?

### Soundness

**Score:** 3

**Description:** 3 = good; the methodology is theoretically sound and shows promise, but there are concerns regarding convergence and the robustness of the results.

### Presentation

**Score:** 3

**Description:** 3 = good; the paper is generally well-written and structured, though some convoluted notations could be improved for better clarity.

### Contribution

**Score:** 3

**Description:** 3 = good; the paper contributes a new approach to module-wise training which could have significant implications, although further validation on larger datasets is needed.

### Rating

**Score:** 5

**Description:** 5 = borderline accept; the paper is technically solid with reasons to accept outweighing reasons to reject, though it requires further evaluation on larger datasets.

### Paper Decision

**Decision:** Accept (poster)

**Reasons:** The paper presents an innovative training approach that addresses significant issues in neural network training. While there are some weaknesses, particularly in the convergence analysis and the scalability of the results, the overall contributions are valuable and the authors have effectively addressed several reviewer concerns during the rebuttal phase. The strengths of the paper, including its theoretical foundation and practical applications, support the decision to accept it for presentation.