Module-wise Training of Neural Networks via the Minimizing Movement Scheme
Authors: Skander Karkar
Ibrahim Ayed
Emmanuel de Bezenac
patrick gallinari
Venue: NeurIPS 2023 poster (2023)
Keywords: Deep learning
greedy layerwise training
memory
optimal transport
TL;DR: We use the minimizing movement scheme for gradient flows in distribution space to regularize and improve the performance of greedy module-wise training of neural networks, requiring less memory than end-to-end training.

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 simple 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.

Module-wise Training of Neural Networks via the Minimizing Movement Scheme
Skander Karkar Criteo, Sorbonne Université
Ibrahim Ayed Sorbonne Université, Thales
Emmanuel de Bézenac ETH Zurich
Patrick Gallinari Criteo, Sorbonne Université
Abstract
Greedy layer-wise or module-wiss wise training of neural networks compelling in constrained and on-device settings where memory is limited, as it cireumvents number of problems of end-to-end back procp--apagation However, it suffers from a stagnation problem, wherehy early layers overfit and deeper layers stop increasing the test accuracy after a certain depth. We propose to solve this issue by introducing module-wise regularization inspired by the minimizing movement scheme for gradient flows in distribution space. We call the method TRGL for Transport Regu- larized Greedy Learning and study it theoretically, proving that leads to greedy modules that are regular and that progressively solve the task. Experimentally, we show improved accuracy of module-wise 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 end-to-ent training, with as much as 60% less memory usage.
1 Introduction
End-to-end backpropagation is the standard training method of neural networks. However, 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, 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. the same time batch-wise), consumes much less memory than end-to-end training 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-wil wise training also allows for parallel training of the modules 88. Despite its
it does
simplicity, module-wine wise training has been recently shown to scale well g
60 45 outperforming
more complicated alternatives to end-to-ent training such as synthetic 33 and delayed [32 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 more layers. Each module takes input the output of the previous module, and each module has an auxiliary classifier so that local loss can be computed, with backpropagation happening only inside the modules and not between them (see Figure below).
The main drawback of module-wise training is the well-documented stagnation problem observed in 60 47 whereby early modules overfit and learn more discriminative features than end-to-end
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

training, destroying tele-aelevant information, and deeper modules don't '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 and Section 4.4 To tackle this issue, InfoPro [60] propose maximize the mutual information that each module keeps with the input, 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 different perspective, leveraging the analogy between residual connections and the Euler scheme for ODEs [6]. To preserve input information, 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.
()
1(22. y)
1(23, y)
input
ø1
22
23
y label
1 loss
Classifier
Classifier F2
Classifier
Ik output of Tp
F1
F3
2k prediction of Fi
D2
33
Forward-prop
x
Module Ti
Module T2
Module T3
Back-prop
Figure 1: Module-wise training.
Our approach is particularly well-adapted to networks that use residual connections such as ResNets
27 28 their variants (e.g. ResNeXt 622, 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 optimal transport map making it more regular and stable. We also show that amounts to 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 aadule-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 introduce and that performs better in many cases. In partieular, our regularization makes parallel module-wise training superior or comparable in accuracy end-to-end training, while consuming 10% 60% less memory.
2 Transport-regularized module-wise training
The typical setting of (sequential) module-wise training for minimizing loss L, is, given dataset D, to solve one after the other, for I<kSK, Problems
(Te, Fi) E arg min T.F
Σ
LFFT Gx-1

---

where G₁ = Ti for I<kSK, Go=id, Ta is the module (one many layers) and Fk. is an auxiliary classifier. Module T₁ receives the output of module -T- and auxiliary classifier Fk. computes the prediction from the output of TE so the loss can be computed. The inputs are and I. has access to their lahels y to calculate the loss. i.e. [[F, T G (x)) = I( o T ( y) where 1 is machine learning loss such as cross-entropy. See Figure The final network trained this way is F Gg. But, at inference, we can stop at any depth & and use F o Gg 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
60
We propose below in (2) 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
2.1 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
function T=id+-r, which includes many transformer architectures [39 17),, its kinetic energy is simply the squared norm of its residue r= which we add to the loss in the target of the greedy problems [T. All layers that have the same input and output dimension can be rewritten residual blocks and the analysis applies to large variety architectures such as VGG
Given T>0, we now solve, for 1sks K, Problems
(T) E arg min Σ ((F,TOGG-(0)) T,F RED
1 IT 2T
(x)11
(2)
where oT. for ISKSK and -id The final network F 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 (xx) and have the same dimension so that T(x) can be computed.
To facilitate the theoretical analysis, we rewrite the method in more general formulation using data distribution p. which can be discrete or continuous, and the distribution loss L that arises from the point-wise loss L. Then Problem (2) is equivalent to Problem
2r [ IT(x) - - dx)(x))
(T PF) E arg min L(F, (1Ph) +
T.F
with (1P) pI=p and ((F,T()) = f (p)))(()(())))) dpt(x) = - 2(F,2) dT))(()))
2.2 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 2(u) == minF L(F, ()), which
is the loss of the best classifier. If we are limited to linear classifiers. Z(pZ) is the linear separability of the representation p at module of the data distribution p. The MMS, introduced in [24] 23), is a
metric counterpart to Euclidean gradient descent for minimizing functionals over distributions. In our case, Z is the functional we want to minimize. We define the MMS below in Definition 2..
The distribution space we work in is the metric Wasserstein space (2 = (P(S2), W2), where Ω C Idd is convex compact set, P(2) is the set of probability distributions over Q and W2 is the Wasserstein distance over P(2) derived from the optimal transport problem with Euclidean cost:
/
W (a. 3)
min Tst. T,a=8
[[((x)
x]| da(zz)
where we assume that 02 is negligible and that the distributions absolutely continous.
3

---

Definition 2.1. Given 2: (2 - R. and starting from pi E P(Ω), the Minimizing Movement
Scheme (MMS) takes proximal steps for minimizing 2. s given by
Po+t E arg min PEP(N)
2(p) + 2T W (p)
(5)
The MMS (5) can be seen as a non-Euclidean implicit Euler step for following the gradient flow of 2, and P converges to a minimizer of Z under some conditions (see the end of this section).
So under the mentioned assumptions on 52 and absolute continuity of the distributions, we have that Problems (3) are equivalent to the minimizing movement scheme 5
Proposition 2.2. The distributions P = (T )apk where the functions T ure found by solving
3 and pi = p is the data distribution, coincide with the MMS 5 for = minp L(F,-).
Proof. The minimizing movement scheme 5 is equivalent to taking Pī+1 = (T PPp where
1 T E arg min p(Tpp + 2T (TsP.. T-R-9
(6)
under conditions that guarantee the existence of a transport map between Pi. and any other measure, and absolute continuity of P suffices, and the loss can ensure that P is also absolutely continuous. Among the functions 7 that solve problem (6), is the optimal transport map from P to Pk+1 To solve specifically for this optimal transport map, we have to solve the equivalent Problem
T E arg min ((TP)) + 27 / [T(x)) x||² dpī(()
(7)
T
Problems (6) and 7 have the same minimum value, but the minimizer of (7) is now an optimal
transport map between PE and Pk+1- This is immediate from the definition (4) of the W2 distance.
Equivalently minimizing first over F and then over 7 in (3). it follows from the definition of 2 that Problems and 77 are equivalent, which concludes.
Since we solve Problems (3) over neural networks, their representation power shown by universal
approximation theorems [3
29 is important to get close to equivalence between (5) and (3) as we
need to approximate optimal transport map. We also know that the training of each module, if it is
shallow, converges 55 6 34 22 18].
If Z is lower-semi continuous then Problems 5 always admit solution because P(22) is compact. If Z also X-geodesically convex for A=0, we have convergence of P as k=000 and 7-0 to a minimizer of Z, potentially under more technical conditions (see Appendix B). Even though machine learning loss will usually not satisfy these conditions, this analysis ofters 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 as small as possible and many modules.
2.3 Regularity result
As secondary result, we show that Problem (3) has solution and that the solution module T5 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 F, that pi is absolutely continuous, that L is continuous and non-negative that Q is convex and compact and that On is negligible.
Proposition 2.3. Problem (3) has a minimizer (TH. FF) such that Tg is an optimal transport map. And for any minimizer (Tg. PE), TV is an optimal transport map.
The proof is in Appendix C OT maps have regularity properties under some boundedness assumptions. Given Theorem A. in Appendix A taken from (0]. T is Hölder continuous almost everywhere and if the optimization algorithm we use to solve the discretized problem (2) returns an approximate solution pair (FT T7) such that Ty is an e-optimal transport map, i.c. I1T, T 1| S e, then we have (using the triangle inequality) the following stability property of the module TT:
( T (v)
2e + Cl| y/1"
(8)

---

for almost every y E supp(p?) and C>00. Composing these stability bounds on T and Tg allows to get bounds for the composition networks T% and G o.. TT.
To summarize Section 2 the transport regularization makes each module more regular and allows the modules to build on each other as k increases to solve the task, which is the property we desire.
3 Practical implementation
3.1 Multi-block modules
For simplicity, we presented in (2 the case where each module is 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 transport regularization, which means that the theoretical results in Propositions 2.2 and 2.3 still apply.
If each module Tę is made up of M residual blocks, i.e. applies Im+1 =m +T (Im) ) for 05m< M, then its total discrete kinetic energy for single data point zg is the sum of squared residue norms Σ Im (2m)]l since a residual network can be seen as discrete Euler scheme for an ordinary differential equation [6]] with velocity field r:
Tm+11
Zm + I'm (1m)
0 (21)
2|2 is then the discretization of the total kinetic energy In (x) di of the ODE. If when denotes the position of point after m residual blocks, then regularizing the kinetic energy of multi-block modules now means solving
and Σ |I (2mm)
M-1
(Th, FT) E arg min T,F EDD
Σ (LL(F,TTGG
Σ rl 13))
(10)
2T
s.t. T (id + .... (id + ro); v/0 = - vm+1 = V. + Tm (v%)
where oT. for I<kSK and =ī=id. We also minimize this of squared residue norms instead of [x) - = 1| (the two no longer coincide) as it works better in practice, which we assume because it offers a more localized control the transport. As expressed can be seen an Euler scheme of an ODE and Problem (10) is then the discretization of
(9), residual network
2T /."
1
(TE,P) PE) E arg min ((F,T;PR) +
|| ( de
(11)
s.t. == 0,0, = (6)) 0 = id
where +i+1 (I% Dapk and rm is the discretization of vector field i, at time m/M. Here, distributions PT are pushed forward through the maps T% which correspond to the flow á at time velocity field vt. We recognize in the second term in the target of ITT the optimal transport problem in its dynamic formulation (15) from ( and given the equivalence
t 1 of the pnpeticall--reullarie
between the Monge OT problem (4) and the dynamic OT problem (15) Appendix [ Problem (1I)
is in fact equivalent to the original continuous formulation (3) and the theoretical results in Section
follow immediately (see also the proof of Proposition 2.3 in Appendix C).
3.2 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 complete forward pass on each batch but without full backward pass, rather backward pass that only updates the current module T and its auxiliary classifier Fÿ, meaning that Ty forwards its output to T 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 the network. We propose variant of sequential module-wist training that we call multi-lap sequential module-wise training, 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
5

---

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 a time) this provides 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-wiss training, our method improves the test accuracy of all three module-wise training regimes.
3.3 Varying the regularization weight
The discussion in Section 2.2 suggests taking a fixed weight r for the transport cost that is as small as possible. However, instead of using fixed т, we might want to vary it along the depth & to further constrain with a smaller The 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 Dito formalize this varying weight T across modules lt and SGD iterations by using scheme inspired by the method of multipliers to solve Problems (2) and (TO). However, it works best in only one experiment in practice. The observed dynamics of The suggest simply finding fixed value of that is multiplied by 2 for the second half of the network, which works best in all the other experiments (see Appendix
Experiments
We call our method TRGL for Transport-Regularre Greedy Learning. For the auxiliary classifiers, we use the architecture from DGL 87. that is a convolution followed by average pooling and a fully connected layer, which is very similar to that used by InfoPro [60], except for the Swin Transformer where 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/ htt///lookkwiisee weikwwwii//mmiiieieeeeeeee htt//moolleeeeeeeee -wis and implementation details are in Appendix E
linear layer. We call vanilla greedy module-wise training with the
4.1 Parallel module-wise training
To compare with other methods, focus first on parallel training, it performs better than sequential training and has been more explored recently. The first experiment is training in parallel residual architectures and VGG-19 [52] divided into modules of equal depth TinylmageNet. We compare in Table I our results in this setup to three of the best recent parallel module-wise training methods: DGL (8. PredSim (45] and Sedona 477, from Table 2 [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-ent tod training on the three ResNets. Parallel TRGL in this case with modules consumes 10 to 21% less memory than end-to-end training (with a batch size of 256).
Table 1:
Test accuracy of parallel TRGL with 4 modules (average and 95% confidence interval over 5 runs) on TinyImageNer compared to DGL, PredSim, Sedona and E2E from Table in (7), with memory saved compared to E2E as percentage of E2E memory consumption in red.
Architecture
Parallel VanGL
Parallel TRGL (ours)
PredSim
DGL
Sedona
E2E
VGG-19
56.17 = 0.29 (L 27%)
57.28 0.20 (L 21%)
44.70
51.40
56.56
58.74
ResNet-50
58.43 = 0.45 (1 26%)
60.30 + 0.58 (1 20%)
47.48
53.96
54.40
58.10
ResNet-101
63.64 = 0.30 (1 24%)
63.71 + 0.40
(1 11%)
53.92
53.80
59.12
62.01
ResNet-152
63.87 0.16 21%)
64.23 0.14 10%)
51.76
57.64
64.10
62.32
The second experiment is training in parallel two ResNets divided into 2 modules on CIFAR100 (37] We compare Table 2 our results in this setup the two delayed gradient methods DDG B2 and FR [31] from Table 2 in 3I 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.

---

Table 2: Test accuracy of parallel TRGL with 2 modules (average and 95% confidence interval over 3
runs) on CIFARI00, compared to DDG, FR and E2E from Table 2 in [1
Architecture
Parallel VanGL
Parallel TRGL (ours)
DDG
FR
E2E
ResNet- 101
77.31 0.27
77.87 0.44
75.75
76.90
76.52
ResNet-1152
75.40 0.75
76.55 1.90
73.61
76.39
74.80
The third experiment is training in parallel a ResNet-110 10 divided into two, four, eight and sixteen modules on STLI0 [12] We compare in Table our results this setup the recent methods
InfoPro [60] and DGL [ from Table 2 in [60] TRGL largely outperforms the other methods. It also outperforms end-to-end training all but one case (that with 16 modules). With batch size of 64, memory savings of parallel TRGL compared to end-to-end training reach 48% and 58, 58.5% with 8 and 16 modules respectively, with comparable test accuracy. With 4 modules, TRGL training weighs 24% less than and-to-end-traiin
and has a test accuracy that is better by 2 percentage points (see Section 4.2 and Table for a detailed memory usage comparison with InfoPro).
Table 3: Test accuracy of Parallel (Par) TRGL with K modules (average and 95% confidence interval over runs) using a ResNet-110 110 on STL10, compared to DGL, two variants of InfoPro and E2E from Table 2 in [60).
K
Par VanGL
Par TRGL (ours)
DGL
InfoPro S
InfoPro C
E2B
79.85 + 0.93
80.04 = 0.85
75.03 + 1.18
78.98 0.51
79.01 0.64
77.73 + 1.61
77.11 + 2.31
79.72 + 0.81
73.23
: 0.64
78.72 + 0.27
77.27 + 0.40
77.73 + 1,61
75.71 + 0.55
77.82 0.73
72.67
1 0.24
76.40 :::00:0 0.49
74.85 0.52
77.73 + 1.61
16
73.57 + 0.95
77.22 1.20
72.27
+ 0.58
73.95 0.71
73.73 0.48
77.73 + 1.61
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 results with those of InfoPro [60] and InfoProL, a variant of InfoPro proposed in [46]. TRGL. outperforms the other module-wiss wise training methods. It does not outperform end-to-un 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 44
Table 4: Test accuracy of parallel TRGL with 4 modules (average and 95% confidence interval over runs) a Swin-Tiny Transformer, compared InfoPro, InfoProL and E2E from Table in [46). with memory saved compared to E2E as a percentage of E2E memory consumption in red.
Dataset
Parallel VanGI.
Parallel TRGL (ours)
InfoPro
InfoProl.
F2F
STL10
67.00 1.36 (L 55%)
67.92 ::: 1.12 (1 50%)
64.61 (L 38%)
66.89
45%)
72.19
CIFARIO
83.94 + 0.42 (1 33%)
86.48 = 0.54 (L 29%)
83.38 (L 38%)
86.28 (1 45%)
91.37
CIFAR100
69.34 + 0.91 (1 33%)
74.11 = 0.31 (1 29%)
68.36 (L 38%)
73.00 ( 45%)
75.03
Finally, compare method to InfoPro, DGL and Sedona in Table below large scale experiment on ImageNet.
Table 5: Top 1 test accuracy of parallel TRGL with 2 modules ResNet-101 trained on ImageNet, compared to VanGL (baseline vanilla module-wise training), DGL and Sedona from [47 and InfoPro from [60] and end-to-end training.
Dataset
Parallel VanGL
Parallel TRGL (ours)
DGL
Sedona
InfoPro
E2E
ImageNet
78.11
79.41
78.47
79.28
78.15
78.71
4.2 Memory savings and training time
As seen above, parallel TRGL. lighter than end-to-end training by up to almost 60%. The extra memory consumed by regularization compared to parallel VanGL is between 2 and 13% of
7

---

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 m usage [3, and Sedona does not claim to save memory, but rather to speed up training 47 DGL is architecture-wise essentially identical to VanGI. and consumes the same memory.
We compare in Tab the memory consumption of our method to that of InfoPro 160 on a ResNet- 110 on STLI0 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 <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 [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 bello 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.
Table 6: Memory savings using a ResNet- on STL10 split into K modules trained in parallel with batch size of 64, a percentage of the weight of end-to-end training. Average test accuracy over runs is between brackets. Test accuracy of end-to-end training is 77.73%.
Equally deep modules
Equally heavy modules
K
Par VanGL
Par TRGL (ours)
InfoPro
Par VanGL*
Par TRGL* (ours)
InfoPro*
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)
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 comparable to InfoPro which reports time overhead between and 27% compared end-to-end training. See Appendix For details.
4.3 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 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 77. We check first in Table [00 in Appendix G that this architecture works well with parallel module-wise training with 2 modules by comparing it favorably on CIFARIO 10 (37] with methods DGL. 8 InfoPro (0) and DDG [32 that use a ResNet-1 110 with the same number of parameters.
We then train& a 10-block ResNet block-wise on CIFARI00. In Tables I1 and 12 in Appendix G we see that MLS training improves the accuracy of sequential training by 0.8 percentage points when the trainset 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 percentage point. While parallel module-wiss 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-wing training on MNIST
[38] and CIFARIO 37 in Tables
1 and [4 in Appendix G
8

---

The 88% accuracy of sequential training on CIFARIO 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.
4.4 Accuracy after each module
Finally, we verify that our method avoids the stagnation or collapse 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 than vanilla greedy learning early, but surpasses it later modules, indicating that it does avoid early overfitting. On the right, from sequential block-wise training experiments from Table [3] see large decline in performance that the regularization avoids. We see similar patterns in Figure in Appendix G with parallel and MLS block-wise training.
Aags allnd tod baaiig
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 ResNet- 10 with 16 modules on STL10 (Table B). Right: sequential vanilla ( (VanGL, in blue) and regularized (TRGL, in red) block-wiss training 10-block ResNet on 2% of CIFAR10 (Table 13).
5
Limitations
The results in Appendix G show 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 't test their methods 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 [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 memory savings) and reaches almost percentage points. We show Appendix H that our method is not very sensitive to the choice of hyperparameter over large scale.
6 Related work
Layer-wise training was initially considered as a pre-training and initialization method 10 4% and
was recently shown be competitive with end-to-end training [5] Many papers consider using
a different auxiliary loss, instead of or in addition to the classification loss: kernel similarity [4
-ormmationttheory-iipipeddd losses [53 44 60 and biologically plausible losses (53
45 25
Methods (7), PredSim [5), DGL 8. Sedona 47 and InfoPro [60] report the best module- wise
training results. 7 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. method called ResIST [19] that similar
to block-wise training of ResNets randomly assigns ResBlocks to one of up 16 modules that

---

are trained independently and reassembled before another random partition. More of a distributed training method, it is only compared 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 DNI 33
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 address forward locking, which is what (55) do. But they observe training issues with more than 5 modules [I.. This makes them compare unfavorably to module-wis training g The high dimension of the predicted gradient, which scales with the
size of the network, makes (33 [4] challenging in practice. Therefore, despite its simplicity, greedy module-wiss 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 [6]] Transport regularization of ResNets in particular is motivated by the observation that they are naturally biased towards minimally modifying their input [35] 266.. 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 M. 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, [2 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.
use the explicit gradient flow of f f-divergences
7 Conclusion
We introduced a transport regularization for modulc-wise training that theoretically links module-wise training 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-ena 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, for 2=2,, i.e. regularizing with a norm ||-1Ip with p=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. Birkhäuser 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 NeurlPS (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] BELILOVSKY, E., EICKENBERG, M., AND OYALLON, E. Greedy layerwise learning can scale to imagenet. In ICML (2019).
[8] BELILOVSKY, F., 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

---

[10] BENGIO, Y., LAMBLIN, P., POPOVICI, D., AND LAROCHELLE, H. Greedy layer-wise training of deep networks. In NeurlPS (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 sigmoidal function. Mathematics of Control, Signals and Systems (1989), 303-314.
[14] CZARNECKI, W. M., SwIRSZCZ. G., JADERBERG, M., OSINDERO, S., VINYALS, O., AND KAYUKCUOGLU, 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 NeurlPS (2020).
[16] DE BÉZENAC, E., AYED, I., AND GALLINARI, P. Optimal unsupervised domain translation. arXiv (2019).
[17] DOSOVITSKIY, A., BEYER, L., KOLESNIKOV, A., WEISSENBORN, D., ZHAI, X., UN- TERTUINER, T., DEUIHAANI, 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, 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] GF. R., LEE, J. D., AND MA, T. Learning one-hiddennllyer neural networks with landscape design. In International Conference on Learning Representations (2018).
[23] GIANAZZA, U., GOBBINO, M., AND SAVARÉ, G. Evolution problems and minimizing movements. Atti della Accademia Nazionale dei Lincei. Classe di Scienze Fisiche, Matematiche Naturali. Rendiconti Lincei. Matematica 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 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 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., 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. NeurlPS (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. ICMI. (2017).
11

---

[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 BÉZENAC, 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. wom.leecun.come aacum.com/exdb/nniss (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., IMŞEKLI, 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 KLEUN, W. B. The hsic bottleneck: Deep learning without back-propagation. In AAA/ (2020).
[42] MANDAR KULKARNI, S. K. Layer-wise training of deep networks using kemel similarity. In DLPR workshop, ICPR (2016).
[43] MARQUEZ, E. S., HARE, S., AND NIRANIAN, M. Deep cascade learning. IEEE Transactions on Neural Networks and Learning Systems (2018).
[44] NGUYEN, T., AND CHO1, J. Layer-wise learning of stochastic neural networks with information bottleneck. Entropy 21 (2019).
[45]
NøKLAND, 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., HAII, 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 Patlern Recognition, CVPR 2018, Sall 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. Birkhäuser, 2015.
[50] SANTAMBROGIO, F. Euclidean, metric, and wasserstein gradient flows: 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 Learning Representations (ICLR 2014) (2014). [53] SINDY LÖWF, PETER O' C'ONNNOR, B. V. Putting an end to end-to-end: Gradient-isslated learning of representations. In NewrIPS (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 -xiliary-variablle networks. arXiv (2021).
[56] TAN, M., AND LE, Q. EfficientNet:
Rethinking model scaling for convolutional neural networks. In Proceedings of the 36th International Conference 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).
12

---

[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: 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., 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 giant leap. IJCA (2019).
[65] ZHANG, J., ZHANG, X., MA, K., GUPTA, R., SALTZ, J., VAKALOPOULOU, M., AND SAMARAS, D. Gigapixel whole-slide images classification using locally supervised learning. In
Medical Image Computing and Computer Assisted Intervention MICCAI 2022 (Cham, 2022),
Springer Nature Switzerland, pp. 192-201.
13

---

A Background on optimal transport
The Wasserstein space (2( with Q a convex and compact subset of E is the space P(l2) of probability measures over St, equipped with the distance W2 given by the solution to the optimal transport problem
min W *ETT(..8)) Jxx || y1l dy(2,
(12)
where IT .9 is the set of probability distribution over f x f with first marginal and second
marginal 8. i.e. 1]) = {r E P(Ω x (2) Øigy = a, T2;7 = [8} where #i (2., y) x and T2(2, 3) = y. The optimal transport problem can be seen as looking for 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 W2 can be shown to be geodesic distance (see for example [49 59). If is absolutely continuous and an is a-negligible then the problem in (12) (called the Kantorovich problem) has unique solution and is equivalent to the Monge problem, i.c.
/ W (a. B) min Tst. Taa= Jn IT((x) ==| da(xx)
(13)
and this problem has unique solution T linked to the solution y" of (12 through y" (id, T )aa. Another equivalent formulation of the optimal transport problem in this setting the dynamical formulation [9. Here, instead of directly pushing samples of to 3 using T, we can equivalently displace mass, according to continuous flow with velocity 14 [R RRd This implies that the density (n at time satisfies the continuity equation 0,a, + V (aaw) = 0, assuming that initial and final conditions are given by ag = a and æ1 = B respectively. In this case, the optimal displacement is the one that minimizes the total action caused by
W2 3) minn I'" |1v (1l:1200)) dt
(14)
s.t. 0,a, + V (a.t) = 0, og = a, 01 =
Instead of describing the density's s evolution through the continuity equation, we can describe the paths of taken by particles at position from when displaced along the flow u. Here of is the position at time t of the particle that was at x ev a at time 0. The continuity equation is then equivalent to O, Ur (df)). See chapters and 5 of [49] for details. Rewriting the conditions as necessary, Problem [14 becomes
(2 3) = min / ||0 (((2((00)))) df v
(15)
s.t. 0, (df), do = id, (0i);a
and the optimal transport map T that solves (3) is in fact (z) of for & that solves the continuity equation together with the optimal v² 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 a and B be absolutely continuous measures on &Rd and T the optimal transport map between and 3 for the Euclidean cost. Suppose there are bounded open sets X and Y, such that the density of a (respectively of (3 is null on X (respectively Y*) and bounded away from zero and infinity on X (respectively Y).
Then there exists two relatively closed sets of null measure C X and B C Y, such that T is q Hölder continuous from X A to Y \ B, Le.. V x, E X A we have
[[T(x) - T())l < Cll yll" for constants n,C C²
B Background on gradient flows
We follow [50 2 for this background on gradient flows Given a function L: Rdd - R and an initial point 2o E Rdd. gradient flow is curve [0, co|+ Rd that solves the Cauchy problem
x' (t) v[(((t))
(16)
x00)= = 2o
14

---

A solution exists and is unique if V Lipschitz or L is convex. Given > 0 and zij = 20 define sequence (x through the minimizing movement scheme:
+ E arg min ((x) + 3/ 1
(17)
#ER
2T
C2||2|12 guarantees existence a solution of 177 for small enough. L X-convex meets these conditions and also provides uniqueness of the solution
lower semi-continous and ((x) 2 C1
because of strict convexity of the target. See 49 50 2].
We interpret the point oi as the value of a curve at time kr. We can then construct a curve x² as the piecewise constant interpolation of the points zt.- We can also construct curve #' as the affine interpolation of the points z/.-
If ((ro)) < oo and inf -c then (x*) and (x*) converge uniformly the same curve x goes to zero (up to extracting a subsequence). If is C²,, then the limit curve æ is a solution of (16) (i.e. a gradient flow of L). the subdifferential of C, i.. satisfies x' (t) E -d[(x(t)) for almost every t.
differentiable then x solution of the problem defined using
If L is X-convex with À 0, then the solution (16) converges exponentially the unique minimizer of L (which exists by coercivity). So taking - 0 and k: + co, we tend towards the minimizer of C.
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 s distance. In the (geodesic) metric space (2( with Ω convex and compact, for L (2( - R U {oo} lower semi-continuous for the weak convergence of measures in duality with C(2) (equivalent to lower semi-continuous with respect to the distance W2) and P Po E P(2), the minimizing movement scheme (17) becomes
Ph++ E arg min
C(p)
27 (2.
(18)
PEP(O)
This problem has solution because the objective is lower semi-continuous and the minimization is over ((22) which is compact by Banach-Alaoglu.
We can construct piecewise constant interpolation between the measures P or a geodesic inter- polation where we travel along a geodesic between PE and Pk+1 in W2( constructed using the optimal transport map between these measures. Again, ((x)) interpolations converge uniformly to a limit curve ý as т goes to zero. Under further conditions on C, mainly X-geodesic convexity (i.e. X-convexity along geodesics) for À >> we can prove stability and convergence of p(t) to a minimizer of C as t + co, see 49 50
and inf C > -oo then both
C Proof of Proposition 2.3
Proof. Take a minimizing sequence (F,, T.), i.e. such that C(F,, T,) - min C, where 2 0 is the
target function in ( and denote 3: -T:Pl- k- Then by compacity F; + F* and 3; - 3 in duality
with (S2) by Banach-Alaoglu. There exists T optimal transport map between P% and f Then
C(F*,T) S lim C(F, T.) min by continuity and because
IT æ=1² dx)(z)
(2
lim (P
lim [T](()) - dp)
5
W2 metrizes weak convergence of measures. We take (FE,Tİ) - (F*,T*), is also immediate
that for any minimizing pair, the transport map has to be optimal. Taking minimizing sequence
F. o and the corresponding induced maps T; we get the same result for Problem (1) Problems
3 and [I] are equivalent by the equivalence between Problems (3) and [5)
D Varying the regularization weight
The discussion in Section 2.2 suggests taking a fixed weight 7 for the transport cost that is as small as possible. However, instead of using fixed might want to vary it along the depth & to further
15

---

constrain with a smaller тhe the earlier modules to avoid that they overfit or the later modules to maintain the accuracy 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 Th across modules and SGID iterations we use scheme inspired by the method of multipliers to solve Problems 2 and (TOp. To simplify the notations, we will instead consider the weight ==, given to the loss. We denote Oz the parameters of both Tj and F& at SGD iteration i. We also denote L(0,x)) x)) and ( x)) respectively the loss and the transport regularization as functions of parameters 0 and data z. We now increase the weight Au, of the loss every iterations of SGD by a value that is proportional to the current loss. Given increase factor h>0, initial parameters 0x,1, initial weights 20, learning rates (7:) and batches (zz). we apply for module lo and i21
[0R,4+1 = -O2, TI V L(OK, x4) + W (Ox,s, (.)) (1Liit Au.i (1.(O.,+11, Bi+1) if mod = 0 else XR,i
The weights As. will vary along modules & because they will evolve differently with iterations i for each k. They will increase more slowly with i for larger because deeper modules will have smaller loss. This method can be seen as method of multipliers for the problem of minimizing the transport under the constraint of zero loss. Therefore
immediate by slightly adapting the proof of Proposition 2.3 or from [36] that we are still solving 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 X1 = AR,1 for all modules so that this method requires choosing three nyper-parameter (h, 8 and X1). In practice (see Section 4.1 and Appendix [F). it works best only one experiment. Simply manually finding a value of that is multiplied by 2 for the second half of the network works best in all the other experiments.
E Implementation details
We 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 Resnet-152 on TinylmageNet in Table D takes about 36 hours. Training a Resnet- 152 on CIFAR 100 in Table I takes about 11 hours. Training a ResNet-110 on STL10 in Table 3 takes about 3 hours. Training Swin-Tiny Transformer in Table take between 40 minutes and hour.
For sequential and multi-lap sequential training, we use SGD with learning rate of 0.007. With the exception of the Swin Transformer in Table we use SGD for parallel training with learning rate of 0.003 in all Tables but Table Suble/whee where the learning rate is 0.002. For the Swin Transformer in Table 4 we use the Adam W 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 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
For experiments in Section 4.1 use batch size of 256, orthogonal initialization [51) with gain of 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
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 I.
we use 500000 for the first two modules and then double it for the last two modules
for TRGL. In Table we use AR,1 = 1, h = 1 and s 50 for TRGL. In Table 3 we use r = 0.5 and double it at the midpoint, expect for the first line where 7 50.
F Memory savings and training time
We compare in Table 2 the memory consumption of our method to that of InfoPro [60 on a ResNet- 1 split into modules trained in parallel on STL.00 with a batch size of 64 (so the same setting as in 3 Section 4.1 We report in Table the memory saved as percentage of the 6230 MiB memory required by end-to-end training with the same batch size. VanGI. refers to our architecture trained without the regularization. InfoPro (60) also propose to split the network into K modules that
16

---

have the same weight but not necessarily the same number of layers. They only implement this for K<4 modules. When the modules are even weight and not in depth, we call the training methods VanGI.* TRGL.* and InfoPro*. In practice, this leads to shallower early modules which slightly hurts performance according to [47]. We verify this in Table (to compared with Table 3 in Section 4.1p. However, TRGL* still outperforms InfoPro and end-to-end training, and it leads to even bigger memory savings. We see in Table that 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* lighter than InfoPro*, with better accuracy. We also that the added memory cost of the regularization compared to vanilla greedy learning is small. However, parallel module-wise 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 time overhead between and 27% compared to end-to-end training.
Table 7: Memory savings using ResNet- 110 on STL10 split into K modules trained in parallel with a batch size of 64, as percentage of the weight of end-to-end training. Average test accuracy over runs is between brackets. Test accuracy of end-to-end training is 77.73%.
Equally deep modules
Equally heavy modules
K
Par VanGL
Par TRGL (ours)
InfoPro
Par VanGL*
Par TRGL* (ours)
InfoPro
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)
Table 8: Test accuracy of parallel (Par) TRGL* with K modules (average and 95% confidence interval over runs) on a ResNet- 110 trained on STL10, compared InfoPro* and E2E training from Table 3 in [60]
K
Par VanGL.*
Par TRGL* (ours)
InfoPro"
79.05 + 1.33
79.47 1.36
79.05 0.57
4
77.14 1: 1.23
78.94 1 1.13
78.78 + 0.72
Note that methods DDG [32] and FR (31], being delayed gradient methods and not module-wise training methods, do 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 [ is architecture-wise essentially identical to VanGL and consumes the same memory.
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 modules on CIFAR 100, sequentially and multi-lap sequentially. Results are in T 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 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).
Table 9: Test accuracy of sequential (Seq) and multi-lap sequential (MLS) TRGL and VanGL with 2 modules on CIFARI00 using 101 (average of 2 runs).
Seq VanGI.
Seq TRGL.
MLS VanGI.
MLS TRGL.
73.31
73.61
74.34
74.78
17

---

We now focus on full block-wise training, meaning that each module is a single ResBlock, mostly sequentially, We propose here 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 network is divided in K modules of M ResBlocks each, we call the network K-M ResNet. These wider shallower architectures are well-adapted to layer-wise training as seen in 77. We check in Table 10 that this architecture works well with parallel module-wite training by comparing favorably on CIFARIO (37) 2-7 ResNet with DGL, InfoPro (160)) and DDG [32]. The 2-7 ResNet has 45 millions parameters, which is about the same the ResNet- 1 10 divided in two used by the other methods, and performs better when trained in parallel.
Table 10: Average test accuracy and 95% confidence interval of 2-7 ResNet over 10 runs on CIFARIO with parallel TRGL and VanGL, compared to DGL and DDG from [8] and InfoPro from 60 that split ResNet-110 module-wise moduue--wis--pparalle--tt parallel-trained modules.
Parallel VanGL (ours)
Parallel TRGL. (ours)
DGL
DDG
InfoPro
94.01 + .17
94.05 = .18
93.50
93.41
93.58
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 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 05k5 10. block 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+2kk 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 I1 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 MI.S 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 [2 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.
Table 11: Average highest test accuracy and 95% confidence interval of 10-1 ResNet over 10 runs on CIFAR 100 with different train sizes and sequential (Seq), multi-lap sequential (MLS) and parallel (Par) TRGL and VanGL, compared to E2E.
Train size
Seq VanGI.
Seq TRGL.
ML. VanGL.
ML. TRGL.
E2E
50000
68.74 + 0.45
68.79 + 0.56
69.48 = 0.53
69.95 : 0.50
75,85 + 0.70
25000
60.48 + 0.15
60.59 = 0.14
61.33 = 0.23
61.71 = 0.32
65.36 + 0.31
12,00
51.64 + 0.33
51.74 0.26
51.30 0.22
51.89 0.30
52.39 0.97
5000
36.37 + 0.33
36.40 0.40
33.68 + 0.48
34.61 + 0.59
36.38 + 0.31
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.
Train size
Par VanGL
Par TRGL
E2E
50000
72.59
+ 0.40
72.63 : 0.40
75.85 + 0.70
25000
64.84
0.19
65.01 = 0.27
65.36 0.31
12500
55.13
0.24
55.40 0.35
52.39
1 0.97
5000
39.45
0.23
40.36 = 0.23
36.38 + 0.31
18

---

We report further results of block-wise training on CIFARIO 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 the regularization as the training set gets smaller, gaining to 6 percentage points.
Table 13: Average last block test accuracy and 95% confidence interval of 10-1 ResNet 10 on CIFAR 10 with different train sizes and sequential TRGL and VanGL, compared to E2E.
Train size
Seq VanGL
Seq TRGL
E2E
50000
88.02 = .18
88.20 + .24
91.88 = .18
25000
83,95 + .13
84.28
.22
88.75
-27
10000
76.00
77.18
82.61
.35
5000
67.74 1 .49
69.67
.44
73.93 1 .67
1000
45.67 = .88
51.34 + .90
50.63 = .88
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 MNIST with different train sizes and parallel TRGL and VanGL, compared to E2E.
Train size
Par VanGL
Par TRGL
E2E
60000
99.07 + _04
99.08 + .04
99.30 1: .03
30000
98.90 1 .05
98.93 1 .06
99.22 + .03
12000
98.52 1
98.59 .06
1 98.96
6000
98.05 = .09
98.16 = .07
98.62 = .66
1500
96.34 .112
96.91 = .07
97.19 = .08
1200
95,80 .12
96.58 = .09
96.88 = .99
600
91.35 .99
95.16 -15
95.30 + .17
300
89.81 +
92.86 .24
92.87 + 28
150
81.84 + 1.22
87.48 + .42
87.82 + .59
The 88% accuracy of sequential training CIFARIO in Table 13 the same for sequential training table of (7), which is the best method for layer-wise sequential training available, with VGG networks of comparable depth and width.
100 ecconcc
Uaiills muttitep itariisaie Dequlariznd mati- trockontes trantin
block
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, 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.
19

---

H Sensitivity to hyper-parameters
We show in Figure below that TRGL still performs better than VanGL (in the same setting in Table in in Section for values of from 0.03 to 100 and is still roughly equivalent to it for values up to 5000.
TRGL
Test accuracy in %
Figure 4: Average test accuracy over 5 runs of parallel TRGL using ResNetl 10 STL10 with 16 modules with different values of r (in red), and VanGL (blue line).
I Broader impact
Less memory usage has a positive environmental impact and allows organizations with less resources use deep learning.
20


Reviews:
[{"review_id": "qlewAhYZQ1", "reviewer": "Reviewer_Ad5i", "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.", "questions": "See weakness", "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.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n\n(2) The authors provide theoretical analysis, proving that TRGL leads to more regular and stable greedy modules that progressively minimize the loss.\n\n(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?\n\n(2) what is VanGL? is it just a module-wise training without the regularization terms?\n\n(3) In my opinion, the performance gain is marginal compared  to vanilla VanGL.\n\n(4) The principle of finding Tau is vague.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer Ad5i,\n\nThank you for your valuable review. We answer your remarks below.\n\n**Weaknesses**\n\n**(1) Training time**\n\nOf course like any other method that adds a term to the local loss (so most other methods), training time is slightly increased in favor of saving memory, which is the real hard constraint when training on a small device. We do discuss this in Section 4.2 but probably not prominently enough and quantitatively compare to other methods. Starting line 259: \n\n“Vanilla parallel module-wise (denoted VanGL in the paper) training does slightly slow down training. Epoch time increases by 6% with 2 modules and by 16% with 16 modules compared to end-to-end training. By adding our regularization, TRGL is only slower than VanGL by 2% for any number of modules. This is comparable to InfoPro which reports a time overhead between 1 and 27% compared to end-to-end training. “\n\nNote that InfoPro is the most comparable method to ours in this regard as it also saves memory at the price of slowing down training a little bit, while having performances that are close to or better than end-to-end training. The real limitation for training on mobile devices is memory since the increase in training time is small as stated above: 2% compared to VanGL which has already been used for on-device training in [50,51] as indicated in the paper. \n\n\n**(2) Meaning of VanGL**\n\nYes this is as you say. “VanGL” is defined on line 209: \"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\"\n\n**(3) Ablation**\n\nThe gains are not marginal on parallel training in Section 4.1. They are higher than 1 percentage point in 9 out of the 13 experiments, and higher than 2 percentage points in 6 out of the 13 experiments. This is confirmed by experiments on the larger ImageNet dataset that we added in the general answer to all reviewers above, where the improvement in Top 1 accuracy is 1.3 percentage points.\n\nThey are indeed marginal on sequential training in the full data regime, but quite big in the small data regime (up to 6 percentage points gained, Tables 10 and 12 in Appendix G). \n\nCompared to other methods, our experiments cover more cases (parallel and sequential with many or few modules) and show that our method is more reliable as at least it never hurts performance. This is already discussed in Section 5 where we say starting line 294: \"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 fail in different settings from the original one in subsequent papers (e.g. delayed gradients methods when the number of modules increases [25] and PredSim in [40]). Also, for parallel training in Section 4.1, the improvement from the regularization compared to VanGL increases with the number of modules (so as the memory savings increase and module-wise training becomes more useful).\"\n \n**(4) Choosing hyper-parameter $\\tau$**\n\nWe experimented with two strategies for choosing the parameter** $\\tau$. A method of multipliers is used to adaptively change $\\tau$ during training differently for each module and is detailed in Section 3.3 and Appendix D. A reference is made to [29] which also uses this method to adaptively change a regularization weight during training. This is more principled than other methods that also add a term to the local loss. For instance, InfoPro [53] have two hyper-parameters per module and they simply assume that they change linearly from the first to the last module and then choose the 4 needed values manually from a fixed set.\n\nBeside, we also experimented with a simple heuristic that works well in practice and involves setting manually a single hyper-parameter: we use a value $\\tau$ for the first $K/2$ modules and double it (i.e use $2 \\tau$) for the last $K/2$ modules. So only $\\tau$ has to be chosen. Again this is mentioned in Section 3.3.\n\n**Limitations**\n\nThe real limitation for training on mobile devices is memory since the increase in training time is small as stated above: 2% compared to VanGL which has already been used for on-device training [50,51] as indicated in the paper. \n"}, {"review_id": "VNHsMxzHLw", "reviewer": "Reviewer_U1e3", "summary": "The paper proposed a new regularization for module-wise training via the distance of the input and output of the module.", "questions": "1. Why is the memory saving of the proposed method often less than other methods? Where this overhead comes from?\n2. Some notations are not introduced properly. For example, the `#' notation, and what is the difference of $T_k$ and $T_k^{\\tau}$.", "limitations": "Limitations are provided.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The proposed regularization is quite straightforward and easy to apply.\n2. The paper connects the proposed regularization with optimal transport via theoretical analysis.\n3. 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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer U1e3,\n\nThank you for your valuable review, we answer your remarks below.\n\n**Weaknesses**\n\n**1. Convergence**\n\nIndeed this is an important question and this will be further discussed in the final version. In Section 2.2, we have assumed that the training of the individual modules converges.\n\nIndeed, the training of shallow networks is known to converge [A,B,C,D,E]. Papers [5,6,24,F,G] show the convergence of vanilla module-wise (both parallel and sequential) training by chaining existing convergence results for shallow modules. We have not considered the impact of the regularization on this, but it is a simple quadratic regularization that makes the target more convex and [29] show the beneficial effects of this transport regularization on training a network end-to-end. \n\nOn the theoretical side, assuming therefore that the training of each shallow module converges with the number of epochs, the assumptions under which we have convergence as the regularization weight parameter $\\tau$ goes to zero and the number of modules $k$ increases are discussed in Section 2.2 (lines 130 to 136). Convergence as the number of modules increases is the central point in module-wise training as we precisely want the modules to build upon each other to solve the task, and this is exactly what we demonstrate.\n\nIn practice, experiments reported on Fig 2 and Fig. 3 show the importance of the regularization term for the convergence to a solution as the number of modules increases.\n\n[A] *Understanding Deep Neural Networks with Rectified Linear Units*, Arora et al., ICLR 2018\n\n[B] *Breaking the Curse of Dimensionality with Convex Neural Networks*, Bach, JMLR 2017\n\n[C] *Beating the Perils of Non-Convexity: Guaranteed Training of Neural Networks using Tensor Methods*, Janzamin et al, arXiv 2016\n\n[D] *Learning One-hidden-layer Neural Networks with Landscape Design*, Ge et al., ICLR 2018\n\n[E] *Improved Learning of One-hidden-layer Convolutional Neural Networks with Overlaps*, Du et al., arXiv 2018\n\n[F] *A Provably Correct Algorithm for Deep Learning that Actually Works*, Malach et al., arXiv 2018\n\n[G] *Provable Bounds for Learning Some Deep Representations*, Arora et al., ICML 2014\n    \n\n\n**2. Setting the regularization strength**\n\nIt is true that in [28] the early modules are *less biased* towards small displacement values, but regularizing them the same as for deeper layers was beneficial in [29]. \n\nWe discuss in Section 3.3 and Appendix D the possible need to have a different regularization strength for each module (i.e a different weight $\\tau_k$ is given to module $k$), and use an algorithm inspired by the method of multipliers that was also used in [29]  to adjust these values $\\tau_k$. \n\nIndeed as you indicate \"the selection of the regularization strength becomes more complex when increasing the number of modules $K$\", but this is true for all module-wise training methods that add a term to the local loss. The method of multipliers algorithm we use allows us to start all $\\tau_k$ from the same initial value and then let them change at different rates during training (Section 3.3 and Appendix D). **This way the number of hyper-parameters required by our method is 3 and does not increase with the number of modules $K$**. \n\nThis is more principled than other methods that also add a term to the local loss. For instance, InfoPro [53] have two hyper-parameters per module and they simply assume that they change linearly from the first to the last module and then choose the 4 needed values manually from a fixed set.\n\nAdditionally, while observing the behaviour of this algorithm, we noticed that a simple pattern for the values $\\tau_k$ works well in most cases and deduced a simple heuristic: use a value $\\tau$ for the first $K/2$ modules and double it (i.e use $2 \\tau$) for the last $K/2$ modules. A **single hyper-parameter $\\tau$** has then to be fixed and this can be done manually or through any search or cross validation method.\n\n**3. Large scale experiments**\n\nSince submitting the paper we have run experiments on ImageNet (ResNet-101 split into 2 modules trained in parallel. Results are below and in the additional PDF file. This validates that the conclusion drawn on the smaller datasets still holds for large datasets such as ImageNet. We did not experiment with more than 2 modules on ImageNet due to resource limitations.\n\n| Dataset  | Parallel VanGL | Parallel TRGL (ours) | DGL   | Sedona | InfoPro | E2E   |\n|----------|----------------|----------------------|-------|--------|---------|-------|\n| ImageNet | 78.11          | **79.41**            | 78.47 | 79.28  | 78.15   | 78.71 |\n\n\n**Questions**\n\n**1. Memory savings**\n\nThere might be a misunderstanding here. The memory savings are not less than the other methods. They are most often higher. InfoPro saves more memory in only 2 out of 7 cases (see Tables 4 and 5). InfoProL saves more memory in 2 out of three cases but many of the other methods (Sedona, DDG, FR) don't claim to save memory at all compared to end-to-end training. \n Please note that VanGL is simply vanilla module-wise training without any added regularization or auxiliary network, so it is normal that it uses less memory than all other methods.\n\n**2. Notations**\n\nSorry about that, we will add the definition of the pushforward measure $\\#$. $T_k$ is simply the module trained in vanilla module-wise training (VanGL), so there is no dependence on $\\tau$ which is the weight given to the regularization in TRGL.\n"}, {"review_id": "9X53I1VBvC", "reviewer": "Reviewer_X5WU", "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.", "questions": "N/A", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. Establishing a connection between neural network training and optimal transport theory is not only intellectually stimulating, but it also presents intriguing theoretical insights.\n\n3. 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.\n\n\n", "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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer X5WU,\n\nThank you for your valuable review. We answer your remarks below.\n\n**Weaknesses**\n\n**1. Convergence**\n\nIndeed this is an important question and this will be further discussed in the final version. In Section 2.2, we have assumed that the training of the individual modules converges.\n\nIndeed, the training of shallow networks is known to converge [A,B,C,D,E]. Papers [5,6,24,F,G] show the convergence of vanilla module-wise (both parallel and sequential) training by chaining existing convergence results for shallow modules. We have not considered the impact of the regularization on this, but it is a simple quadratic regularization that makes the target more convex and [29] show the beneficial effects of this transport regularization on training a network end-to-end. \n\nAssuming therefore that the training of each shallow module converges with the number of epochs, the assumptions under which we have convergence as the regularization weight parameter $\\tau$ goes to zero and the number of modules $k$ increases are discussed in Section 2.2 (lines 130 to 136). Convergence as the number of modules increases is the central point in module-wise training as we precisely want the modules to build upon each other to solve the task, and this is exactly what we demonstrate.\n\n[A] *Understanding Deep Neural Networks with Rectified Linear Units*, Arora et al., ICLR 2018\n\n[B] *Breaking the Curse of Dimensionality with Convex Neural Networks*, Bach, JMLR 2017\n\n[C] *Beating the Perils of Non-Convexity: Guaranteed Training of Neural Networks using Tensor Methods*, Janzamin et al, arXiv 2016\n\n[D] *Learning One-hidden-layer Neural Networks with Landscape Design*, Ge et al., ICLR 2018\n\n[E] *Improved Learning of One-hidden-layer Convolutional Neural Networks with Overlaps*, Du et al., arXiv 2018\n\n[F] *A Provably Correct Algorithm for Deep Learning that Actually Works*, Malach et al., arXiv 2018\n\n[G] *Provable Bounds for Learning Some Deep Representations*, Arora et al., ICML 2014\n    \n\n**2. Impact of variations in the regularization penalty**\n\nWe already provide an exploration of the impact of variations in the regularization penalty in Figure 4 in Appendix H (this is mentioned in Section 5) on the STL10 dataset. We found that \"TRGL performs better than VanGL for values of $\\tau$ from 0.03 to 100 and is then roughly equivalent to it for values up to 5000.\" \n\nWe have run a similar experiment using the Swin-Tiny Transformer split in 4 modules trained in parallel on CIFAR100 (line 4 in Table 4), with similar results. The figure is in the additional PDF joined to the general answer above.\n\n"}, {"review_id": "fnZfBXHbde", "reviewer": "Reviewer_N1H7", "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.\n\n-------After Rebuttal-------\n\nI thank the authors for the answers to my comments and questions. Most of my concerns are properly addressed. Therefore I raise my score.\n\nWhile 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.", "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.", "limitations": "See weaknesses and questions.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Overall, the idea is new and the proposed method is theoretically sound.\n2. 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.\n\nAs 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 \ncapacity of the first few layers. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer N1H7,\n\nThank you for your valuable review. We answer your remarks below.\n\n**Weaknesses**\n\n**1. Evidence for early overfitting**\n\nThis is indeed an interesting question. Our experiments in Figures 2 and 3 demonstrate that vanilla module-wise training performs very well in the early layers but stagnates and gets overtaken later for the datasets considered in the paper. \n\nMost importantly, we cite 4 papers [36,5,53,40] that make the same observation on line 33. The arguments developed in each paper are detailed below. We therefore considered that this was well-established in the literature. We propose to add this discussion of the evidence in the literature for this early overfitting phenomenon to the paper.\n\n[36] compare the accuracy after each layer to that of end-to-end training and find that layer-wise training overfits and outperforms E2E in early layers, but stagnates later while E2E improves dramatically. \n\n[53] explore this in more detail. They find that \"greedy SL (module-wise training) contributes to dramatically more discriminative features with the first one (or few) local module(s), but is only able to slightly improve the performance with all the consequent modules. In contrast, the E2E learned network progressively boosts the linear separability of features throughout the whole network with even more significant effects in the later layers, surpassing greedy SL eventually\". [53] also show that module-wise training detroys a larger amount of mutual information between the learned features and the inputs (and between the learned features and the outputs) compared to end-to-end training in the early modules.\n\nAlso, [40] find that the first module \"may learn only too coarse features to classify complex images, which may be less meaningful representations for upper layers\". They solve this by making the first module deeper, which reduces the memory savings of the method, and [5] solve the problem by making the auxiliary network deeper. \n\n**2. Large scale experiments**\n\nSince submitting the paper we have run experiments on ImageNet (ResNet-101 split into 2 modules trained in parallel). Results are below and in the additional PDF file. This validates that the conclusion drawn on the smaller datasets still holds for large datasets such as ImageNet. We did not experiment with more than 2 modules on ImageNet due to resource limitations.\n\n| Dataset  | Parallel VanGL | Parallel TRGL (ours) | DGL   | Sedona | InfoPro | E2E   |\n|----------|----------------|----------------------|-------|--------|---------|-------|\n| ImageNet | 78.11          | **79.41**            | 78.47 | 79.28  | 78.15   | 78.71 |\n\n\n\n**Questions**\n\n**1. Training time**\n\nThis was included in the paper but probably not prominently enough, in section 4.2 starting line 259:\n\nVanilla parallel module-wise (denoted VanGL in the paper) training does slightly increase  training time. Epoch time increases by 6% with 2 modules and by 16% with 16 modules compared to end-to-end training. By adding our regularization, TRGL is only slower than VanGL by 2% for any number of modules. This is comparable to InfoPro baseline which reports a time overhead between 1 and 27% compared to end-to-end training. \n\nNote that InfoPro is the most comparable method to ours in this regard as it also saves memory at the price of slowing down training a little bit, while having performances that are close to or better than end-to-end training.\n"}]

Author rebuttal:
Thanks to all the reviewers for their valuable insights. We have addressed in the individual answers the points raised by the reviewers. Since some comments are shared by different reviewers, we summarise here the main answers:

**1. Experiments on ImageNet (Reviewers U1e3 and N1H7)**

The reviewers asked for experiments on larger datasets. Since submitting the paper we have run experiments on ImageNet (ResNet-101 split into 2 modules trained in parallel). Results are below and in the additional PDF file. This validates that the conclusion drawn on the smaller datasets still holds for large datasets such as ImageNet. We did not experiment with more than 2 modules on ImageNet due to resource limitations.

| Dataset  | Parallel VanGL | Parallel TRGL (ours) | DGL   | Sedona | InfoPro | E2E   |
|----------|----------------|----------------------|-------|--------|---------|-------|
| ImageNet | 78.11          | **79.41**            | 78.47 | 79.28  | 78.15   | 78.71 |

**2. Training time (Reviewers Ad5i and N1H7)**

The reviewers asked about the increase in wall-time training cost caused by our regularization. This was included in the paper but not prominently enough. In section 4.2 we say:

Vanilla parallel module-wise (denoted VanGL in the paper) training does slightly increase training time. Epoch time increases by 6% with 2 modules and by 16% with 16 modules compared to end-to-end training. By adding our regularization, TRGL is only slower than VanGL by 2% for any number of modules. This is comparable to InfoPro baseline which reports a time overhead between 1 and 27% compared to end-to-end training. 

Note that InfoPro is the most comparable method to ours in this regard as it also saves memory at the price of slowing down training a little bit, while having performances that are close to or better than end-to-end training.


**3. Convergence (Reviewers U1e3 and X5WU)**

Indeed this is an important question and this will be further discussed in the final version. In Section 2.2, we have assumed that the training of the individual modules converges.

Indeed, the training of shallow networks is known to converge [A,B,C,D,E]. Papers [5,6,24,F,G] show the convergence of vanilla module-wise (both parallel and sequential) training by chaining existing convergence results for shallow modules. We have not considered the impact of the regularization on this, but it is a simple quadratic regularization that makes the target more convex and [29] show the beneficial effects of this transport regularization on training a network end-to-end. 

Assuming therefore that the training of each shallow module converges with the number of epochs, the assumptions under which we have convergence as the regularization weight parameter $\tau$ goes to zero and the number of modules (denoted $k$ ln 131) increases are discussed in Section 2.2 (lines 130 to 136). Convergence as the number of modules increases is the central point in module-wise training as we precisely want the modules to build upon each other to solve the task, and this is exactly what we demonstrate.

In practice, experiments reported on Fig 2 and Fig. 3 show the importance of the regularization term for the convergence to a solution as the number of modules increases.

[A] *Understanding Deep Neural Networks with Rectified Linear Units*, Arora et al., ICLR 2018

[B] *Breaking the Curse of Dimensionality with Convex Neural Networks*, Bach, JMLR 2017

[C] *Beating the Perils of Non-Convexity: Guaranteed Training of Neural Networks using Tensor Methods*, Janzamin et al, arXiv 2016

[D] *Learning One-hidden-layer Neural Networks with Landscape Design*, Ge et al., ICLR 2018

[E] *Improved Learning of One-hidden-layer Convolutional Neural Networks with Overlaps*, Du et al., arXiv 2018

[F] *A Provably Correct Algorithm for Deep Learning that Actually Works*, Malach et al., arXiv 2018

[G] *Provable Bounds for Learning Some Deep Representations*, Arora et al., ICML 2014
    

**4. Choosing hyper-parameter $\tau$ (Reviewers Ad5i and U1e3)**

We discuss in Section 3.3 and Appendix D the possible need to have a different regularization strength for each module (i.e a different weight $\tau_k$ is given to module $k$), and use an algorithm inspired by the method of multipliers (that was also used in [29]) to adjust these values $\tau_k$. 

The method of multipliers algorithm we use allows us to start all values $\tau_k$ from the same initial value and then let them change at different rates during training. **This way the number of hyper-parameters required by our method is 3 and does not increase with the number of modules $K$**. 

This is more principled than other methods that also add a term to the local loss. For instance, InfoPro [53] have two hyper-parameters per module and they simply assume that they change linearly from the first to the last module and then choose the 4 needed values manually from a fixed set.

Additionally, while observing the behaviour of this algorithm, we noticed that a simple pattern for the values $\tau_k$ works well in most cases and deduced a simple heuristic: use a value $\tau$ for the first $K/2$ modules and double it (i.e use $2 \tau$) for the last $K/2$ modules. A **single hyper-parameter $\tau$** has then to be fixed and this can be done manually or through any search or cross validation method.

**5. Evidence for early overfitting (Reviewer N1H7)**

Besides our experiments in Figures 2 and 3, we cite 4 papers [36,5,53,40] on line 33 that make the same observation. The arguments developed in each paper are detailed below in the individual answer to Reviewer N1H7. We therefore considered that this was well-established in the literature. We propose to add the discussion in the individual answer to Reviewer N1H7 of the evidence in the literature for this phenomenon to the paper.

**6. Further sensitivity to hyper-parameter experiment (Reviewer X5WU)**

See Appendix H and a new experiment in the additional PDF





Meta-review: The submission proposes a training method for greedy layer-/module- wise training of neural networks. In contrast to standard (full-network) backpropagation, which has the requirement of keeping the full network and computation graph in the memory during the training, module-wise training relies on chopping up the network into modules that can be trained more independently. This is useful in resource-bounded settings such as on-device training, but the approximation comes with certain endemic issues. The goal of the submissions is to alleviate some of them.  

All of the reviewers recommend the submission for acceptance. The rebuttal was particularly useful and changed the opinion of several of the reviewers toward acceptance. The newly reported results on larger-scale experiments were also found helpful. The authors are recommended to incorporate the comments raised during the review process in preparing their camera-ready.
Decision: Accept (poster)