SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To compare with other methods, we focus first on parallel training, as it performs better than sequential training and has been more explored recently. The first experiment is training in parallel 3 residual architectures and a VGG-19 [52] divided into 4 modules of equal depth on TinyImageNet. We compare in Table 1 our results in this setup to three of the best recent parallel module-wise training methods: DGL [8], PredSim [45] and Sedona [47], from Table 2 in [47]. We find that our TRGL has a much better test accuracy than the three other methods, especially on the smaller architectures. It also performs better than end-to-end training on the three ResNets. Parallel TRGL in this case with 4 modules consumes 10 to  $21\%$  less memory than end-to-end training (with a batch size of 256).
Table 1: Test accuracy of parallel TRGL with 4 modules (average and  $95\%$  confidence interval over 5 runs) on TinyImageNet, compared to DGL, PredSim, Sedona and E2E from Table 2 in [47], with memory saved compared to E2E as a percentage of E2E memory consumption in red.
Architecture	Parallel VanGL	Parallel TRGL (ours)	PredSim	DGL	Sedona	E2E
VGG-19	56.17 ± 0.29 (↓ 27%)	57.28 ± 0.20 (↓ 21%)	44.70	51.40	56.56	58.74
ResNet-50	58.43 ± 0.45 (↓ 26%)	60.30 ± 0.58 (↓ 20%)	47.48	53.96	54.40	58.10
ResNet-101	63.64 ± 0.30 (↓ 24%)	63.71 ± 0.40 (↓ 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 in Table 2 our results in this setup to the two delayed gradient methods DDG [32] and FR [31] from Table 2 in [31]. Here again, parallel TRGL has a better accuracy than the other two methods and than end-to-end training. With only two modules, the memory gains from less backpropagation are neutralized by the weight of the extra classifier and there are negligible memory savings compared to end-to-end training. However, parallel TRGL has a better test accuracy by up to almost 2 percentage points.
Table 2: Test accuracy of parallel TRGL with 2 modules (average and  $95\%$  confidence interval over 3 runs) on CIFAR100, compared to DDG, FR and E2E from Table 2 in [31].
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-152	75.40 ± 0.75	76.55 ± 1.90	73.61	76.39	74.80
The third experiment is training in parallel a ResNet-110 divided into two, four, eight and sixteen modules on STL10 [12]. We compare in Table 3 our results in this setup to the recent methods InfoPro [60] and DGL [8] from Table 2 in [60]. TRGL largely outperforms the other methods. It also outperforms end-to-end training in all but one case (that with 16 modules). With a batch size of 64, memory savings of parallel TRGL compared to end-to-end training reach  $48\%$  and  $58.5\%$  with 8 and 16 modules respectively, with comparable test accuracy. With 4 modules, TRGL training weighs  $24\%$  less than end-to-end-training, and has a test accuracy that is better by 2 percentage points (see Section 4.2 and Table 6 for a detailed memory usage comparison with InfoPro).
Table 3: Test accuracy of Parallel (Par) TRGL with  $K$  modules (average and 95% confidence interval over 5 runs) using a ResNet-110 on STL10, compared to DGL, two variants of InfoPro and E2E from Table 2 in [60].
K	Par VanGL	Par TRGL (ours)	DGL	InfoPro S	InfoPro C	E2E
2	79.85 ± 0.93	80.04 ± 0.85	75.03 ± 1.18	78.98 ± 0.51	79.01 ± 0.64	77.73 ± 1.61
4	77.11 ± 2.31	79.72 ± 0.81	73.23 ± 0.64	78.72 ± 0.27	77.27 ± 0.40	77.73 ± 1.61
8	75.71 ± 0.55	77.82 ± 0.73	72.67 ± 0.24	76.40 ± 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 4 our results with those of InfoPro [60] and InfoProL, a variant of InfoPro proposed in [46]. TRGL outperforms the other module-wise training methods. It does not outperform end-to-end training in this case, but consumes  $29\%$  less memory on CIFAR10 and CIFAR100 and  $50\%$  less on STL10, compared to  $38\%$  for InfoPro and  $45\%$  for InfoProL in [46].
Table 4: Test accuracy of parallel TRGL with 4 modules (average and  $95\%$  confidence interval over 5 runs) on a Swin-Tiny Transformer, compared to InfoPro, InfoProL and E2E from Table 3 in [46], with memory saved compared to E2E as a percentage of E2E memory consumption in red.
Dataset	Parallel VanGL	Parallel TRGL (ours)	InfoPro	InfoProL	E2E
STL10	67.00 ± 1.36 (↓ 55%)	67.92 ± 1.12 (↓ 50%)	64.61 (↓ 38%)	66.89 (↓ 45%)	72.19
CIFAR10	83.94 ± 0.42 (↓ 33%)	86.48 ± 0.54 (↓ 29%)	83.38 (↓ 38%)	86.28 (↓ 45%)	91.37
CIFAR100	69.34 ± 0.91 (↓ 33%)	74.11 ± 0.31 (↓ 29%)	68.36 (↓ 38%)	73.00 (↓ 45%)	75.03
Finally, we compare our method to InfoPro, DGL and Sedona in Table 5 below on a large scale experiment on ImageNet.
Table 5: Top 1 test accuracy of parallel TRGL with 2 modules on a 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 is lighter than end-to-end training by up to almost  $60\%$ . The extra memory consumed by our regularization compared to parallel VanGL is between 2 and  $13\%$  of
end-to-end memory. Memory savings depend then mainly on the size of the auxiliary classifier, which can easily be adjusted. Note that delayed gradients method DDG and FR increase memory usage [31], and Sedona does not claim to save memory, but rather to speed up training [47]. DGL is architecture-wise essentially identical to VanGL and consumes the same memory.
We compare in Table 6 the memory consumption of our method to that of InfoPro [60] on a ResNet-110 on STL10 with a batch size of 64 (so the same setting as in Table 3). InfoPro [60] also propose to split the network into modules that have the same weight but not necessarily the same number of layers. They only implement this for  $K \leq 4$  modules. When the modules are even in weight and not in depth, we call the training methods VanGL*, TRGL* and InfoPro*. In practice, this leads to shallower early modules which slightly hurts performance according to [47], and as seen below. However, TRGL* still outperforms InfoPro and end-to-end training, and it leads to even bigger memory savings than InfoPro*. We see in Table 6 below that TRGL saves more memory than InfoPro in two out of three cases (4 and 8 modules), and about the same in the third case (16 modules), with much better test accuracy in all cases. Likewise, TRGL* is lighter than InfoPro*, with better accuracy.
Table 6: Memory savings using a ResNet-110 on STL10 split into  $K$  modules trained in parallel with a batch size of 64, as a percentage of the weight of end-to-end training. Average test accuracy over 5 runs is between brackets. Test accuracy of end-to-end training is  ${77.73}\%$  .
K	Equally deep modules	Equally heavy modules
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)
However, parallel module-wise training does slightly slow down training. Epoch time increases by  $6\%$  with 2 modules and by  $16\%$  with 16 modules. TRGL is only slower than VanGL by  $2\%$  for all number of modules due to the additional regularization term. This is comparable to InfoPro which reports a time overhead between 1 and  $27\%$  compared to end-to-end training. See Appendix F for details.
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 as on-device training [58, 57]. We therefore test our regularization in this section in this setting, with more details in Appendix G.
We propose here to use shallower ResNets that are initially wider. These architectures are well-adapted to layer-wise training as seen in [7]. We check first in Table 10 in Appendix G that this architecture works well with parallel module-wise training with 2 modules by comparing it favorably on CIFAR10 [37] with methods DGL [8], InfoPro [60] and DDG [32] that use a ResNet-110 with the same number of parameters.
We then train a 10-block ResNet block-wise on CIFAR100. In Tables 11 and 12 in Appendix G, we see that MLS training improves the accuracy of sequential training by 0.8 percentage points when the trainset is full, but works less well on small train sets. Of the two, the regularization mainly improves the test accuracy of MLS training. The improvement increases as the training set gets smaller and reaches 1 percentage point. While parallel module-wise training performs quite close to end-to-end training in the full data regime and much better in the small data regime, sequential and multi-lap sequential training are competitive with end-to-end training in the small data regime. Combining the multi-lap trick and the regularization improves the accuracy of sequential training by 1.2 percentage points when using the entire trainset. We report further results for full block-wise training on MNIST [38] and CIFAR10 [37] in Tables 13 and 14 in Appendix G.
The  $88\%$  accuracy of sequential training on CIFAR10 in Table 13 is the same as in Table 2 of [7], which is the best method for layer-wise sequential training available, with VGG networks of comparable depth and width.
4.4 Accuracy after each module
Finally, we verify that our method avoids the stagnation or collapse in accuracy with depth. In Figure 2 below, we show the accuracy of each module with and without the regularization.
On the left, from parallel module-wise training experiments from Table 3, TRGL performs worse than vanilla greedy learning early, but surpasses it in later modules, indicating that it does avoid early overfitting. On the right, from sequential block-wise training experiments from Table 13, we see a large decline in performance that the regularization avoids. We see similar patterns in Figure 3 in Appendix G with parallel and MLS block-wise training.
Figure 2: Test accuracy after each module averaged over 10 runs with  $95\%$  confidence intervals. Left: parallel vanilla (VanGL, in blue) and regularized (TRGL, in red) module-wise training of a ResNet-110 with 16 modules on STL10 (Table 3). Right: sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training of a 10-block ResNet on  $2\%$  of CIFAR10 (Table 13).
The results in Appendix G show a few limitations of our method, as the improvements from the regularization are sometimes minimal on sequential training. However, the results show that our approach works in all settings (parallel and sequential with many or few modules), whereas other papers don't test their methods in all settings, and some show problems in other settings than the original one in subsequent papers (e.g. delayed gradients methods when the number of modules increases [31] and PredSim in [47]). Also, for parallel training in Section 4.1, the improvement from the regularization compared to VanGL is larger and increases with the number of modules (so with the memory savings) and reaches almost 5 percentage points. We show in Appendix H that our method is not very sensitive to the choice of hyperparameter  $\tau$  over a large scale.
Layer-wise training was initially considered as a pre-training and initialization method [10, 43] and was recently shown to be competitive with end-to-end training [7, 45]. Many papers consider using a different auxiliary loss, instead of or in addition to the classification loss: kernel similarity [42], information-theory-inspired losses [53, 44, 41, 60] and biologically plausible losses [53, 45, 25, 11]. Methods [7], PredSim [45], DGL [8], Sedona [47] and InfoPro [60] report the best module-wise training results. [7, 8] do it simply through the architecture choice of the auxiliary networks. Sedona applies architecture search to decide on where to split the network into modules and what auxiliary classifier to use before module-wise training. Only BoostResNet [30] also proposes a block-wise training idea geared for ResNets. However, their results only show better early performance and end-to-end fine-tuning is required to be competitive. A method called ResIST [19] that is similar to block-wise training of ResNets randomly assigns ResBlocks to one of up to 16 modules that
are trained independently and reassembled before another random partition. More of a distributed training method, it is only compared to local SGD [54]. These methods can all be combined with our regularization, and we do use the auxiliary classifier from [7, 8].
Besides module-wise training, methods such as DNI [33, 14], DDG [32] and FR [31], solve the update and backward locking problems with an eye towards parallelization by using delayed or predicted gradients, or even predicted inputs to address forward locking, which is what [55] do. But they observe training issues with more than 5 modules [31]. This makes them compare unfavorably to module-wise training [8]. The high dimension of the predicted gradient, which scales with the size of the network, makes [33, 14] challenging in practice. Therefore, despite its simplicity, greedy module-wise training is more appealing when working in a constrained setting.
Viewing ResNets as dynamic transport systems [16, 36] followed from their view as a discretization of ODEs [61]. Transport regularization of ResNets in particular is motivated by the observation that they are naturally biased towards minimally modifying their input [35, 26]. We further linked this transport viewpoint with gradient flows in the Wasserstein space to apply it in a principled way to module-wise training. Gradient flows on the data distribution appeared recently in deep learning. In [1], the focus is on functionals of measures whose first variations are known in closed form and used, through their gradients, in the algorithm. This limits the scope of their applications to transfer learning and similar tasks. Likewise, [21, 40, 4, 3] use the explicit gradient flow of  $f$ -divergences and other distances between measures for generation and generator refinement. In contrast, we use the discrete minimizing movement scheme which does not require computation of the first variation and allows to consider classification.
We introduced a transport regularization for module-wise training that theoretically links module-wise training to gradient flows of the loss in probability space. Our method provably leads to more regular modules and experimentally improves the test accuracy of module-wise parallel, sequential and multi-lap sequential (a variant of sequential training that we introduce) training. Through this simple method that does not complexify the architecture, we make module-wise training competitive with end-to-end training while benefiting from its lower memory usage. Being a regularization, the method can easily be combined with other layer-wise training methods. Future work can experiment with working in Wasserstein space  $W_{p}$  for  $p \neq 2$ , i.e. regularizing with a norm  $\lVert .\rVert_p$  with  $p \neq 2$ .
[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 NeurIPS (2019).
[5] ARORA, R., BASU, A., MIANJY, P., AND MUKHERJEE, A. Understanding deep neural networks with rectified linear units. In International Conference on Learning Representations (2018).
[6] BACH, F. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research 18, 19 (2017), 1-53.
[7] BELIrovsky, E., EICKENBERG, M., AND OYALLON, E. Greedy layerwise learning can scale to imagenet. In ICML (2019).
[8] BELIrovsky, E., EICKENBERG, M., AND OYALLON, E. Decoupled greedy learning of cnns. In ICML (2020).
[9] BENAMOU, J., AND BRENIER, Y. A computational fluid mechanics solution to the mongekantorovich mass transfer problem. Numerische Mathematik (2000).
[58] TENG, Q., WANG, K., ZHANG, L., AND HE, J. The layer-wise training convolutional neural networks using local loss for sensor-based human activity recognition. IEEE Sensors Journal (2020).
[59] VILLANI, C. Optimal Transport: Old and New. Springer-Verlag, 2008.
[60] WANG, Y., NI, Z., SONG, S., AND LE YANG, G. H. Revisiting locally supervised learning: an alternative to end-to-end training. In ICLR (2021).
[61] WEINAN, E. A proposal on machine learning via dynamical systems. Commun. Math. Stat (2017).
[62] XIE, S., ET AL. Aggregated residual transformations for deep neural networks. In CVPR (2017).
[63] ZAGORUYKO, S., AND KOMODAKIS, N. Wide residual networks. In BMVC (2016).
[64] ZHANG, J., ET AL. Towards robust resnet: A small step but a giant leap. In *IJCAI* (2019).
[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.
A Background on optimal transport
The Wasserstein space  $\mathbb{W}_2(\Omega)$  with  $\Omega$  a convex and compact subset of  $\mathbb{R}^d$  is the space  $\mathcal{P}(\Omega)$  of probability measures over  $\Omega$ , equipped with the distance  $W_2$  given by the solution to the optimal transport problem
$$
W _ {2} ^ {2} (\alpha , \beta) = \min  _ {\gamma \in \Pi (\alpha , \beta)} \int_ {\Omega \times \Omega} \| x - y \| ^ {2} d \gamma (x, y) \tag {12}
$$
where  $\Pi (\alpha ,\beta)$  is the set of probability distribution over  $\Omega \times \Omega$  with first marginal  $\alpha$  and second marginal  $\beta$ , i.e.  $\Pi (\alpha ,\beta) = \{\gamma \in \mathcal{P}(\Omega \times \Omega)\mid \pi_{1\sharp}\gamma = \alpha ,\pi_{2\sharp}\gamma = \beta \}$  where  $\pi_1(x,y) = x$  and  $\pi_2(x,y) = y$ . The optimal transport problem can be seen as looking for a transportation plan minimizing the cost of displacing some distribution of mass from one configuration to another. This problem indeed has a solution in our setting and  $W_{2}$  can be shown to be a geodesic distance (see for example [49, 59]). If  $\alpha$  is absolutely continuous and  $\partial \Omega$  is  $\alpha$ -negligible then the problem in (12) (called the Kantorovich problem) has a unique solution and is equivalent to the Monge problem, i.e.
$$
W _ {2} ^ {2} (\alpha , \beta) = \min  _ {T \text {s . t .} T _ {\sharp} \alpha = \beta} \int_ {\Omega} \| T (x) - x \| ^ {2} \mathrm {d} \alpha (x) \tag {13}
$$
and this problem has a unique solution  $T^{\star}$  linked to the solution  $\gamma^{\star}$  of (12) through  $\gamma^{\star} = (\mathrm{id}, T^{\star})_{\sharp} \alpha$ . Another equivalent formulation of the optimal transport problem in this setting is the dynamical formulation [9]. Here, instead of directly pushing samples of  $\alpha$  to  $\beta$  using  $T$ , we can equivalently displace mass, according to a continuous flow with velocity  $v_t: \mathbb{R}^d \to \mathbb{R}^d$ . This implies that the density  $\alpha_t$  at time  $t$  satisfies the continuity equation  $\partial_t \alpha_t + \nabla \cdot (\alpha_t v_t) = 0$ , assuming that initial and final conditions are given by  $\alpha_0 = \alpha$  and  $\alpha_1 = \beta$  respectively. In this case, the optimal displacement is the one that minimizes the total action caused by  $v$ :
$$
W _ {2} ^ {2} (\alpha , \beta) = \min  _ {v} \int_ {0} ^ {1} \| v _ {t} \| _ {L ^ {2} \left(\alpha_ {t}\right)} ^ {2} \mathrm {d} t \tag {14}
$$
$$
\mathrm {s . t .} \partial_ {t} \alpha_ {t} + \nabla \cdot (\alpha_ {t} v _ {t}) = 0, \alpha_ {0} = \alpha , \alpha_ {1} = \beta
$$
Instead of describing the density's evolution through the continuity equation, we can describe the paths  $\phi_t^x$  taken by particles at position  $x$  from  $\alpha$  when displaced along the flow  $v$ . Here  $\phi_t^x$  is the position at time  $t$  of the particle that was at  $x \sim \alpha$  at time 0. The continuity equation is then equivalent to  $\partial_t \phi_t^x = v_t (\phi_t^x)$ . See chapters 4 and 5 of [49] for details. Rewriting the conditions as necessary, Problem (14) becomes
$$
W _ {2} ^ {2} (\alpha , \beta) = \min  _ {v} \int_ {0} ^ {1} \| v _ {t} \| _ {L ^ {2} \left(\left(\phi_ {t}\right) _ {\sharp} \alpha\right)} ^ {2} \mathrm {d} t \tag {15}
$$
$$
\mathrm {s . t .} \partial_ {t} \phi_ {t} ^ {x} = v _ {t} \left(\phi_ {t} ^ {x}\right), \dot {\phi_ {0}} = \mathrm {i d}, \left(\dot {\phi_ {1}}\right) _ {\sharp} \alpha = \beta
$$
and the optimal transport map  $T^{\star}$  that solves (13) is in fact  $T^{\star}(x) = \phi_1^x$  for  $\phi$  that solves the continuity equation together with the optimal  $v^{\star}$  from (15). We refer to [49, 59] for these results on optimal transport.
Optimal transport maps have some regularity properties under some boundedness assumptions. We mention the following result from [20]:
Theorem A.1. Let  $\alpha$  and  $\beta$  be absolutely continuous measures on  $\mathbb{R}^d$  and  $T$  the optimal transport map between  $\alpha$  and  $\beta$  for the Euclidean cost. Suppose there are bounded open sets  $X$  and  $Y$ , such that the density of  $\alpha$  (respectively of  $\beta$ ) is null on  $X^c$  (respectively  $Y^c$ ) and bounded away from zero and infinity on  $X$  (respectively  $Y$ ).
Then there exists two relatively closed sets of null measure  $A \subset X$  and  $B \subset Y$ , such that  $T$  is  $\eta$ -Hölder continuous from  $X \setminus A$  to  $Y \setminus B$ , i.e.  $\forall x, y \in X \setminus A$  we have
$$
\left\| T (x) - T (y) \right\| \leq C \| x - y \| ^ {\eta} f o r c o n s t a n t s \eta , C > 0
$$
B Background on gradient flows
We follow [50, 2] for this background on gradient flows. Given a function  $\mathcal{L}:\mathbb{R}^d\to \mathbb{R}$  and an initial point  $x_0\in \mathbb{R}^d$ , a gradient flow is a curve  $x:[0,\infty [\rightarrow \mathbb{R}^d$  that solves the Cauchy problem
$$
\left\{ \begin{array}{l} x ^ {\prime} (t) = - \nabla \mathcal {L} (x (t)) \\ x (0) = x _ {0} \end{array} \right. \tag {16}
$$
A solution exists and is unique if  $\nabla \mathcal{L}$  is Lipschitz or  $\mathcal{L}$  is convex. Given  $\tau > 0$  and  $x_0^\tau = x_0$  define a sequence  $(x_k^\tau)_k$  through the minimizing movement scheme:
$$
x _ {k + 1} ^ {\tau} \in \arg \min  _ {x \in \mathbb {R} ^ {d}} \mathcal {L} (x) + \frac {1}{2 \tau} \| x - x _ {k} ^ {\tau} \| ^ {2} \tag {17}
$$
$\mathcal{L}$  lower semi-continuous and  $\mathcal{L}(x)\geq C_1 - C_2\| x\|^2$  guarantees existence of a solution of (17) for  $\tau$  small enough.  $\mathcal{L}$ $\lambda$  convex meets these conditions and also provides uniqueness of the solution because of strict convexity of the target. See [49, 50, 2].
We interpret the point  $x_{k}^{\tau}$  as the value of a curve  $x$  at time  $k\tau$ . We can then construct a curve  $x^{\tau}$  as the piecewise constant interpolation of the points  $x_{k}^{\tau}$ . We can also construct a curve  $\tilde{x}^{\tau}$  as the affine interpolation of the points  $x_{k}^{\tau}$ .