SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Skander Karkar
Criteo, Sorbonne Université
Ibrahim Ayed
Sorbonne Université, Thales
Emmanuel de Bezenac
Patrick Gallinari
Greedy layer-wise or module-wise training of neural networks is compelling in constrained and on-device settings where memory is limited, as it circumvents a number of problems of end-to-end back-propagation. However, it suffers from a stagnation problem, whereby early layers overfit and deeper layers stop increasing the test accuracy after a certain depth. We propose to solve this issue by introducing a module-wise regularization inspired by the minimizing movement scheme for gradient flows in distribution space. We call the method TRGL for Transport Regularized Greedy Learning and study it theoretically, proving that it leads to greedy modules that are regular and that progressively solve the task. Experimentally, we show improved accuracy of module-wise training of various architectures such as ResNets, Transformers and VGG, when our regularization is added, superior to that of other module-wise training methods and often to end-to-end training, with as much as  $60\%$  less memory usage.
End-to-end backpropagation is the standard training method of neural networks. However, it requires storing the whole model and computational graph during training, which requires large memory consumption. It also prohibits training the layers in parallel. Dividing the network into modules, a module being made up of one or more layers, accompanied by auxiliary classifiers, and greedily solving module-wise optimization problems sequentially (i.e. one after the other fully) or in parallel (i.e. at the same time batch-wise), consumes much less memory than end-to-end training as it does not need to store as many activations, and when done sequentially, only requires loading and training one module (so possibly one layer) at a time. Module-wise training has therefore been used in constrained settings in which end-to-end training can be impossible such as training on mobile devices [58, 57] and dealing with very large whole slide images [65]. When combined with batch buffers, parallel module-wise training also allows for parallel training of the modules [8]. Despite its simplicity, module-wise training has been recently shown to scale well [8, 47, 60, 45], outperforming more complicated alternatives to end-to-end training such as synthetic [33, 14] and delayed [32, 31] gradients, while having superior memory savings.
In a classification task, module-wise training splits the network into successive modules, a module being made up of one or more layers. Each module takes as input the output of the previous module, and each module has an auxiliary classifier so that a local loss can be computed, with backpropagation happening only inside the modules and not between them (see Figure 1 below).
The main drawback of module-wise training is the well-documented stagnation problem observed in [43, 7, 60, 47], whereby early modules overfit and learn more discriminative features than end-to-end
training, destroying task-relevant information, and deeper modules don't improve the test accuracy significantly, or even degrade it, which limits the deployment of module-wise training. We further highlight this phenomenon in Figures 2 and 3 in Section 4.4. To tackle this issue, InfoPro [60] propose to maximize the mutual information that each module keeps with the input, in addition to minimizing the loss. [7] make the auxiliary classifier deeper and Sedona [47] make the first module deeper. These last two methods lack a theoretical grounding, while InfoPro requires a second auxiliary network for each module besides the classifier. We propose a different perspective, leveraging the analogy between residual connections and the Euler scheme for ODEs [61]. To preserve input information, we minimize the kinetic energy of the modules along with the training loss. Intuitively, this forces the modules to change their input as little as possible. We leverage connections with the theories of gradient flows in distribution space and optimal transport to analyze our method theoretically.
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 [62], Wide ResNet [63], EfficientNet [56] and MobileNetV2 [48]) and vision transformers that are made up essentially of residual connections [39, 17], but is immediately usable on any network where many layers have the same input and output dimension such as VGG [52]. Our contributions are the following:
- We propose a new method for module-wise training. Being a regularization, it is lighter than many recent state-of-the-art methods (PredSim [45], InfoPro [60]) that train another auxiliary network besides the auxiliary classifier for each module.
- We theoretically justify our method, proving that it is a transport regularization that forces the module to be an optimal transport map making it more regular and stable. We also show that it amounts to a discretization of the gradient flow of the loss in probability space, which means that the modules progressively minimize the loss and explains why the method avoids the accuracy collapse observed in module-wise training.
- Experimentally, we consistently improve the test accuracy of module-wise trained networks (ResNets, VGG and Swin-Transformer) beating 8 other methods, in sequential and parallel module-wise training, and also in multi-lap sequential training, a variant of sequential module-wise training that we introduce and that performs better in many cases. In particular, our regularization makes parallel module-wise training superior or comparable in accuracy to end-to-end training, while consuming  $10\%$  to  $60\%$  less memory.
2 Transport-regularized module-wise training
The typical setting of (sequential) module-wise training for minimizing a loss  $L$ , is, given a dataset  $\mathcal{D}$ , to solve one after the other, for  $1 \leq k \leq K$ , Problems
$$
\left(T _ {k}, F _ {k}\right) \in \arg \min  _ {T, F} \sum_ {x \in \mathcal {D}} L \left(F, T \circ G _ {k - 1} (x)\right) \tag {1}
$$
where  $G_{k} = T_{k} \circ \ldots \circ T_{1}$  for  $1 \leq k \leq K$ ,  $G_{0} = \mathrm{id}$ ,  $T_{k}$  is the module (one or many layers) and  $F_{k}$  is an auxiliary classifier. Module  $T_{k}$  receives the output of module  $T_{k-1}$ , and auxiliary classifier  $F_{k}$  computes the prediction from the output of  $T_{k}$  so the loss can be computed. The inputs are  $x$  and  $L$  has access to their labels  $y$  to calculate the loss. i.e.  $L(F, T \circ G_{k-1}(x)) = l(F \circ T \circ G_{k-1}(x), y)$  where  $l$  is a machine learning loss such as cross-entropy. See Figure 1. The final network trained this way is  $F_{K} \circ G_{K}$ . But, at inference, we can stop at any depth  $k$  and use  $F_{k} \circ G_{k}$  if it performs better. Indeed, an intermediate module often performs as well or better than the last module because of the early overfitting and subsequent stagnation or collapse problem of module-wise training [43, 7, 60, 47].
We propose below in (2) a regularization that avoids the destruction of task-relevant information by the early modules by forcing them to minimally modify their input. Proposition 2.2 proves that by using our regularization (2), we are indeed making the modules build upon each other to solve the task, which is the property we desire in module-wise training, as the modules now act as successive proximal optimization steps in the minimizing movement scheme optimization algorithm for maximizing the separability of the data representation. The background on optimal transport (OT), gradient flows and the minimizing movement scheme is in Appendices A and B.
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 is a function  $T = \mathrm{id} + r$ , which includes many transformer architectures [39, 17]), its kinetic energy is simply the squared norm of its residue  $r = T - \mathrm{id}$ , which we add to the loss  $L$  in the target of the greedy problems (1). All layers that have the same input and output dimension can be rewritten as residual blocks and the analysis applies to a large variety of architectures such as VGG [52]. Given  $\tau > 0$ , we now solve, for  $1 \leq k \leq K$ , Problems
$$
\left(T _ {k} ^ {\tau}, F _ {k} ^ {\tau}\right) \in \arg \min  _ {T, F} \sum_ {x \in \mathcal {D}} L \left(F, T \circ G _ {k - 1} ^ {\tau} (x)\right) + \frac {1}{2 \tau} \| T \circ G _ {k - 1} ^ {\tau} (x) - G _ {k - 1} ^ {\tau} (x) \| ^ {2} \tag {2}
$$
where  $G_{k}^{\tau} = T_{k}^{\tau} \circ \dots \circ T_{1}^{\tau}$  for  $1 \leq k \leq K$  and  $G_{0}^{\tau} = \mathrm{id}$ . The final network is  $F_{K}^{\tau} \circ G_{K}^{\tau}$ . Intuitively, this biases the modules towards moving the points as little as possible, thus at least keeping the performance of the previous module. Residual connections are already biased towards small displacements and this bias is desirable and should be encouraged [35, 64, 26, 15, 36]. But the method can be applied to any module where  $T(x)$  and  $x$  have the same dimension so that  $T(x) - x$  can be computed.
To facilitate the theoretical analysis, we rewrite the method in a more general formulation using data distribution  $\rho$ , which can be discrete or continuous, and the distribution-wide loss  $\mathcal{L}$  that arises from the point-wise loss  $L$ . Then Problem (2) is equivalent to Problem
$$
\left(T _ {k} ^ {\tau}, F _ {k} ^ {\tau}\right) \in \arg \min  _ {T, F} \mathcal {L} (F, T _ {\sharp} \rho_ {k} ^ {\tau}) + \frac {1}{2 \tau} \int_ {\Omega} \| T (x) - x \| ^ {2} d \rho_ {k} ^ {\tau} (x) \tag {3}
$$
with  $\rho_{k + 1}^{\tau} = (T_k^{\tau})_{\sharp}\rho_k^{\tau}$ ,  $\rho_1^\tau = \rho$  and  $\mathcal{L}(F,T_{\sharp}\rho_k^\tau) = \int L(F,T(x))\mathrm{d}\rho_k^\tau (x) = \int L(F,z)\mathrm{d}T_{\sharp}\rho_k^\tau (x)$ .
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  $\mathcal{Z}(\mu)\coloneqq \min_F\mathcal{L}(F,\mu)$ , which is the loss of the best classifier. If we are limited to linear classifiers,  $\mathcal{Z}(\rho_k^\tau)$  is the linear separability of the representation  $\rho_{k}^{\tau}$  at module  $k$  of the data distribution  $\rho$ . The MMS, introduced in [24, 23], is a metric counterpart to Euclidean gradient descent for minimizing functionals over distributions. In our case,  $\mathcal{Z}$  is the functional we want to minimize. We define the MMS below in Definition 2.1
The distribution space we work in is the metric Wasserstein space  $\mathbb{W}_2(\Omega) = (\mathcal{P}(\Omega), W_2)$ , where  $\Omega \subset \mathbb{R}^d$  is a convex compact set,  $\mathcal{P}(\Omega)$  is the set of probability distributions over  $\Omega$  and  $W_2$  is the Wasserstein distance over  $\mathcal{P}(\Omega)$  derived from the optimal transport problem with Euclidean cost:
$$
W _ {2} ^ {2} (\alpha , \beta) = \min  _ {T \text {s . t .} T _ {\sharp} \alpha = \beta} \int_ {\Omega} \| T (x) - x \| ^ {2} \mathrm {d} \alpha (x) \tag {4}
$$
where we assume that  $\partial \Omega$  is negligible and that the distributions are absolutely continuous.
Definition 2.1. Given  $\mathcal{Z}:\mathbb{W}_2(\Omega)\to \mathbb{R}$ , and starting from  $\rho_1^\tau \in \mathcal{P}(\Omega)$ , the Minimizing Movement Scheme (MMS) takes proximal steps for minimizing  $\mathcal{Z}$ . It is s given by
$$
\rho_ {k + 1} ^ {\tau} \in \arg \min  _ {\rho \in \mathcal {P} (\Omega)} \mathcal {Z} (\rho) + \frac {1}{2 \tau} W _ {2} ^ {2} (\rho , \rho_ {k} ^ {\tau}) \tag {5}
$$
The MMS (5) can be seen as a non-Euclidean implicit Euler step for following the gradient flow of  $\mathcal{Z}$ , and  $\rho_k^\tau$  converges to a minimizer of  $\mathcal{Z}$  under some conditions (see the end of this section).
So under the mentioned assumptions on  $\Omega$  and absolute continuity of the distributions, we have that Problems (3) are equivalent to the minimizing movement scheme (5):
Proposition 2.2. The distributions  $\rho_{k + 1}^{\tau} = (T_k^\tau)_{\sharp}\rho_k^\tau$ , where the functions  $T_{k}^{\tau}$  are found by solving (3) and  $\rho_1^\tau = \rho$  is the data distribution, coincide with the MMS (5) for  $\mathcal{Z} = \min_F\mathcal{L}(F,.)$
Proof. The minimizing movement scheme (5) is equivalent to taking  $\rho_{k + 1}^{\tau} = (T_k^\tau)_{\sharp}\rho_k^\tau$  where
$$
T _ {k} ^ {\tau} \in \arg \min  _ {T: \Omega \rightarrow \Omega} \mathcal {Z} \left(T _ {\sharp} \rho_ {k} ^ {\tau}\right) + \frac {1}{2 \tau} W _ {2} ^ {2} \left(T _ {\sharp} \rho_ {k} ^ {\tau}, \rho_ {k} ^ {\tau}\right) \tag {6}
$$
under conditions that guarantee the existence of a transport map between  $\rho_{k}^{\tau}$  and any other measure, and absolute continuity of  $\rho_{k}^{\tau}$  suffices, and the loss can ensure that  $\rho_{k + 1}^{\tau}$  is also absolutely continuous. Among the functions  $T_{k}^{\tau}$  that solve problem (6), is the optimal transport map from  $\rho_{k}^{\tau}$  to  $\rho_{k + 1}^{\tau}$ . To solve specifically for this optimal transport map, we have to solve the equivalent Problem
$$
T _ {k} ^ {\tau} \in \arg \min  _ {T} \mathcal {Z} \left(T _ {\sharp} \rho_ {k} ^ {\tau}\right) + \frac {1}{2 \tau} \int_ {\Omega} \| T (x) - x \| ^ {2} \mathrm {d} \rho_ {k} ^ {\tau} (x) \tag {7}
$$
Problems (6) and (7) have the same minimum value, but the minimizer of (7) is now an optimal transport map between  $\rho_{k}^{\tau}$  and  $\rho_{k + 1}^{\tau}$ . This is immediate from the definition (4) of the  $W_{2}$  distance. Equivalently minimizing first over  $F$  and then over  $T$  in (3), it follows from the definition of  $\mathcal{Z}$  that Problems (3) and (7) are equivalent, which concludes.
Since we solve Problems (3) over neural networks, their representation power shown by universal approximation theorems [13, 29] is important to get close to equivalence between (5) and (3), as we need to approximate an optimal transport map. We also know that the training of each module, if it is shallow, converges [5, 6, 34, 22, 18].
If  $\mathcal{Z}$  is lower-semi continuous then Problems (5) always admit a solution because  $\mathcal{P}(\Omega)$  is compact. If  $\mathcal{Z}$  is also  $\lambda$ -geodesically convex for  $\lambda > 0$ , we have convergence of  $\rho_k^\tau$  as  $k \to \infty$  and  $\tau \to 0$  to a minimizer of  $\mathcal{Z}$ , potentially under more technical conditions (see Appendix B). Even though a machine learning loss will usually not satisfy these conditions, this analysis offers hints as to why our method avoids in practice the problem of stagnation or collapse in performance of module-wise training as  $k$  increases, as we are making proximal local steps in Wasserstein space to minimize the loss. This convergence discussion also suggests taking  $\tau$  as small as possible and many modules.
2.3 Regularity result
As a secondary result, we show that Problem (3) has a solution and that the solution module  $T_{k}^{\tau}$  is an optimal transport map between its input and output distributions, which means that it comes with some regularity. [36] show that these networks generalize better and overfit less in practice. We assume that the minimization in  $F$  is over a compact set  $\mathcal{F}$ , that  $\rho_{k}^{\tau}$  is absolutely continuous, that  $\mathcal{L}$  is continuous and non-negative, that  $\Omega$  is convex and compact and that  $\partial \Omega$  is negligible.
Proposition 2.3. Problem (3) has a minimizer  $(T_k^\tau, F_k^\tau)$  such that  $T_k^\tau$  is an optimal transport map. And for any minimizer  $(T_k^\tau, F_k^\tau)$ ,  $T_k^\tau$  is an optimal transport map.
The proof is in Appendix C. OT maps have regularity properties under some boundedness assumptions. Given Theorem A.1 in Appendix A taken from [20],  $T_{k}^{\tau}$  is  $\eta$ -Hölder continuous almost everywhere and if the optimization algorithm we use to solve the discretized problem (2) returns an approximate solution pair  $(\tilde{F}_k^\tau, \tilde{T}_k^\tau)$  such that  $\tilde{T}_k^\tau$  is an  $\epsilon$ -optimal transport map, i.e.  $\| \tilde{T}_k^\tau - T_k^\tau \|_\infty \leq \epsilon$ , then we have (using the triangle inequality) the following stability property of the module  $\tilde{T}_k^\tau$ :
$$
\left\| \tilde {T} _ {k} ^ {\tau} (x) - \tilde {T} _ {k} ^ {\tau} (y) \right\| \leq 2 \epsilon + C \| x - y \| ^ {\eta} \tag {8}
$$
for almost every  $x, y \in \mathrm{supp}(\rho_k^\tau)$  and  $C > 0$ . Composing these stability bounds on  $T_k^\tau$  and  $\tilde{T}_k^\tau$  allows to get bounds for the composition networks  $G_k^\tau$  and  $\tilde{G}_k^\tau = \tilde{T}_k^\tau \circ \dots \circ \tilde{T}_1^\tau$ .
To summarize Section 2, the transport regularization makes each module more regular and it allows the modules to build on each other as  $k$  increases to solve the task, which is the property we desire.
3 Practical implementation
3.1 Multi-block modules
For simplicity, we presented in (2) the case where each module is a single residual block. However, in practice, we often split the network into modules that are made-up of many residual blocks each. We show here that regularizing the kinetic energy of such modules still amounts to a transport regularization, which means that the theoretical results in Propositions 2.2 and 2.3 still apply.
If each module  $T_{k}$  is made up of  $M$  residual blocks, i.e. applies  $x_{m + 1} = x_{m} + r_{m}(x_{m})$  for  $0 \leq m < M$ , then its total discrete kinetic energy for a single data point  $x_{0}$  is the sum of its squared residue norms  $\sum \| r_m(x_m)\|^2$ , since a residual network can be seen as a discrete Euler scheme for an ordinary differential equation [61] with velocity field  $r$ :
$$
x _ {m + 1} = x _ {m} + r _ {m} \left(x _ {m}\right) \longleftrightarrow \partial_ {t} x _ {t} = r _ {t} \left(x _ {t}\right) \tag {9}
$$
and  $\sum \| r_m(x_m)\|^2$  is then the discretization of the total kinetic energy  $\int_0^1\| r_t(x)\|^2\mathrm{d}t$  of the ODE. If  $\psi_{m}^{x}$  denotes the position of a point  $x$  after  $m$  residual blocks, then regularizing the kinetic energy of multi-block modules now means solving
$$
\left(T _ {k} ^ {\tau}, F _ {k} ^ {\tau}\right) \in \arg \min  _ {T, F} \sum_ {x \in \mathcal {D}} \left(L \left(F, T \left(G _ {k - 1} ^ {\tau} (x)\right) + \frac {1}{2 \tau} \sum_ {m = 0} ^ {M - 1} \| r _ {m} \left(\psi_ {m} ^ {x}\right) \| ^ {2}\right) \right. \tag {10}
$$
$$
\mathrm {s . t .} T = (\mathbf {i d} + r _ {M - 1}) \circ \ldots \circ (\mathbf {i d} + r _ {0}), \psi_ {0} ^ {x} = G _ {k - 1} ^ {\tau} (x), \psi_ {m + 1} ^ {x} = \psi_ {m} ^ {x} + r _ {m} (\psi_ {m} ^ {x})
$$
where  $G_{k}^{\tau} = T_{k}^{\tau} \circ \dots \circ T_{1}^{\tau}$  for  $1 \leq k \leq K$  and  $G_{0}^{\tau} = \mathrm{id}$ . We also minimize this sum of squared residue norms instead of  $\| T(x) - x\|^{2}$  (the two no longer coincide) as it works better in practice, which we assume is because it offers a more localized control of the transport. As expressed in (9), a residual network can be seen as an Euler scheme of an ODE and Problem (10) is then the discretization of
$$
\left(T _ {k} ^ {\tau}, F _ {k} ^ {\tau}\right) \in \arg \min  _ {T, F} \mathcal {L} \left(F, T _ {\sharp} \rho_ {k} ^ {\tau}\right) + \frac {1}{2 \tau} \int_ {0} ^ {1} \| v _ {t} \| _ {L ^ {2} \left(\left(\phi_ {t} ^ {\cdot}\right) _ {\sharp} \rho_ {k} ^ {\tau}\right)} ^ {2} \mathrm {d} t \tag {11}
$$
$$
\mathrm {s . t .} T = \dot {\phi_ {1}}, \partial_ {t} \phi_ {t} ^ {x} = v _ {t} \left(\phi_ {t} ^ {x}\right), \dot {\phi_ {0}} = \mathrm {i d}
$$
where  $\rho_{k + 1}^{\tau} = (T_k^\tau)_{\sharp}\rho_k^\tau$  and  $r_m$  is the discretization of vector field  $v_{t}$  at time  $t = m / M$ . Here, distributions  $\rho_{k}^{\tau}$  are pushed forward through the maps  $T_{k}^{\tau}$  which correspond to the flow  $\phi$  at time  $t = 1$  of the kinetically-regularized velocity field  $v_{t}$ . We recognize in the second term in the target of (11) the optimal transport problem in its dynamic formulation (15) from [9], and given the equivalence between the Monge OT problem (4) and the dynamic OT problem (15) in Appendix A, Problem (11) is in fact equivalent to the original continuous formulation (3), and the theoretical results in Section 2 follow immediately (see also the proof of Proposition 2.3 in Appendix C).
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 a complete forward pass on each batch but without a full backward pass, rather a backward pass that only updates the current module  $T_{k}^{\tau}$  and its auxiliary classifier  $F_{k}^{\tau}$ , meaning that  $T_{k}^{\tau}$  forwards its output to  $T_{k+1}^{\tau}$  immediately after it computes it. We call this parallel module-wise training. It is called decoupled greedy training in [8], which shows that combining it with batch buffers solves all three locking problems and allows a linear training parallelization in the depth of the network. We propose a variant of sequential module-wise training that we call multi-lap sequential module-wise training, in which instead of training each module for  $N$  epochs, we train each module from the first to the last sequentially for  $N / R$  epochs, then go back and train from the first module to
the last for  $N / R$  epochs again, and we do this for  $R$  laps. For the same total number of epochs and training time, and the same advantages (loading and training one module at a time) this provides a non-negligible improvement in accuracy over normal sequential module-wise training in most cases, as shown in Section 4. Despite our theoretical framework being that of sequential module-wise training, our method improves the test accuracy of all three module-wise training regimes.
3.3 Varying the regularization weight
The discussion in Section 2.2 suggests taking a fixed weight  $\tau$  for the transport cost that is as small as possible. However, instead of using a fixed  $\tau$ , we might want to vary it along the depth  $k$  to further constrain with a smaller  $\tau_{k}$  the earlier modules to avoid that they overfit or the later modules to maintain the accuracy of earlier modules. We might also want to regularize the network further in earlier epochs when the data is more entangled. We propose in Appendix D to formalize this varying weight  $\tau_{k,i}$  across modules  $k$  and SGD iterations  $i$  by using a scheme inspired by the method of multipliers to solve Problems (2) and (10). However, it works best in only one experiment in practice. The observed dynamics of  $\tau_{k,i}$  suggest simply finding a fixed value of  $\tau$  that is multiplied by 2 for the second half of the network, which works best in all the other experiments (see Appendix E).
We call our method TRGL for Transport-Regularized Greedy Learning. For the auxiliary classifiers, we use the architecture from DGL [7, 8], that is a convolution followed by an average pooling and a fully connected layer, which is very similar to that used by InfoPro [60], except for the Swin Transformer where we use a linear layer. We call vanilla greedy module-wise training with the same architecture but without our regularization VanGL, and we include its results in all tables for ablation study purposes. The code is available at github.com/block-wise/module-wise and implementation details are in Appendix E.
4.1 Parallel module-wise training