SlowGuess/NeurIPS_2023_OCR huggingface dataset:
If  $\mathcal{L}(x_0) < \infty$  and  $\inf \mathcal{L} > -\infty$  then  $(x^{\tau})$  and  $(\tilde{x}^{\tau})$  converge uniformly to the same curve  $x$  as  $\tau$  goes to zero (up to extracting a subsequence). If  $\mathcal{L}$  is  $\mathcal{C}^1$ , then the limit curve  $x$  is a solution of (16) (i.e. a gradient flow of  $\mathcal{L}$ ). If  $\mathcal{L}$  is not differentiable then  $x$  is solution of the problem defined using the subdifferential of  $\mathcal{L}$ , i.e.  $x$  satisfies  $x'(t) \in -\partial \mathcal{L}(x(t))$  for almost every  $t$ .
If  $\mathcal{L}$  is  $\lambda$ -convex with  $\lambda > 0$ , then the solution to (16) converges exponentially to the unique minimizer of  $\mathcal{L}$  (which exists by coercivity). So taking  $\tau \to 0$  and  $k \to \infty$ , we tend towards the minimizer of  $\mathcal{L}$ .
The advantage of the minimizing movement scheme (17) is that it can be adapted to metric spaces by replacing the Euclidean distance by the metric space's distance. In the (geodesic) metric space  $\mathbb{W}_2(\Omega)$  with  $\Omega$  convex and compact, for  $\mathcal{L}:\mathbb{W}_2(\Omega)\to \mathbb{R}\cup \{\infty \}$  lower semi-continuous for the weak convergence of measures in duality with  $\mathcal{C}(\Omega)$  (equivalent to lower semi-continuous with respect to the distance  $W_{2}$ ) and  $\rho_0^\tau = \rho_0\in \mathcal{P}(\Omega)$ , the minimizing movement scheme (17) becomes
$$
\rho_ {k + 1} ^ {\tau} \in \arg \min  _ {\rho \in \mathcal {P} (\Omega)} \mathcal {L} (\rho) + \frac {1}{2 \tau} W _ {2} ^ {2} (\rho , \rho_ {k} ^ {\tau}) \tag {18}
$$
This problem has a solution because the objective is lower semi-continuous and the minimization is over  $\mathcal{P}(\Omega)$  which is compact by Banach-Alaoglu.
We can construct a piecewise constant interpolation between the measures  $\rho_k^\tau$ , or a geodesic interpolation where we travel along a geodesic between  $\rho_k^\tau$  and  $\rho_{k + 1}^\tau$  in  $\mathbb{W}_2(\Omega)$ , constructed using the optimal transport map between these measures. Again, if  $\mathcal{L}(x_0) < \infty$  and inf  $\mathcal{L} > -\infty$  then both interpolations converge uniformly to a limit curve  $\tilde{\rho}$  as  $\tau$  goes to zero. Under further conditions on  $\mathcal{L}$ , mainly  $\lambda$ -geodesic convexity (i.e.  $\lambda$ -convexity along geodesics) for  $\lambda > 0$ , we can prove stability and convergence of  $\tilde{\rho}(t)$  to a minimizer of  $\mathcal{L}$  as  $t \to \infty$ , see [49, 50, 2].
C Proof of Proposition 2.3
Proof. Take a minimizing sequence  $(\tilde{F}_i, \tilde{T}_i)$ , i.e., such that  $\mathcal{C}(\tilde{F}_i, \tilde{T}_i) \to \min \mathcal{C}$ , where  $\mathcal{C} \geq 0$  is the target function in (3) and denote  $\beta_i = \tilde{T}_{i\sharp} \rho_k^\tau$ . Then by compactness  $\tilde{F}_i \to F^\star$  and  $\beta_i \rightharpoonup \beta^\star$  in duality with  $\mathcal{C}_b(\Omega)$  by Banach-Alaoglu. There exists  $T^\star$  an optimal transport map between  $\rho_k^\tau$  and  $\beta^\star$ . Then  $\mathcal{C}(F^\star, T^\star) \leq \lim \mathcal{C}(\tilde{F}_i, \tilde{T}_i) = \min \mathcal{C}$  by continuity of  $\mathcal{L}$  and because
$$
\begin{array}{l} \int_ {\Omega} \| T ^ {\star} (x) - x \| ^ {2} \mathrm {d} \rho_ {k} ^ {\tau} (x) = W _ {2} ^ {2} (\rho_ {k} ^ {\tau}, \beta^ {\star}) \\ = \lim  W _ {2} ^ {2} \left(\rho_ {k} ^ {\tau}, \beta_ {i}\right) \\ \leq \lim  \int_ {\Omega} \| \tilde {T} _ {i} (x) - x \| ^ {2} d \rho_ {k} ^ {\tau} (x) \\ \end{array}
$$
as  $W_{2}$  metrizes weak convergence of measures. We take  $(F_k^\tau, T_k^\tau) = (F^\star, T^\star)$ . It is also immediate that for any minimizing pair, the transport map has to be optimal. Taking a minimizing sequence  $(\tilde{F}_i, \tilde{v}^i)$  and the corresponding induced maps  $\tilde{T}_i$  we get the same result for Problem (11). Problems (3) and (11) are equivalent by the equivalence between Problems (13) and (15).
D Varying the regularization weight
The discussion in Section 2.2 suggests taking a fixed weight  $\tau$  for the transport cost that is as small as possible. However, instead of using a fixed  $\tau$ , we might want to vary it along the depth  $k$  to further
constrain with a smaller  $\tau_{k}$  the earlier modules to avoid that they overfit or the later modules to maintain the accuracy of earlier modules. We might also want to regularize the network further in earlier epochs when the data is more entangled. To unify and formalize this varying weight  $\tau_{k,i}$  across modules  $k$  and SGD iterations  $i$ , we use a scheme inspired by the method of multipliers to solve Problems (2) and (10). To simplify the notations, we will instead consider the weight  $\lambda_{k,i} := 2\tau_{k,i}$  given to the loss. We denote  $\theta_{k,i}$  the parameters of both  $T_{k}$  and  $F_{k}$  at SGD iteration  $i$ . We also denote  $L(\theta, x)$  and  $W(\theta, x)$  respectively the loss and the transport regularization as functions of parameters  $\theta$  and data  $x$ . We now increase the weight  $\lambda_{k,i}$  of the loss every  $s$  iterations of SGD by a value that is proportional to the current loss. Given increase factor  $h > 0$ , initial parameters  $\theta_{k,1}$ , initial weights  $\lambda_{k,1} \geq 0$ , learning rates  $(\eta_{i})$  and batches  $(x_{i})$ , we apply for module  $k$  and  $i \geq 1$ :
$$
\left\{ \begin{array}{l l} \theta_ {k, i + 1} & = \theta_ {k, i} - \eta_ {i} \nabla_ {\theta} (\lambda_ {k, i} L (\theta_ {k, i}, x _ {i}) + W (\theta_ {k, i}, x _ {i})) \\ \lambda_ {k, i + 1} & = \lambda_ {k, i} + h L (\theta_ {k, i + 1}, x _ {i + 1}) \text {i f} i \bmod s = 0 \text {e l s e} \lambda_ {k, i} \end{array} \right.
$$
The weights  $\lambda_{k,i}$  will vary along modules  $k$  because they will evolve differently with iterations  $i$  for each  $k$ . They will increase more slowly with  $i$  for larger  $k$  because deeper modules will have smaller loss. This method can be seen as a method of multipliers for the problem of minimizing the transport under the constraint of zero loss. Therefore it is immediate by slightly adapting the proof of Proposition 2.3 or from [36] that we are still solving a problem that admits a solution whose non-auxiliary part is an optimal transport map with the same regularity as stated above. We use the same initial value  $\lambda_1 = \lambda_{k,1}$  for all modules so that this method requires choosing three hyper-parameters  $(h,s$  and  $\lambda_1)$ . In practice (see Section 4.1 and Appendix E), it works best in only one experiment. Simply manually finding a value of  $\tau$  that is multiplied by 2 for the second half of the network works best in all the other experiments.
We use standard data augmentation and standard implementations for VGG-19, ResNet-50, ResNet-101, ResNet-110, ResNet-152 and Swin-Tiny Transformer (the same as for the other methods in Section 4.1). We use NVIDIA Tesla V100 16GB GPUs for the experiments. Training a Resnet-152 on TinyImageNet in Table 1 takes about 36 hours. Training a Resnet-152 on CIFAR100 in Table 1 takes about 11 hours. Training a ResNet-110 on STL10 in Table 3 takes about 3 hours. Training a Swin-Tiny Transformer in Table 4 take between 40 minutes and 1 hour.
For sequential and multi-lap sequential training, we use SGD with a learning rate of 0.007. With the exception of the Swin Transformer in Table 4, we use SGD for parallel training with learning rate of 0.003 in all Tables but Table 3 where the learning rate is 0.002. For the Swin Transformer in Table 4, we use the AdamW optimizer with a learning rate of 0.007 and a CosineLR scheduler.
For end-to-end training we use a learning rate of 0.1 that is divided by five at epochs 120, 160 and 200. Momentum is always 0.9. For parallel and end-to-end training, we train for 300 epochs. For sequential and multi-lap sequential training, the number of epochs varies per module (see Section G).
For experiments in Section 4.1, we use a batch size of 256, orthogonal initialization [51] with a gain of 0.1, label smoothing of 0.1 and weight decay of 0.0002. The batch size changes to 64 for Table 3 and to 1024 for Table 4.
For experiments in Section 4.3, we use a batch size of 128, orthogonal initialization with a gain of 0.05, no label smoothing and weight decay of 0.0001.
In Table 1, we use  $\tau = 500000$  for the first two modules and then double it for the last two modules for TRGL. In Table 2, we use  $\lambda_{k,1} = 1$ ,  $h = 1$  and  $s = 50$  for TRGL. In Table 3, we use  $\tau = 0.5$  and double it at the midpoint, expect for the first line where  $\tau = 50$ .
F Memory savings and training time
We compare in Table 7 the memory consumption of our method to that of InfoPro [60] on a ResNet-110 split into  $K$  modules trained in parallel on STL10 with a batch size of 64 (so the same setting as in Table 3 in Section 4.1). We report in Table 7 the memory saved as a percentage of the 6230 MiB memory required by end-to-end training with the same batch size. VanGL refers to our architecture trained without the regularization. InfoPro [60] also propose to split the network into  $K$  modules that
have the same weight but not necessarily the same number of layers. They only implement this for  $K \leq 4$  modules. When the modules are even in weight and not in depth, we call the training methods VanGL*, TRGL* and InfoPro*. In practice, this leads to shallower early modules which slightly hurts performance according to [47]. We verify this in Table 8 (to be compared with Table 3 in Section 4.1). However, TRGL* still outperforms InfoPro and end-to-end training, and it leads to even bigger memory savings. We see in Table 7 that TRGL saves more memory than InfoPro in two out of three cases (4 and 8 modules), and about the same in the third case (16 modules), with much better test accuracy in all cases. Likewise, TRGL* is lighter than InfoPro*, with better accuracy. We also see that the added memory cost of the regularization compared to vanilla greedy learning is small. However, parallel module-wise training does slightly slow down training. Epoch time increases by  $6\%$  with 2 modules and by  $16\%$  with 16 modules. TRGL is only slower than VanGL by  $2\%$  for all number of modules due to the additional regularization term. This is comparable to InfoPro which report a time overhead between 1 and  $27\%$  compared to end-to-end training.
Table 7: Memory savings using a ResNet-110 on STL10 split into  $K$  modules trained in parallel with a batch size of 64, as a percentage of the weight of end-to-end training. Average test accuracy over 5 runs is between brackets. Test accuracy of end-to-end training is  ${77.73}\%$  .
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)
Table 8: Test accuracy of parallel (Par) TRGL* with  $K$  modules (average and 95% confidence interval over 5 runs) on a ResNet-110 trained on STL10, compared to InfoPro* and E2E training from Table 3 in [60]
K	Par VanGL*	Par TRGL* (ours)	InfoPro*
2	79.05 ± 1.33	79.47 ± 1.36	79.05 ± 0.57
4	77.14 ± 1.23	78.94 ± 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 no save memory (they actually increase memory usage, see FR [31]). Sedona [47] also does not claim to save memory, as their first module (the heaviest) is deeper than the others, but rather to speed up computation. Finally, DGL [8] is architecture-wise essentially identical to VanGL and consumes the same memory.
G Sequential full block-wise training
To show that our method works well with all types of module-wise training when using few modules, we train a ResNet-101 split in 2 modules on CIFAR100, sequentially and multi-lap sequentially. Results are in Table 9. We see that our idea of multi-lap sequential training adds one percentage point of accuracy to sequential training, and that the regularization further improves the accuracy by about half a percentage point. As only one module has to be trained at a time, these two training methods require only around half the memory end-to-end training requires (the size of the heaviest module and its classifier more exactly).
Table 9: Test accuracy of sequential (Seq) and multi-lap sequential (MLS) TRGL and VanGL with 2 modules on CIFAR100 using ResNet-101 (average of 2 runs).
Seq VanGL	Seq TRGL	MLS VanGL	MLS TRGL
73.31	73.61	74.34	74.78
We now focus on full block-wise training, meaning that each module is a single ResBlock, mostly sequentially. We propose here to use shallower and initially wider ResNets with a downsampling and 256 filters initially and a further downsampling and doubling of the number of filters at the midpoint, no matter the depth. In these ResNets, we use the ResBlock from [27] with two convolutional layers. If such a network is divided in  $K$  modules of  $M$  ResBlocks each, we call the network a  $K - M$  ResNet. These wider shallower architectures are well-adapted to layer-wise training as seen in [7]. We check in Table 10 that this architecture works well with parallel module-wise training by comparing favorably on CIFAR10 ([37]) a 2-7 ResNet with DGL, InfoPro ([60]) and DDG [32]. The 2-7 ResNet has 45 million parameters, which is about the same as the ResNet-110 divided in two used by the other methods, and performs better when trained in parallel.
Table 10: Average test accuracy and  $95\%$  confidence interval of 2-7 ResNet over 10 runs on CIFAR10 with parallel TRGL and VanGL, compared to DGL and DDG from [8] and InfoPro from [60] that split a ResNet-110 in 2 module-wise-parallel-trained modules.
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 to show that our method works in all settings (parallel or sequential with many or few splits), which other methods don't do. For sequential training, block  $k$  is trained for  $50 + 10k$  epochs where  $0 \leq k \leq 10$ , block 0 being the encoder. This idea of increasing the number of epochs along with the depth is found in [43]. For MLS training, block  $k$  is trained for  $10 + 2k$  epochs, and this is repeated for 5 laps. In block-wise training, the last block does not always perform the best and we report the accuracy of the best block. In Table 11, we see that MLS training improves the test accuracy of sequential training by around 0.8 percentage points when the training dataset is full, but works less well on small training sets. Of the two, the regularization mainly improves the test accuracy of MLS training. The improvement increases as the training set gets smaller and reaches 1 percentage point. That is also the case for parallel module-wise training in Table 12, which already performs quite close to end-to-end training in the full data regime and much better in the small data regime. Combining the multi-lap trick and the regularization improves the performance of sequential training by 1.2 percentage points.
Table 11: Average highest test accuracy and  $95\%$  confidence interval of 10-1 ResNet over 10 runs on CIFAR100 with different train sizes and sequential (Seq), multi-lap sequential (MLS) and parallel (Par) TRGL and VanGL, compared to E2E.
Train size	Seq VanGL	Seq TRGL	MLS VanGL	MLS 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
12500	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 ± 0.97
5000	39.45 ± 0.23	40.36 ± 0.23	36.38 ± 0.31
We report further results of block-wise training on CIFAR10 in Table 13 and on MNIST [38] in Table 14, but now we report the accuracy of the last block. We see again greater improvement due to the regularization as the training set gets smaller, gaining up to 6 percentage points.
Table 13: Average last block test accuracy and  $95\%$  confidence interval of 10-1 ResNet over 10 runs on CIFAR10 with different train sizes and sequential TRGL and VanGL, compared to E2E.
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 ± .39	77.18 ± .34	82.61 ± .35
5000	67.74 ± .49	69.67 ± .44	73.93 ± .67
1000	45.67 ± .88	51.34 ± .90	50.63 ± .98
Table 14: Average last block test accuracy and  $95\%$  confidence interval of 20-1 ResNet (32 filters, fixed encoder, same classifier) over 20/50 runs on MNIST with different train sizes and parallel TRGL and VanGL, compared to E2E.
Train size	Par VanGL	Par TRGL	E2E
60000	99.07 ± .04	99.08 ± .04	99.30 ± .03
30000	98.90 ± .05	98.93 ± .06	99.22 ± .03
12000	98.52 ± .06	98.59 ± .06	98.96 ± .06
6000	98.05 ± .09	98.16 ± .07	98.62 ± .06
1500	96.34 ± .12	96.91 ± .07	97.19 ± .08
1200	95.80 ± .12	96.58 ± .09	96.88 ± .09
600	91.35 ± .99	95.16 ± .15	95.30 ± .17
300	89.81 ± .73	92.86 ± .24	92.87 ± .28
150	81.84 ± 1.22	87.48 ± .42	87.82 ± .59
The  $88\%$  accuracy of sequential training on CIFAR10 in Table 13 is the same as for sequential training in table 2 of [7], which is the best method for layer-wise sequential training available, with VGG networks of comparable depth and width.
Figure 3: Test accuracy after each block of 10-1 ResNet averaged over 10 runs with  $95\%$  confidence intervals. Left: multi-lap sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) blockwise 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.
H Sensitivity to hyper-parameters
We show in Figure 4 below that TRGL still performs better than VanGL (in the same setting as in Table 3 in Section 4.1) for values of  $\tau$  from 0.03 to 100 and is still roughly equivalent to it for values up to 5000.
Figure 4: Average test accuracy over 5 runs of parallel TRGL using a ResNet110 on STL10 with 16 modules with different values of  $\tau$  (in red), and of VanGL (blue line).
I Broader impact
Less memory usage has a positive environmental impact and allows organizations with less resources to use deep learning.