# Unveiling the Transport Dynamics of Neural Networks: A Least Action Principle for Deep Learning

Skander Karkar

► To cite this version:

Skander Karkar. Unveiling the Transport Dynamics of Neural Networks: A Least Action Principle for Deep Learning. Artificial Intelligence [cs.AI]. Sorbonne Université, 2023. English. NNT: 2023SORUS306. tel-04294849

HAL Id: tel-04294849

https://hal.science/tel-04294849

Submitted on 20 Nov 2023

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
Thèse de doctorat

# Unveiling the Transport Dynamics of Neural Networks:
A Least Action Principle for Deep Learning

Skander Karkar

soutenue publiquement le 13 octobre 2023 devant le jury composé de

<table><tr><td>Mme Laetitia Chapel</td><td>Rapportrice</td></tr><tr><td>Professeure, Institut Agro Rennes-Angers</td><td></td></tr><tr><td>M. Ievgen Redko</td><td>Rapporteur</td></tr><tr><td>Principal research scientist, Huawei</td><td></td></tr><tr><td>M. Alasdair Newson</td><td>Examinateur</td></tr><tr><td>Maître de conférences, Sorbonne Université</td><td></td></tr><tr><td>M. David Picard</td><td>Président</td></tr><tr><td>Senior research scientist, École des Ponts</td><td></td></tr><tr><td>M. Patrick Gallinari</td><td>Directeur</td></tr><tr><td>Professeur, Sorbonne Université</td><td></td></tr><tr><td>M. Alain Rakotomamonjy</td><td>Directeur</td></tr><tr><td>Principal researcher, Criteo</td><td></td></tr></table>
<figure><img src="image_2.png" /></figure>
Abstract

Residual connections are ubiquitous in deep learning, since besides residual networks [He et al., 2016b, He et al., 2016a] and their variants, they are also present in Transformers [Liu et al., 2021, Dosovitskiy et al., 2021]. The dynamic view of residual networks views them as analogous to a forward Euler scheme for an ordinary differential equation [Weinan, 2017]. We can then say that residual networks transport input points in space, time being represented by the depth of the network. This viewpoint has, for example, led to new architecture inspired by other numerical schemes for differential equations [Lu et al., 2018, Haber et al., 2019]. On the other hand, a bias of residual networks towards small perturbations of the input has been observed [Greff et al., 2016, Hauser, 2019, Jastrzebski et al., 2018, De and Smith, 2020, Chang et al., 2018b]. In the context of the dynamic view of residual networks mentioned above, this means a bias towards a small transport cost. In a first paper [Karkar et al., 2020], we experimentally verify that this bias is beneficial and should be encouraged and we show that forcing the network to approximate an optimal transport map by regularizing its transport cost improves its generalization ability and training stability. In a second paper [Karkar et al., 2023a], we show that applying this transport regularization to successive neural modules that don't back-propagate to each other amounts to following a gradient flow for minimizing the loss in distribution space, thus improving the performance of module-wise training, which consumes a lot less memory than end-to-end training. In a third paper [Karkar et al., 2023b], we propose a detector of adversarial and out-of-distribution samples that is based on the view of residual networks as discrete dynamical systems and show that transport regularization makes adversarial detection easier.
<figure><img src="image_3.png" /></figure>
# Contents

<table><tr><td>1</td><td>Introduction</td><td>1</td></tr><tr><td>1.1</td><td>Context</td><td>1</td></tr><tr><td>1.2</td><td>Motivation</td><td>1</td></tr><tr><td>1.3</td><td>Contributions</td><td>2</td></tr><tr><td>1.4</td><td>Outline</td><td>3</td></tr><tr><td>2</td><td>Background</td><td>5</td></tr><tr><td>2.1</td><td>Residual Networks</td><td>5</td></tr><tr><td>2.2</td><td>Residual Block Architectures</td><td>6</td></tr><tr><td>2.3</td><td>Interpretations of Residual Networks</td><td>9</td></tr><tr><td>2.4</td><td>The Iterative Refinement View of Residual Networks</td><td>10</td></tr><tr><td>2.5</td><td>Numerical Methods for Differential Equations</td><td>11</td></tr><tr><td>2.6</td><td>Residual Networks as Discrete Dynamical Systems</td><td>12</td></tr><tr><td>2.7</td><td>Neural Networks as Transport Systems</td><td>14</td></tr><tr><td>2.8</td><td>Optimal Transport Theory</td><td>15</td></tr><tr><td>2.9</td><td>Optimal Transport for the Study of Neural Networks</td><td>18</td></tr><tr><td>3</td><td>A Least Action Principle for the Training of Neural Networks</td><td>21</td></tr><tr><td>3.1</td><td>Introduction</td><td>21</td></tr><tr><td>3.2</td><td>Empirical Analysis of Transport Dynamics in Residual Networks</td><td>22</td></tr><tr><td>3.2.1</td><td>Transport Cost and Generalization in Image Classification</td><td>22</td></tr><tr><td>3.2.2</td><td>Visualizing Network Dynamics in 2 Dimensions</td><td>23</td></tr><tr><td>3.3</td><td>Method and Theory</td><td>25</td></tr><tr><td>3.3.1</td><td>General Setting</td><td>25</td></tr><tr><td>3.3.2</td><td>Formulation</td><td>26</td></tr><tr><td>3.3.3</td><td>Regularity</td><td>28</td></tr><tr><td>3.4</td><td>Practical Implementation</td><td>29</td></tr><tr><td>3.5</td><td>Experiments</td><td>30</td></tr><tr><td>3.5.1</td><td>Generalization</td><td>30</td></tr><tr><td>3.5.2</td><td>Transport Visualization</td><td>32</td></tr><tr><td>3.5.3</td><td>Stability</td><td>32</td></tr><tr><td>3.5.4</td><td>Training Dynamics</td><td>33</td></tr><tr><td>3.6</td><td>Discussion and Conclusion</td><td>34</td></tr><tr><td>4</td><td>Module-wise Training via the Minimizing Movement Scheme</td><td>37</td></tr><tr><td>4.1</td><td>Introduction</td><td>37</td></tr><tr><td>4.2</td><td>Related Work</td><td>40</td></tr><tr><td>4.3</td><td>Method and Theory</td><td>42</td></tr><tr><td>4.3.1</td><td>Formulation</td><td>42</td></tr><tr><td>4.3.2</td><td>Background on the Minimizing Movement Scheme</td><td>43</td></tr><tr><td>4.3.3</td><td>Link with the Minimizing Movement Scheme</td><td>44</td></tr><tr><td>4.3.4</td><td>Regularity Result</td><td>45</td></tr><tr><td>4.4</td><td>Practical implementation</td><td>46</td></tr><tr><td>4.4.1</td><td>Multi-block Modules</td><td>46</td></tr></table>
<table>
  <tbody>
    <tr>
      <td>4.4.2</td>
      <td>Solving the Module-wise Problems</td>
      <td>47</td>
    </tr>
    <tr>
      <td>4.4.3</td>
      <td>Varying the Regularization Weight</td>
      <td>47</td>
    </tr>
    <tr>
      <td>4.5</td>
      <td>Experiments</td>
      <td>48</td>
    </tr>
    <tr>
      <td>4.5.1</td>
      <td>Parallel Module-wise Training with Few Modules</td>
      <td>49</td>
    </tr>
    <tr>
      <td>4.5.2</td>
      <td>Ablation Study and Sensitivity to Hyperparameter τ</td>
      <td>51</td>
    </tr>
    <tr>
      <td>4.5.3</td>
      <td>Memory Usage of Parallel TRGL</td>
      <td>52</td>
    </tr>
    <tr>
      <td>4.5.4</td>
      <td>Training Time of Parallel TRGL</td>
      <td>53</td>
    </tr>
    <tr>
      <td>4.5.5</td>
      <td>Sequential Full Block-wise Training</td>
      <td>53</td>
    </tr>
    <tr>
      <td>4.5.6</td>
      <td>Avoiding Early Overfitting</td>
      <td>55</td>
    </tr>
    <tr>
      <td>4.6</td>
      <td>Discussion and Conclusion</td>
      <td>56</td>
    </tr>
    <tr>
      <td>5</td>
      <td>Adversarial Detection Through Transport Dynamics</td>
      <td>57</td>
    </tr>
    <tr>
      <td>5.1</td>
      <td>Introduction</td>
      <td>57</td>
    </tr>
    <tr>
      <td>5.2</td>
      <td>Related Work</td>
      <td>58</td>
    </tr>
    <tr>
      <td>5.3</td>
      <td>Method and Theory</td>
      <td>62</td>
    </tr>
    <tr>
      <td>5.3.1</td>
      <td>Detection</td>
      <td>62</td>
    </tr>
    <tr>
      <td>5.3.2</td>
      <td>Regularization</td>
      <td>63</td>
    </tr>
    <tr>
      <td>5.4</td>
      <td>Experiments</td>
      <td>67</td>
    </tr>
    <tr>
      <td>5.4.1</td>
      <td>Preliminary Experiments</td>
      <td>68</td>
    </tr>
    <tr>
      <td>5.4.2</td>
      <td>Detection of Seen Attacks</td>
      <td>70</td>
    </tr>
    <tr>
      <td>5.4.3</td>
      <td>Detection of Unseen Attacks</td>
      <td>72</td>
    </tr>
    <tr>
      <td>5.4.4</td>
      <td>Detection of Out-of-Distribution Samples</td>
      <td>73</td>
    </tr>
    <tr>
      <td>5.4.5</td>
      <td>Attacking the Detector</td>
      <td>75</td>
    </tr>
    <tr>
      <td>5.4.6</td>
      <td>Time Comparison</td>
      <td>76</td>
    </tr>
    <tr>
      <td>5.5</td>
      <td>Discussion and Conclusion</td>
      <td>76</td>
    </tr>
    <tr>
      <td>6</td>
      <td>Conclusion</td>
      <td>79</td>
    </tr>
    <tr>
      <td>6.1</td>
      <td>Summary</td>
      <td>79</td>
    </tr>
    <tr>
      <td>6.2</td>
      <td>Future Directions</td>
      <td>79</td>
    </tr>
    <tr>
      <td>6.2.1</td>
      <td>Adaptive Depth During Training</td>
      <td>79</td>
    </tr>
    <tr>
      <td>6.2.2</td>
      <td>Similarity of Neural Network Representations</td>
      <td>82</td>
    </tr>
    <tr>
      <td>6.2.3</td>
      <td>Optimal Transport between Different Dimensions</td>
      <td>82</td>
    </tr>
    <tr>
      <td>6.2.4</td>
      <td>Simulating Continuity Partial Differential Equations</td>
      <td>82</td>
    </tr>
    <tr>
      <td>6.3</td>
      <td>Conclusion</td>
      <td>83</td>
    </tr>
    <tr>
      <td colspan="2">References</td>
      <td>85</td>
    </tr>
    <tr>
      <td colspan="2">List of Acronyms</td>
      <td>111</td>
    </tr>
    <tr>
      <td colspan="2">List of Tables</td>
      <td>115</td>
    </tr>
    <tr>
      <td colspan="2">List of Figures</td>
      <td>117</td>
    </tr>
    <tr>
      <td colspan="2">A Background on Numerical Methods for Differential Equations</td>
      <td>119</td>
    </tr>
    <tr>
      <td>A.1</td>
      <td>The Euler Method</td>
      <td>119</td>
    </tr>
    <tr>
      <td>A.2</td>
      <td>Runge-Kutta Methods</td>
      <td>120</td>
    </tr>
    <tr>
      <td>A.3</td>
      <td>Adaptive Runge-Kutta Methods</td>
      <td>122</td>
    </tr>
  </tbody>
</table>
<table>
  <tbody>
    <tr>
      <td>B</td>
      <td>Background on Optimal Transport Theory</td>
      <td>125</td>
    </tr>
    <tr>
      <td>B.1</td>
      <td>Monge and Kantorovich Problems</td>
      <td>125</td>
    </tr>
    <tr>
      <td>B.2</td>
      <td>Duality</td>
      <td>128</td>
    </tr>
    <tr>
      <td>B.3</td>
      <td>The Wasserstein Space</td>
      <td>129</td>
    </tr>
    <tr>
      <td>B.4</td>
      <td>The Dynamic Formulation of Optimal Transport</td>
      <td>133</td>
    </tr>
    <tr>
      <td>B.5</td>
      <td>Regularity of Optimal Transport Maps</td>
      <td>136</td>
    </tr>
    <tr>
      <td>C</td>
      <td>Background on Gradient Flows</td>
      <td>139</td>
    </tr>
    <tr>
      <td>C.1</td>
      <td>Gradient Flows in Euclidean Space</td>
      <td>139</td>
    </tr>
    <tr>
      <td>C.2</td>
      <td>Gradient Flows in Metric Spaces</td>
      <td>140</td>
    </tr>
    <tr>
      <td>C.3</td>
      <td>Gradient Flows in the Wasserstein Space</td>
      <td>141</td>
    </tr>
    <tr>
      <td>D</td>
      <td>Background on Module-wise Training</td>
      <td>145</td>
    </tr>
    <tr>
      <td>D.1</td>
      <td>Additional Loss Terms</td>
      <td>145</td>
    </tr>
    <tr>
      <td>D.2</td>
      <td>Architectural Approaches</td>
      <td>146</td>
    </tr>
    <tr>
      <td>E</td>
      <td>Background on Adversarial Attacks</td>
      <td>149</td>
    </tr>
    <tr>
      <td>E.1</td>
      <td>Adversarial Attacks</td>
      <td>149</td>
    </tr>
    <tr>
      <td>E.1.1</td>
      <td>White-box Attacks</td>
      <td>149</td>
    </tr>
    <tr>
      <td>E.1.2</td>
      <td>Black-box Attacks</td>
      <td>150</td>
    </tr>
    <tr>
      <td>E.2</td>
      <td>Adversarial Detectors</td>
      <td>151</td>
    </tr>
    <tr>
      <td>E.2.1</td>
      <td>Statistical Approaches</td>
      <td>151</td>
    </tr>
    <tr>
      <td>E.2.2</td>
      <td>Computer Vision Approaches</td>
      <td>152</td>
    </tr>
    <tr>
      <td>F</td>
      <td>Additional Experiments from Section 3</td>
      <td>155</td>
    </tr>
    <tr>
      <td>F.1</td>
      <td>Implementation Details</td>
      <td>155</td>
    </tr>
    <tr>
      <td>F.2</td>
      <td>Transport Visualization</td>
      <td>155</td>
    </tr>
    <tr>
      <td>F.3</td>
      <td>Experiments with Fixed λ</td>
      <td>159</td>
    </tr>
    <tr>
      <td>F.4</td>
      <td>Impact on Weight Distribution of Classifier</td>
      <td>160</td>
    </tr>
    <tr>
      <td>F.5</td>
      <td>Additional Experiments in 2 Dimensions</td>
      <td>161</td>
    </tr>
    <tr>
      <td>F.5.1</td>
      <td>Comparison to Exact Wasserstein Distances</td>
      <td>161</td>
    </tr>
    <tr>
      <td>F.5.2</td>
      <td>Comparison to Batch Normalization</td>
      <td>161</td>
    </tr>
    <tr>
      <td>F.5.3</td>
      <td>Additional Transport Visualizations in 2 Dimensions</td>
      <td>163</td>
    </tr>
    <tr>
      <td>G</td>
      <td>Additional Experiments from Section 4</td>
      <td>169</td>
    </tr>
    <tr>
      <td>G.1</td>
      <td>Implementation Details</td>
      <td>169</td>
    </tr>
    <tr>
      <td>G.2</td>
      <td>Additional Experiments</td>
      <td>169</td>
    </tr>
    <tr>
      <td>G.3</td>
      <td>Memory Usage</td>
      <td>172</td>
    </tr>
    <tr>
      <td>G.4</td>
      <td>Sensitivity to Hyperparameter τ</td>
      <td>173</td>
    </tr>
    <tr>
      <td>H</td>
      <td>Additional Experiments from Section 5</td>
      <td>175</td>
    </tr>
    <tr>
      <td>H.1</td>
      <td>Implementation Details</td>
      <td>175</td>
    </tr>
    <tr>
      <td>H.2</td>
      <td>Detection of Seen Attacks</td>
      <td>175</td>
    </tr>
    <tr>
      <td>H.3</td>
      <td>Detection of Unseen Attacks</td>
      <td>176</td>
    </tr>
    <tr>
      <td>H.4</td>
      <td>Detection Rate of Successful Adversarial Samples</td>
      <td>177</td>
    </tr>
    <tr>
      <td>H.5</td>
      <td>False Positive Rate</td>
      <td>178</td>
    </tr>
  </tbody>
</table>
<table><tr><td>H.6 AUROC</td><td>179</td></tr><tr><td>H.7 Detection of Out-of-Distribution Samples</td><td>179</td></tr></table>
# 1 Introduction

## 1.1 Context

Machine learning is the prevailing current approach in the field of artificial intelligence. It focuses on designing systems that automatically learn from observed data [Russell and Norvig, 2020]. In the past twenty years, deep learning, i.e. the use of deep artificial neural networks, has emerged as a crucial advancement in machine learning [Goodfellow et al., 2016]. It has become the state-of-the-art in nearly all tasks, from classification to generation, in the fields of computer vision [Krizhevsky et al., 2012], natural language [Wu et al., 2016], sound [Purwins et al., 2019], speech [Mehrish et al., 2023] and video [Sharma et al., 2021] processing, and has been applied successfully to other tasks such as dynamical system forecasting [de Bezenac et al., 2018], biology [Jumper et al., 2021] and reinforcement learning [Silver et al., 2016].

In all these applications, the neural network architectures used rely on a set of heuristics, techniques and components that are omnipresent. These components, and especially the interactions between them, are not yet fully understood theoretically. They include back-propagation [Rumelhart et al., 1986] and stochastic gradient descent [Bottou, 2012, Amari, 1967], or a variant of it [Bottou et al., 2018], for training the model, non-linear activations (ReLU [Fukushima, 1975] very often), normalization layers (most often batch normalization [Ioffe and Szegedy, 2015] or layer normalization [Ba et al., 2016]) and the importance of a good initialization [Saxe et al., 2014, He et al., 2015]. Among these ever-present elements is the architectural component known as a skip (or residual) connection [He et al., 2016b, He et al., 2016a].

Indeed, residual networks and their variants are a very popular neural network architecture, especially in computer vision. Their main building block, the residual connection, has become ubiquitous in state-of-the-art architectures such as Transformers used in natural language processing.

## 1.2 Motivation

Many works have tried to explain the success of residual networks through both empirical and theoretical explorations, and to propose new improvements (architectures, initializations, training procedures...) suggested by these explanations.

A point that stands out is that it is precisely their initialization as a perturbation of the identity function that biases residual networks towards finding good solutions to learning problems that only minimally move the inputs.

We position ourselves in this field of research that tries to explain the importance of skip connections and then improve the performances of networks that use them.
1.3 Contributions

The particular point of view of the functioning of residual networks that we
consider is that of discrete dynamical systems. Indeed, residual networks can
be viewed as a discrete Euler scheme for a differential equation. We further
link this to the theory of optimal transport in its dynamical formulation. This
allows us to see residual networks as discrete transport systems that move their
input points through space in time, time being represented by the depth, to
disentangle them before they are classified by a linear classifier.

Through this understanding, we propose a regularization that is particularly
well-adapted to residual networks and that improves their training stability and
generalization ability. Indeed, from this viewpoint, the bias towards small resid-
uals means a bias towards a small transport cost. Since we verify that this bias
is beneficial, a natural idea that follows is to regularize residual networks during
training by penalizing their transport cost, which is the distance by which the
input points travel through the residual blocks, thereby reinforcing the desirable
bias towards remaining close to the identity function. This regularization forces
the solution network to be an optimal transport map in the sense of optimal
transport theory, thus giving it some regularity. In practice, this is showcased by
increased stability (less variance in the results and the ability to train without
batch normalization for example) and better accuracy at test time.

This transport regularization has other interesting properties that we exploit for
two applications: module-wise training and adversarial detection.

Firstly, the quantity we regularize is in fact a distance between probability
distributions (the Wasserstein distance). Therefore, this regularization allows to
split neural networks into modules and to train them module-wise, i.e. without
back-propagation between the modules, in a theoretically sound way and with
good performances. This type of training is interesting in practice in constrained
settings as it does without many of the back-propagated gradients of end-to-end
back-propagation, thus consuming less memory. The transport regularization
keeps the early modules from overfitting and makes the modules take small
proximal steps to minimize the loss in the Wasserstein metric. This allows the
modules to build upon each other in accuracy without back-propagation between
them, which is exactly the property desired in module-wise training.

Secondly, The transport regularization also has the particularity of having
a value for each input image. The transport statistics therefore allow us to
study the behavior of the network on a particular image. We can therefore use
these transport statistics that describe the trajectory of the input in space for
the purpose of detecting adversarial attacks. We propose a novel detector of
adversarial and out-of-distribution samples that tells clean inputs from abnormal
ones by comparing the discrete vector fields they follow throughout the network's
layers before the final classification layer. We also show that minimizing the
transport cost gives the network more regularity only on the support of the
clean data distribution. This makes the network's activations on clean samples
more distinguishable from those on abnormal samples, and thus improves the performance of adversarial detection methods that use the internal embeddings as inputs.

## 1.4 Outline

In Section 2 we present the related work and necessary background on residual networks, their dynamic interpretation and the advances that followed from it, numerical methods for ordinary differential equations and optimal transport theory.

In Section 3, we present the work done in [Karkar et al., 2020]. We first verify experimentally that the bias of residual networks towards minimal transformations is beneficial, and that it means a bias of residual networks towards being good approximators of optimal transport maps. We then formalize the idea of encouraging small transport cost by linking this bias to optimal transport theory in its dynamic formulation and implement it as a transport regularization. This means that we seek to retrieve a network that solves the task at hand while moving the input points as little as possible. We theoretically prove that this leads to regular networks. Experimentally, our method improves the generalization of residual-type networks, especially in small data and overfitting regimes. It also improves training stability, allowing for example to train deep networks without batch normalization.

In Section 4, we present the work done in [Karkar et al., 2020a], where we show that applying this transport regularization to successive neural modules that don't back-propagate to each other amounts to following a gradient flow for minimizing the loss in distribution space. This offers a new formulation and method for greedily training neural networks module-wise, i.e. splitting the network into modules and training these modules without back-propagation between them. This requires less memory than standard end-to-end training and has therefore been used in constrained settings such as training on mobile devices. Experimentally, regularizing the transport cost of each module leads to better performances of these networks by avoiding the problem of stagnation or collapse in accuracy along the depth observed in greedily-trained modules, whereby earlier modules overfit and deeper ones stop improving the test accuracy or even degrade it.

In Section 5, we present the work done in [Karkar et al., 2020b], where we propose a detector of adversarial and out-of-distribution samples that uses the spatiotemporal dynamics of the inputs throughout the network to detect abnormal inputs. We also show that the increased regularity guarantee given by the transport regularization only on the support of the data makes detection of adversarial attacks, that tend to lie outside the data manifold, easier.

In Section 6, we conclude and discuss future avenues of research. The main idea for future work is to extend the analogy between neural networks and numerical schemes for differential equations to include adaptive schemes. Indeed,
architectures that imitate schemes that are of higher order than the Euler
method have been proposed, but the possibilities and advantages of including
adaptive methods, which rely on imbedded schemes of different orders, have
not been sufficiently explored. These possibilities include ensembling different
forward passes at test time to improve performance, and adaptively adjusting
the depth of the network during training to the complexity of the task at hand.
# 2 Background

We present in this section the background and notions that are necessary to all the following sections. The notions that are specific only to a particular section are presented in it, and more detailed background can be found in the Appendix. Here, we define residual networks, stress their importance in modern deep learning and present some interpretations of their functioning, focusing on the iterative refinement and dynamic interpretations. We then discuss recent advances and ideas that were inspired by these interpretations. Finally, we present the essential notions of optimal transport theory that we will use later.

## 2.1 Residual Networks

Residual networks (ResNets) [He et al., 2016b, He et al., 2016a] and their variants such as ResNeXt [Xie et al., 2017] and WideResNet [Zagoruyko and Komodakis, 2016] have become an important architecture in deep learning, especially for computer vision. They allowed for training very deep models: up to 1000 layers, while training compositional (i.e. non-residual) networks deeper than 20 layers remains challenging [Jastrzebski et al., 2018, Balduzzi et al., 2017]. Residual networks and their variants reach state-of-the-art performances on many computer vision tasks such as image classification [Xie et al., 2017, Mahajan et al., 2018, Wightman et al., 2021], unsupervised domain translation [Zhu et al., 2017], image segmentation [Pohlen et al., 2017, Chen et al., 2017a] and action recognition [Tran et al., 2018]. They have also been used with success in speech and natural language processing applications [et al., 2016, Xiong et al., 2017, van den Oord et al., 2016].

Furthermore, residual connections are a basic building block in many state-of-the-art architectures such as EfficientNet [Tan and Le, 2019] and MobileNetV2 [Sandler et al., 2018], which use a so-called inverted residual block. Residual connections have also been added to prior architectures such as Inception to improve their performances [Szegedy et al., 2017], and are used in some U-Net architectures [Janner et al., 2022]. If we call a *residual network* any neural network whose architecture is made up of successive residual blocks (preceded by an initial encoding, separated by pooling layers and followed by a classification head), then vision Transformers such as Swin [Liu et al., 2021] and ViT [Dosovitskiy et al., 2021], video Transformers such as TimeSFormer [Bertasius et al., 2021], and natural language processing Transformers such as GPT [Radford et al., 2018] can also be considered residual networks, with attention and multilayer perceptrons instead of convolutions in their residual functions.

A *residual block* (ResBlock) is an architectural component in a neural network that has as input $x$ (an embedding at a certain depth) and outputs $x + r(x)$. We will call the function $r$ associated with a residual block its *residual function* (or *forcing function*), and $r(x)$ the *residue*. Consecutive residual blocks with
residual functions $r_m$ applied to an input $x_0$ can then be written

$$x_{m+1} = x_m + r_m(x_m) \qquad (1)$$

The addition of $x$ to $r(x)$ is called a *skip connection* or a *residual connection*. More generally, a skip connection saves the value of the input to some layer and uses it later in the network. The residual point of view is different from the traditional one where the data is transformed compositionally by applying $f_1 \circ f_2 \circ f_3 \circ \dots \circ f_m$ to the input. In a residual architecture, the input is transformed in an additive manner by adding residues to it.

This rewriting of course does not change the representation power of the network (i.e. the functions it is able to approximate). Networks with layers that apply $x_{m+1}=f_m(x_m)$ where the input and output of $f_m$ have the same dimension can be reformulated as residual blocks by taking as residual functions $r_m(x_m)=f_m(x_m)-x_m$ so that $x_{m+1}=x_m+r_m(x_m)$. The theory and implementation developed apply therefore directly to networks with stages made up of such layers, which includes for examples VGG networks [Simonyan and Zisserman, 2014]. The particularity of residual blocks in residual networks and Transformers is then that they are initialized to functions of the form $id + r$, which confers certain biases to the training procedure and to the solution that will be found. It is these biases that we aim to uncover and exploit.

## 2.2 Residual Block Architectures

Typically the residual function of a residual block used for a vision task contains two or three convolutions with activations and batch normalisations. The *basic residual block* architecture is in Figure 1 below.

<figure><img src="image_4.png" /><figcaption>Figure 1: Basic residual block architecture proposed in [He et al., 2016a]. BN: batch normalization, ReLU: $x \mapsto (x)_+$. Conv: convolution. d: number of channels.</figcaption></figure>

The other common residual block architecture is called the *bottleneck residual block* [He et al., 2016b, He et al., 2016a]. While the convolutions in the basic block keep the same number of channels, the bottleneck block has three convolutions, the first of which (a $1 \times 1$ convolution) reduces the number of channels. The
third convolution (also 1 × 1) restores it to that of the block's input. This structure, applied across many paths, is present in the residual block used in the ResNeXt architecture. The bottleneck block and the ResNeXt block are illustrated in Figure 2 below.

<figure><img src="image_5.png" /><figcaption>Figure 2: Bottleneck residual block (left) and ResNeXt block (right) architectures. From [Xie et al., 2017]. A convolutional layer is shown as (number of input channels, filter size, number of output channels).</figcaption></figure>

Finally, we mention the *inverted residual block*. It was proposed initially for the MobileNetV2 [Sandler et al., 2018] architecture for efficiency reasons and has been used since for mobile-optimized networks. It uses an inverse structure to the bottleneck block, as the input and output now have few channels, and the first 1 × 1 convolution increases the number of channels before the third convolution reduces it back to that of the input. It also uses a depth-wise convolution for the second 3 × 3 convolution, which reduces the number of parameters. This architecture is compared to the bottleneck architecture in Figure 3 below.

<figure><img src="image_6.png" /><figcaption>Figure 3: Bottleneck residual block (left) and inverted residual block (right) architectures. From [Sandler et al., 2018]. The thickness indicates the number of channels.</figcaption></figure>

Residual blocks can be separated by layers that change the dimensions of the embedding (for example pooling layers and patch merging layers). If the same embedding dimension is kept, we call the network a *single representation* ResNet.
If the embedding dimension changes as in the original ResNet of [He et al., 2016a], we call each group of residual blocks that keep the same dimension a *residual stage*. In Figure 4, we represent a residual stage containing three consecutive basic residual blocks.

<figure><img src="image_7.png" /><figcaption>Figure 4: Three consecutive residual blocks.</figcaption></figure>

Below in Figure 5 is the Swin Transformer architecture with 4 residual stages separated by patch merging. We see that the typical *Transformer block* is made up of two residual blocks, the first containing self-attention and the second containing a multilayer perceptron. The Transformer block in ViT (Vision Transformer, [Dosovitskiy et al., 2021]) is very similar.

<figure><img src="image_8.png" /><figcaption>Figure 5: Swin Transformer architecture with 4 residual stages from [Liu et al., 2021]. MLP: multi-layer perception. LN: layer normalization. W-MSA: window multi-head self-attention. SW-MSA: shifted window multi-head self-attention.</figcaption></figure>
## 2.3 Interpretations of Residual Networks

The initial intuition behind the design of the ResBlock in [He et al., 2016b] is that adding more layers to a network should not decrease its performance, as the deeper network contains the shallower one by simply adding identity mappings.

However, increased depth in compositional networks is observed to cause a decrease in training accuracy. The residual block architecture was proposed as a way of forcing deep networks to only learn perturbations from identity, as learning the zero function is easier than learning the identity function. This also helped with the problems of vanishing and exploding gradients (exacerbated in deeper networks as gradients multiply).

Many perspectives try to further explain how residual connections work and why they perform so well. Once an interesting explanation is found, it is often possible to deduce from it new methods, architectures and improvements to residual networks and to neural networks in general.

In [Balduzzi et al., 2017], the authors identify the problem of *shattered gradients*. They show, experimentally and theoretically, that correlations between gradients of deeper layers at initialization decay exponentially with depth in compositional networks, making them resemble white noise. More precisely, they consider the gradient of the network's output with respect to a fixed neuron's output for different inputs. Because of the chain rule, these gradients are strongly linked (via terms that are independent of the network's architecture) to the gradients with respect to the weights computed during optimization. This decay in correlations between gradients of deeper layers makes training difficult as it leads to averaging white-noise-like gradients across minibatches, hinders optimization algorithms that assume gradients at nearby points are similar (e.g. momentum-based and accelerated methods) and causes the effect of a neuron on the output to become increasingly unstable as the depth increases. The authors show that, in contrast, gradients in residual networks behave more like brown noise and that the correlations between them decrease sublinearly with depth (when batch normalization is used). From this, they propose a new initialization that makes training deep non-residual networks easier.

In the ensemble view [Veit et al., 2016, Huang et al., 2018], residual networks are thought to learn an exponential ensemble of shallower models. For a residual network of depth 3, the output $x_3$ can indeed be written

$$x_3 = x_0 + r_0(x_0) + r_1(x_0 + r_0(x_0)) + r_2(x_0 + r_0(x_0) + r_1(x_0 + r_0(x_0)))$$

The input $x_0$ then has many paths to reach the final classification layer. If this classification layer is linear, then its output can be seen as the sum of its outputs on many different shallower learned representations of the input $x_0$, starting from $x_0$ itself to the deepest representation $r_2(x_0 + r_0(x_0) + r_1(x_0 + r_0(x_0)))$. [Huang et al., 2018] combine this point of view with boosting theory [Freund and Schapire, 1997] to propose an algorithm for sequentially training residual networks, as we do in Section 4.
## 2.4 The Iterative Refinement View of Residual Networks

A view of residual network that is of use to us is that of unrolled iterative estimation, where residual blocks are thought to iteratively refine representations instead of learning new ones [Greff et al., 2016, Liao and Poggio, 2016, Hauser, 2019, De and Smith, 2020]. This view essentially verifies that residual blocks, especially deeper ones, do remain close to the identity function. Paper [Jastrzebski et al., 2018] explores this view further, calling it iterative refinement (or iterative inference). The authors rewrite the loss L of a residual network with M residual blocks by applying recursive Taylor expansions:

$$ 
\begin{aligned}
L(x_M) &= L(x_{M-1} + r_{M-1}(x_{M-1})) \\
&= L(x_{M-1}) + \nabla L(x_{M-1}) \cdot r_{M-1}(x_{M-1}) + \mathcal{O}(r_{M-1}^2(x_{M-1})) \\
&= L(x_M) + \sum_{m \le i < M} \nabla L(x_i) \cdot r_i(x_i) + \mathcal{O}(r_i^2(x_i))
\end{aligned}
$$

which holds for any $0 \le m < M$. This is a good approximation if the magnitude of the residual functions $r_i$ is small. Thus, for these blocks, the training can be thought of as in part minimizing the terms $\nabla L(x_i) \cdot r_i(x_i)$, which can be done by making $r_i(x_i)$ point in the opposite half-space to $\nabla L(x_i)$. The block operations $x_i + r_i(x_i)$ are therefore moving the input towards the half-space of $-\nabla L(x_i)$. A residual block is then said to perform iterative refinement or inference if, on average on the dataset, its residual function $r_i$ has a small magnitude and satisfies $\nabla L(x_i) \cdot r_i(x_i) < 0$. A residual block whose residual function has a high magnitude is said to learn a new representation of the data, which is slowly refined by later blocks performing iterative inference.

It is observed in [Greff et al., 2016, Hauser, 2019, Jastrzebski et al., 2018, Chang et al., 2018b] that only few early residual blocks in each residual stage learn a new representation and that the deeper blocks in a residual stage perform iterative refinement according to the definition above. Figure 6 shows the ratio $\frac{\|r_i(x_i)\|}{\|x_i\|}$ and the cosine loss $\frac{\nabla \mathcal{L}(x_i) \cdot r_i(x_i)}{\|\nabla \mathcal{L}(x_i)\| \|r_i(x_i)\|}$ on CIFAR10 of 4 ResNets.

<figure><img src="image_9.png" /><figcaption>Figure 6: Ratio of residue norm to input norm (first row) and cosine loss (second row) of 4 residual architectures on CIFAR10 from [Jastrzebski et al., 2018].</figcaption></figure>
We see that the peaks in the ratio come in fact in the early residual blocks of a residual stage, especially on the first architecture on the left. This means that residual blocks are indeed biased towards minimally modifying their input. Furthermore, [Zhang and Wynter, 2018, Zhang et al., 2019] add a fixed positive weight h to the residual function, i.e. consider residual blocks of the form $x + h r(x)$, and find that taking a small positive h < 1 makes residual networks more stable. This indicates that making the residual blocks even closer to identity is beneficial.

In [Jastrzebski et al., 2018], two suggested modifications to residual networks follow from this finding that later residual blocks all perform similar iterative refinement. First, they propose to unroll the last block at inference time (i.e. apply it many times) to improve the accuracy of the model on borderline inputs, that is, inputs that are very close to the decision boundary. Secondly, they propose to make the later residual blocks share the same parameters to reduce memory footprint.

The fact that residual blocks only minimally build upon each other instead of learning new representations also explains why stochastic depth training [Huang et al., 2016, Hayou and Ayed, 2021], i.e dropping a random subset of layers at each training epoch and replacing them with the identity function, is particularly effective on residual networks. It increases their test accuracy as it functions as a regularization [Hayou and Ayed, 2021], and of courses reduces training time. In [Huang et al., 2016], stochastic depth training allowed for training residual networks that have more than 1200 layers on CIFAR10 while still improving test accuracy, i.e without overfitting.

Finally, [De and Smith, 2020] also find that the bias towards small residuals is beneficial and allows one to successfully train deeper networks. They further show that batch normalization encourages this bias by downscaling the variance of the residue at initialization compared to the variance of the identity branch of the residual block. From this they develop an initialization that allows to train deep residual networks without batch normalization. Our transport regularization in Section 3 also allows to train deep residual networks without batch normalization.

## 2.5 Numerical Methods for Differential Equations

Before introducing the dynamic interpretation of residual networks, we need to introduce numerical schemes for ordinary differential equations (ODEs). Given a time-dependent vector field $v_t: \mathbb{R}^d \to \mathbb{R}^d$, consider the *initial value problem* of finding a function $x: I \to \mathbb{R}^d$ that satisfies

$$x'(t) = v_t(x(t)) \qquad (2)$$

with initial condition $x(t_0)=x_0$, where I is the interval $[t_0, t_0+T]$.
Equivalently, the integral form of this problem is finding $x$ that satisfies

$$x(t) = x_0 + \int_{t_0}^{t} v_t(x(t)) dt \qquad (3)$$

Finding closed-formed solutions to this problem can be difficult. We wish then to approximate the values $x$ takes at certain time points. Given a subdivision $t_0<t_1<..<t_N=t_0+T$ of $I$, we denote the step-sizes $h_n:=t_{n+1}-t_n$ for $0\le n<N$ and define $h_{\max}:=\max h_n$. The explicit Euler method is a numerical approximation method that approximates the values $(x(t_n))_n$ through the values $(x_n)_n$ given by

$$x_{n+1} = x_n + h_n v_{t_n}(x_n) \qquad (4)$$

starting from the initial condition $x_0$. This method simply comes from a first-order Taylor expansion of $x$, or from a left-hand rectangle method for approximating the integral in (3). The right-hand rectangle method for approximating the integral in (3) leads to the *implicit Euler method*:

$$x_{n+1} = x_n + h_n v_{t_n}(x_{n+1}) \qquad (5)$$

A fixed-point method [Shashkin, 1991] or the Newton-Raphson method [Ypma, 1995] can be used to solve (5) for $x_{n+1}$.

A numerical method converges if

$$\max_n \|x(t_n)-x_n\| \xrightarrow{h_{\max}\to 0} 0$$

where $\max_n \|x(t_n)-x_n\|$ is the *global error* of the method. Euler's method converges if $v$ is $C^0$ and Lipschitz in $x$. If $v$ is also $C^1$ then it converges with speed $O(h_{\max})$. See Appendix A.1 for more details.

## 2.6 Residual Networks as Discrete Dynamical Systems

The view we consider is the dynamic view. It views residual network as an Euler discretization of a differential equation on the interval $I = [0, 1]$.

$$x_{m+1} = x_m + r_m(x_m) \longleftrightarrow x'(t) = v_t(x(t)) \qquad (6)$$

where $r_m$ approximates the velocity field $v_t$ at time $t = t_m$ and we suppose that the time-steps are learned inside $r_m$. This observation was first made in [Weinan, 2017] and allows to draw on the rich theory of differential equations and on the numerical methods developed to simulate them. New network architectures were designed by considering other numerical schemes for the approximation of differential equations such as linear multi-step methods [Lu et al., 2018], implicit-explicit methods [Haber et al., 2019], and Runge-Kutta methods [Zhu et al., 2022]. We present some of these architectures below.

The IMEX architecture [Haber et al., 2019] considers a parametrized differential equation $x'(t) = v(x(t); \theta(t))$ where $v$ is a neural network parametrized by $\theta$
and discretizes it using the implicit-explicit numerical method [Ascher et al., 1995, Ascher et al., 1997, Ascher and Petzold, 1998]. This method rewrites the differential equation as follows $x'(t) = v(x(t); \theta(t)) + Ax(t) - Ax(t)$ where $A$ is a matrix to be chosen or learned. The method then uses an explicit Euler scheme for the first part $v(x(t); \theta(t)) + Ax(t)$ and an implicit Euler scheme for the second part $Ax(t)$. This leads to the following numerical step and block architecture

$$x_{n+1} = (\mathbf{Id} + hA)^{-1}(x_n + hAx_n + hv(x_n, \theta_n)) \quad (7)$$

The matrix $A$ is a trainable group convolution and the left-multiplication by an inverse matrix is done in the Fourier domain. According to [Haber et al., 2019], this architecture has two advantages in theory: the increased stability of implicit methods, and the global field of view connecting all pixels thanks to the multiplication by $(\mathbf{Id} + hA)^{-1}$, since this matrix is dense. Experiments on tasks where the field of view is important such as image segmentation [Minaee et al., 2022] and depth estimation [Masoumian et al., 2022] confirm this second advantage over ResNets.

[Lu et al., 2018] build architectures called LM-ResNet and LM-ResNeXt that mimic linear multi-step numerical schemes [Dahlquist, 1963], while FractalNet [Larsson et al., 2017] and RKCNN [Zhu et al., 2022] mimic Runge-Kutta numerical methods (see Appendix A.2 for a definition of these methods). PolyNet [Zhang et al., 2017b] use a polynomial approximation to the implicit Euler method (5) to design their architecture. The increased stability of implicit methods allows them to increase the step-size and therefore to reduce the depth of the network.

Forward passes that imitate partial differential equations were also proposed in [Ruthotto and Haber, 2018]. The stability of the forward pass was also studied through the lens of dynamical systems, leading to considering particular forms of differential equations to imitate such as Hamiltonian systems [Haber and Ruthotto, 2017], and to enforcing particular conditions on the weight matrices [Ciccone et al., 2018].

Reversibility, which means that the dynamics of a differential equation can also be simulated backwards in time, is an important notion in the study of dynamic systems. The dynamic view was therefore linked to the possible reversibility of residual networks in [Chang et al., 2018a]. The possibility of making residual networks reversible is one of their advantages. The reversible residual network called RevNet designed in [Gomez et al., 2017] allows to reconstruct the activations of a layer from those of the next layer. It is then possible to perform back-propagation without storing the activations in memory, except for those of a handful of non-reversible layers. These residual networks therefore have an activation storage complexity that is independent of depth and an order of magnitude lower than that of non-reversible networks of equal size. Reversible networks called i-RevNets [Jacobsen et al., 2018] are an extension of these RevNets.
Some of the architectures discussed above are illustrated in Figure 7 below.

<figure><img src="image_10.png" /><figcaption>Figure 7: Some architectures inspired by numerical schemes. From [Lu et al., 2018].</figcaption></figure>

## 2.7 Neural Networks as Transport Systems

The dynamic view also allows one to consider that neural networks are transporting their inputs in space, with depth representing time, before a linear classifier is applied (if the task at hand is classification). Here, we discuss other works that share this point of view.

NeuralODE [Chen et al., 2018, Ruiz-Balet and Zuazua, 2021, Yan et al., 2020] take the continuous limit of a residual network and parameterize and learn the velocity field of an ordinary differential equation that is then numerically integrated to get the output. More precisely, since a residual network is analogous to an Euler scheme, taking its limit as the step size goes to zero suggests directly learning the continuous velocity field $v_t$ that moves the input points in space so that they land in a configuration that solves the task. So if $L$ is the loss, and $x_0$ is the input, we look for a vector field $v_t$ (in practice a neural network parametrized by $\theta$) such that the initial value problem $x'(t) = v_t(x_t)$ on the interval $[0, 1]$ with initial condition $x(0) = x_0$ has a solution curve $x$ that ends up at a value $x(1)$ that has a small loss. So we minimize (through gradient descent) in $\theta$

$$L(x(1)) = L \left( x_0 + \int_0^1 v(t, x(t); \theta) dt \right) \approx L(\text{Solver}(v(.,.; \theta), x_0)) \quad (8)$$

where Solver applies a numerical solver such as Euler to the initial value problem and returns an approximation of $x(1)$. If the task is classification, a linear classifier is applied to $x(1)$ before the loss is computed and an embedding might precede going through the numerical solver. The classifier and encoder are then jointly learned with the dynamics network $v$.

Differentiating the loss (8) in $\theta$ can be done in two ways. The first is to discretize then differentiate, i.e. to differentiate $L(\text{Solver}(v(.,.; \theta), x_0))$ with respect to
$\theta$ by simply using automatic differentiation to back-propagate through the operations of the numerical solver (which are all differentiable). The second is to differentiate then discretize, i.e. to use the so-called adjoint state method, which can be thought of as an instantaneous analog of the chain rule, to differentiate $L(x_0 + \int_0^1 v(t, x(t); \theta) dt)$ in $\theta$. The formula for this gradient includes an integral that is then also discretized through a numerical solver.

In [Li et al., 2018], the parameter $\theta$ in (8) is allowed to depend on $t$ and the resulting problem is looked at as the optimal control problem [Athans and Falb, 2007] of finding a control function $\theta(t):[0, 1] \to \Theta$ that is going to guide the solution curve $x$ to the ordinary differential equation $x'(t) = v(t, x(t); \theta(t))$ towards an ending point $x(1)$ with a small loss. The authors then use Pontryagin's maximum principle (PMP, [Pontryagin, 1987]) to state the necessary optimality conditions for this optimal control problem, and then use the method of successive approximations (MSA, [Chernousko and Lyubushin, 1982]) to discretize and numerically solve for these optimality conditions. This leads to an alternative training algorithm for residual networks.

To improve adversarial robustness, [Wang et al., 2019a] add a diffusion term to the transport/continuity equation (equation (11) in Section 2.8 below) that a residual network approximates, which amounts to training an ensemble of Gaussian-perturbed residual networks.

## 2.8 Optimal Transport Theory

We quickly introduce here the theory of optimal transport which we will use in the following sections to analyze residual networks as transport maps. More details are in Appendix B. Optimal transport deals with the problem of displacing mass from one configuration to another while expending the least amount of effort possible. This minimal effort required to move mass from one distribution to another is actually a distance between distributions that we can equip the space of probability distributions with to get the so-called Wasserstein space. Let us first define the push-forward measure:

**Definition 1.** For a function $T: X \to Y$ and a measure $\mu$ on $X$, the push-forward of $\mu$ by $T$ is the measure $T_{\#}\mu$ on $Y$ defined by $T_{\#}\mu(A) = \mu(T^{-1}(A))$ for every measurable set $A \subset Y$.

The metric Wasserstein space $W_2(\Omega) := (P(\Omega), W_2)$, with $\Omega$ a convex and compact subset of $\mathbb{R}^d$, is the space $P(\Omega)$ of probability measures over $\Omega$, equipped with the distance $W_2$ given by the solution to the Kantorovitch problem

$$W_2^2(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\Omega \times \Omega} \|x - y\|^2 d\gamma(x, y) \quad (9)$$

where $\Pi(\mu, \nu)$ is the set of probability distribution over $\Omega \times \Omega$ with first marginal equal to $\mu$ and second marginal equal to $\nu$, i.e

$$\Pi(\mu, \nu) = \{\gamma \in \mathcal{P}(\Omega \times \Omega) \mid \pi_{1\#}\gamma = \mu, \pi_{2\#}\gamma = \nu\}$$
where $\pi_1(x, y) = x$ and $\pi_2(x, y) = y$ are projections. Here, $\gamma \in \Pi(\mu, \nu)$ is called a transport plan between $\mu$ and $\nu$, and $\gamma(x, y)$ indicates how much mass from position $x$ should be sent to position $y$ (see Figure 8 below). This problem has a solution in our setting and $W_2$ is a geodesic distance (see Appendix B).

<figure><img src="image_11.png" /><figcaption>Figure 8: The Kantorovitch problem from $\mu$ to $\nu$. From [Peyre and Cuturi, 2019]. The black dots and curve represent the transport plan $\gamma$.</figcaption></figure>

If $\mu$ is absolutely continuous and $\partial\Omega$ is $\mu$-negligible then the minimization problem in (9) has a unique solution and is equivalent to the Monge problem

$$W_2^2(\mu, \nu) = \min_{T \text{ s.t. } T_*\mu = \nu} \int_{\Omega} \|T(x) - x\|^2 d\mu(x) \quad (10)$$

In the Monge problem, mass is no longer split and all the mass in $x$ goes to $T(x)$ (see Figure 9 below). A function $T$ that satisfies $T_*\mu = \nu$ is called a transport map between $\mu$ and $\nu$.

<figure><img src="image_12.png" /><figcaption>Figure 9: The Monge problem from $\mu$ to $\nu$. From [Kolouri et al., 2017].</figcaption></figure>
The Monge problem has a unique solution $T^*$ linked to the solution $\gamma^*$ of (9) through $\gamma^* = (\text{id}, T^*)_{\#\mu}$, where $(\text{id}, T^*)$ is the function $x \mapsto (x, T^*(x))$. The Monge formulation is more suitable to the classification problems we consider, since we don't have a specific target position $y$ to learn a transport plan $\pi(x, y)$.

Another equivalent formulation of the optimal transport problem in this setting is the dynamical formulation [Benamou and Brenier, 2000]. Here, instead of directly pushing samples of $\mu$ to $\nu$ 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 $\mu_t$ at time $t$ satisfies the continuity equation

$$ \partial_t \mu_t + \nabla \cdot (\mu_t v_t) = 0 \qquad (11) $$

assuming that initial and final conditions are given by $\mu_0 = \mu$ and $\mu_1 = \nu$ respectively. The continuity equation describes the evolution of a probability distribution under a velocity field. In this case, the optimal displacement is the one that minimizes the total action caused by the velocity field $v$:

$$ W_2^2(\mu, \nu) = \min_v \int_0^1 \|v_t\|_{L^2(\mu_t)}^2 dt \qquad (12) \\ \text{s.t.} \quad \partial_t \mu_t + \nabla \cdot (\mu_t v_t) = 0, \quad \mu_0 = \mu, \mu_1 = \nu $$

Instead of describing the evolution of the density through the continuity equation, we can describe the paths $\phi_t^x$ taken by particles at position $x$ from $\mu$ when displaced along the flow $v$. Here $\phi_t^x$ is the position at time $t$ of the particle that was at $x \sim \mu$ at time 0. The continuity equation is then equivalent to $\partial_t \phi_t^x = v_t(\phi_t^x)$. Rewriting the conditions as necessary, Problem (12) becomes

$$ W_2^2(\mu, \nu) = \min_v \int_0^1 \|v_t\|_{L^2((\phi_t)_{\#\mu})}^2 dt \qquad (13) \\ \text{s.t.} \quad \partial_t \phi_t^x = v_t(\phi_t^x), \quad \phi_0 = \text{id}, (\phi_1)_{\#\mu} = \nu $$

and the optimal $T^*$ that solves (10) is in fact $T^*(x) = \phi_1^x$ for $\phi$ that solves $\partial_t \phi_t^x = v_t^*(\phi_t^x)$ together with the optimal $v^*$ from (13). This formulation is the most useful to us, as we will discretize $\partial_t \phi_t^x = v_t^*(\phi_t^x)$ to get a residual architecture.

<figure><img src="image_13.png" /><figcaption>Figure 10: Dynamic transport of $\mu_0$ to $\mu_1$ according to velocity field $v$ ($\phi_t^x$ is $\phi_t(x)$ in the text). From [de Bézenac et al., 2021].</figcaption></figure>
Since we will work with a discrete set of input points, this second dynamic formulation (13) is more useful to us.

Optimal transport maps have some regularity properties under some boundedness assumptions. We include one such result here in Theorem 1, and another we will use in Appendix B.5, along with more background on optimal transport theory. First let us define Hölder continuity, which is equivalent to Lipschitz continuity when $\eta = 1$:

**Definition 2.** A function $f$ is $\eta$-Hölder on a set $X \subset \mathbb{R}^d$ if $\forall a, b \in X$, we have

$$ \|f(a) - f(b)\| \le C\|a - b\|^\eta $$

for some constants $C>0$ and $0<\eta\le 1$. We denote this $f \in C^{0,\eta}(X)$. If all derivatives of $f$ up to the $k$-th order are $\eta$-Hölder on $X$, we write $f \in C^{k,\eta}(X)$.

We can now state the following regularity theorem for optimal transport maps:

**Theorem 1.** Suppose there are $X, Y$, bounded open sets in $\mathbb{R}^d$, such that the densities of $\mu$ and $\nu$ are null in their respective complements and bounded away from zero and infinity over them respectively.

Then, if $Y$ is convex, there exists $0 < \eta \le 1$ such that the optimal transport map $T$ between $\mu$ and $\nu$ is in $C^{0,\eta}(X)$.

If $Y$ isn't convex, there exists two relatively closed sets $A$ in $X$ and $B$ in $Y$ such that $T \in C^{0,\eta}(X \setminus A, Y \setminus B)$, where $A$ and $B$ are of null Lebesgue measure.

Moreover, if the densities are $C^{k,\eta}$, then $C^{0,\eta}$ can be replaced by $C^{k+1,\eta}$ in the conclusions above. In particular, if the densities are smooth, then the transport map is a diffeomorphism.

## 2.9 Optimal Transport for the Study of Neural Networks

We verify in Section 3.2.2 and Appendix F.5.1 that residual networks are indeed good approximators of optimal transport maps. This was also confirmed in [Gai and Zhang, 2021], which was released shortly after our first work [Karkar et al., 2020], and which finds that residual networks are closer to being optimal transport maps between their input and output distributions than non-residual networks.

Also released shortly after our first work, [Finlay et al., 2020] accelerate the NeuralODE model mentioned in Section 2.7 by adding an optimal transport regularization to the loss. This makes the learned dynamics $v$ well-behaved, which means that less numerical integration steps are needed. While NeuralODE has been tested on classification tasks, it is mainly used for generative and dynamic modeling tasks.

Optimal transport theory was used in [Sonoda and Murata, 2019] to analyze deep gaussian denoising autoencoders (not necessarily implemented through
residual networks) as transport systems. In the continuous limit, they are shown to transport the data distribution so as to decrease its entropy.

Finally, [de Bézenac et al., 2021] reformulate the unsupervised domain translation problem tackled by CycleGAN [Zhu et al., 2017, Sim et al., 2020] as an optimal transport problem that is solved by discretizing its dynamic formulation using residual networks.

Wasserstein gradient flows in the data space (as in this thesis) appeared recently in deep learning [Alvarez-Melis and Fusi, 2021, Gao et al., 2019, Liutkus et al., 2019, Arbel et al., 2019, Ansari et al., 2021], but with a focus on transfer learning and generation, while we focus on classification (see Section 4.3.2 for more details). Rather unrelated to our work, [Chizat and Bach, 2018, Chizat, 2020] used Wasserstein gradient flows in the parameter space to prove convergence of the training of two-layer neural networks.
<figure><img src="image_14.png" /></figure>
# 3 A Least Action Principle for the Training of Neural Networks

## Abstract

Neural networks have been achieving high generalization on many tasks despite being highly over-parameterized. Since classical statistical learning theory struggles to explain this behavior, much effort has been focused on uncovering the mechanisms behind it, in the hope of developing a more adequate theoretical framework and having a better control over the trained models. We adopt an alternative perspective, viewing the neural network as a dynamical system displacing input particles over time. We conduct a series of experiments and, by analyzing the network's behavior through its displacements, we show the presence of a low kinetic energy bias in the transport map of the network, and link this bias with generalization performance. From this observation, we reformulate the learning problem as follows: find networks that solve the task while transporting the data as efficiently as possible. This offers a novel formulation of the learning problem which allows us to provide regularity results for the solution network, based on Optimal Transport theory. From a practical viewpoint, this allows us to propose a new learning algorithm, which automatically adapts to the complexity of the task, and leads to networks with a high generalization ability even in low data regimes. The work presented in this section led to the following publication [Karkar et al., 2020] which won the runner-up best student paper award at ECML 2020.

## 3.1 Introduction

Deep neural networks have repeatedly shown their ability to solve a wide range of challenging tasks, while often having many more parameters than there are training samples. Such a performance of over-parametrized models is counter-intuitive. They seem to adapt their complexity to the given task, systematically achieving a low training error without suffering from over-fitting as could be expected [Belkin et al., 2019, Nakkiran et al., 2020, Zhang et al., 2017a]. This is in contradiction with the classical statistical practice of selecting a class of functions complex enough to represent the coherent patterns in the data, and simple enough to avoid spurious correlations [Belkin et al., 2018, Hastie et al., 2001]. Although this behavior has sparked much recent work towards explaining neural networks' success [De Palma et al., 2019, Jacot et al., 2018, Novak et al., 2018, Rahaman et al., 2019], it still remains poorly understood. Among the factors to consider are the implicit biases present in the choices made for the parametrization, the architecture, the initialization and the optimization algorithm, and that contribute all to this success. Our aim is to uncover some of these hidden biases and highlight their link with generalization.

We will focus on residual networks and adopt the dynamical point of view discussed in Section 2.6, which allows us to leverage the theories and mathematical tools developed to study, approximate and apply differential equations. We
first conduct experiments to observe how neural networks displace their inputs,
seen as particles, through time. We measure a strong empirical correlation
between good test performance and neural networks with low kinetic energy
along their transport flow. From this, we reformulate the training problem as
follows: retrieve the network which solves the task using the principle of least
action, i.e. expending as little kinetic energy as possible. This problem, in its
probabilistic formulation, is tightly linked with and inspired by the well-known
problem of finding an optimal transportation map. This yields new insights into
neural networks' generalization capabilities, and provides a novel algorithm that
automatically adapts to the complexity of the data and robustly improves the
network's performance, including in low data regimes, without slowing down
the training. To summarize, our contributions are the following:

- We highlight a low-energy bias in ResNets.

- We formulate a Least Action Principle for the training of neural networks.

- We prove existence and regularity results for networks with minimal energy.

- We provide an algorithm for retrieving minimal energy networks that
leads to better generalization performance on different classification tasks,
without complexifying the architecture.

## 3.2 Empirical Analysis of Transport Dynamics in Residual Networks

Before introducing our framework, we conduct an exploratory analysis of the impact of the network’s inner dynamics on generalization. We present below two experiments. The first one highlights how good generalization performance is closely related to low transport cost for classification tasks on MNIST [LeCun et al., 2010] and CIFAR10 [Krizhevsky, 2009]. This cost therefore appears as a natural characterization of the complexity and disorder of a network. The second experiment, performed on a toy 2D dataset, visualizes the transport induced by the blocks of a ResNet.

We consider ResNets where, after encoding, a data point $x_0$ is transported by
applying $x_{m+1} = x_m + r_m(x_m)$ for $M$ residual blocks and then classified using $F$.
We measure the disorder/complexity of a network by its transport cost which is
the sum $C = \sum_m \|r_m(x_m)\|_2^2$ of the displacements induced by its residual blocks.
This quantity corresponds to the kinetic energy of the displacement.

3.2.1 Transport Cost and Generalization in Image Classification

In order to study the correlation between the transport cost of a residual
network and its generalization ability on image data, we train convolutional 9-
block ResNets with different initializations (orthogonal and normal with different
gains), for 10-class classification tasks MNIST and CIFAR10. In Figure 11, each
point represents a trained network and gives the transport cost $C$ as a function of the test accuracy of the network.

This experiment clearly highlights the strong negative correlation between transport cost and good generalization. This illustrates the importance of the implicit initialization bias and motivates initialization schemes which favor a low kinetic energy. A number of factors contribute to this low energy bias: small initialization gains tend to bias $\|r_m(x_m)\|_2^2$ towards small values, and stochastic gradient descent does not change this much.

<figure><img src="image_15.png" /><figcaption>Figure 11: Test transport cost $C$ against test accuracy of ResNet-9 models on MNIST (left) and CIFAR10 (right) with fitted linear regressions, where each color indicates a different initialization (either orthogonal or normal with varying gains)</figcaption></figure>

### 3.2.2 Visualizing Network Dynamics in 2 Dimensions

We visualize in 2 dimensions the transport dynamics inside a network. The task is 2-class classification of a non-linearly separable dataset (two concentric circles, from `sklearn` [Pedregosa et al., 2011]) that contains 1000 points with a train-test split of 80%-20%, see Figure 12 top left. The network is a ResNet containing 9 residual blocks, followed by a fixed linear classifier. Each residual block contains two fully connected layers separated by a batch normalization and a ReLU activation.

With the cross-entropy loss alone, the behavior of a well trained and carefully initialized network achieving 100% test accuracy is illustrated in the first row of Figure 12. With a $N(0, 5)$ initialization, significantly bigger than a good initialization, the test accuracy drops to 98% (average of 100 runs) and the transport becomes chaotic (Figure 12, second row). Adding the transport cost to the loss improves the test accuracy (99.7% on average) of this badly initialized network and the movement becomes more controlled (third row of Figure 12).

Thus, controlling transport improves the behavior and generalization of the net-
work. This allows to explicitly control the network whereas implicit biases such as good initialization rely on heuristics. In Appendix F.5.2, more experiments show that in other situations that deviate from the ideal setting where the task is perfectly solved, e.g. when using a network which is too large or too small, or a small training set, controlling the transport cost also improves generalization.

<figure><img src="image_16.png" /><figcaption>Figure 12: Transformed circles test set by a ResNet-9 after blocks 1, 5 and 9 after training; first row with good initialization; second row with a $\mathcal{N}(0, 5)$ initialization; third row with a $\mathcal{N}(0, 5)$ initialization and the transport cost added to the loss.</figcaption></figure>

In 2 dimensions, we can compute the $W_2$ distance between successive point clouds before and after each block, and compare it, in Figure 37 in Appendix F.5.1, to the displacement that the block causes. This confirms that ResNets are already close to being optimal transport maps between their input and output distributions, as the distance traveled by the points is close to the minimal distance possible. We see that the transport regularization encourages this bias.
## 3.3 Method and Theory

### 3.3.1 General Setting

In order to better understand the inner workings of a neural network, we adopt a viewpoint in which the different mechanisms become apparent and are decoupled.

**Decomposing a neural network.** We consider the following model of a neural network where computations are separated into the three steps, i.e. $F \circ T \circ E$ (this corresponds to the general structure of recent deep models or to the structure of components of a deep model):

1. **Dimensionality change:** starting from an input distribution $\rho$ on $\mathbb{R}^n$, a transformation $E: \mathbb{R}^n \to \Omega \subset \mathbb{R}^d$ is applied, transforming it into $\mu = E_{\#}\rho$, a distribution on $\Omega$. This corresponds to the first few layers in most architectures and represents a change of dimension. We call $E$ the *encoder*.

2. **Data transport:** then $\mu$ is transformed by a mapping $T: \Omega \to \Omega$, which we see as a transport map. Here, the dimensionality doesn't change and, if this part of the network is a sequence of residual blocks, $T$ can be written as the discretized flow of an ODE.

3. **Task-specific final layers:** a final function $F: \mathbb{R}^d \to \mathcal{Y}$ is applied to $T_{\#}\mu$ in order to compute the loss $\mathcal{L}$ associated with the task at hand, e.g. $F$ could be a linear classifier. Like $E$, $F$ is typically made up of a few layers.

<figure><img src="image_17.png" /><figcaption>Figure 13: Decomposition of a neural network</figcaption></figure>

Our focus is on analyzing the second phase, data transport, and we assume that the encoder $E$ is pretrained and fixed (this will be relaxed in the experiments). To solve a complex non-linear task for which a neural network is needed, the data has to be transformed in a non-trivial way, meaning that this is an essential phase, for example in the case of classification, $T_{\#}\mu$ needs to be linearly separable if $F$ is linear. This model is quite general, as many ResNet-like architectures [He et al., 2016b, He et al., 2016a, Xie et al., 2017, Zagoruyko and Komodakis, 2016] alternate modules that change the dimensionality (step 1, typically pooling or patch merging) and transport modules that keep the dimensionality fixed (step 2) and according to [Jastrzebski et al., 2018], the transport modules have similar
behavior. The model can then be considered as a simplified ResNet, sometimes called a single representation ResNet. Note that [Sandler et al., 2019] finds that networks that keep the same resolution remain competitive.

**The set of admissible targets.** As recent neural architectures have systematically achieved near-zero training error [Belkin et al., 2019, Belkin et al., 2018, Jacot et al., 2018, Zhang et al., 2017a], we place ourselves in this regime, which makes it possible to model this as a hard constraint. For some tasks, this constraint over $T$ is obvious: in a generative setting for example, $T_{\#}\mu$ must be equal to some prescribed distribution $\nu$ which is the target of the generation process. But in general, $T$ is less strictly constrained and the condition depends on $F$ and $\mathcal{L}$. This leads us to define a set of admissible targets for the task:

$$S_{F,\mathcal{L}} = \{\nu \in \mathcal{P}(\Omega) \mid \mathcal{L}(F, \nu) = 0\} \quad (14)$$

In general, $\mathcal{L}$ is fixed while $F$ is learned jointly with $T$. This set is supposed to be non-empty for some $F$ and, in general, it will contain many distributions. The goal of the learning task can then be reformulated as:

$$\text{Find } (T, F) \text{ such that } T_{\#}\mu \in S_{F,\mathcal{L}} \quad (15)$$

An important observation is that, even when $S_{F,\mathcal{L}}$ is reduced to a singleton, the problem is still strongly under-constrained and it is possible to obtain many such $(T, F)$ that lead to poor generalization. One can then ask why this is not the case in practice, as good generalization performance is usually achieved.

**The case of classification.** Even though our framework is general, we focus our experiments on classification tasks, with $\mathcal{L}$ being the cross entropy loss. The task consists in separating $N$ classes. We denote $\mu_i$ the class distributions which are supposed to be distributions in $\mathbb{R}^d$ of mutually disjoint supports, meaning that there is no ambiguity in the class of data points, and such that $\mu = \sum w_i \mu_i$ with $\sum w_i = 1$ and $w_i \in [0, 1[$. We want to find a transformation $T$ of these distributions such that all transported distributions can be correctly classified by a classifier $F$. When $F$ is linear, $S_{F,\mathcal{L}}$ is the set of distributions which have $N$ components that are linearly separated by $F$. Note that we place ourselves in a noiseless ideal setting where perfect classification is possible. In addition to the problem of finding $T$ being strongly under-constrained, finding such a mapping is difficult, especially when only few data samples are available, and one might wonder whether some mappings are easier to calculate than others.

### 3.3.2 Formulation

Given the empirical observations of Section 3.2, a natural way to select a robust model is to select, among the maps which transport $\mu$ to $S_{F,\mathcal{L}}$ and thus solve the task, one with a minimal transport cost. This gives us the following optimization problem:

$$\begin{array}{ll}
\inf_{T,F} & \mathcal{C}(T) = \int_{\Omega} c(x, T(x)) \, d\mu(x) \quad (16) \\
\text{subject to} & T_{\#}\mu \in S_{F,\mathcal{L}}
\end{array}$$
We consider ground costs $c(x, y) = \|x - y\|^p$ with $p > 1$, and suppose $\Omega$ compact with negligible boundary and $\mu$ absolutely continuous. We assume that the space of classifiers is compact, that the loss $\mathcal{L}$ is continuous, that the set $\cup_{F \in \mathcal{F}} S_{F, \mathcal{L}}$ is at a finite $p$-Wasserstein distance $W_p$ from $\mu$ (in particular, it is non-empty) and that all its bounded subsets are totally bounded (i.e. can be covered by finitely many subsets of any fixed size). These properties depend on the choice of the loss $\mathcal{L}$ and of a class of functions $\mathcal{F}$ for the classifier $F$. The equivalent dynamical version of (16) is then, as per Section 2.8,

$$ 
\begin{array}{ll}
\inf_{v,F} & \displaystyle \int_0^1 \|v_t\|_{L^p((\phi_t)_\sharp \mu)}^p \, dt \\
\text{subject to} & \partial_t \phi_t^x = v_t(\phi_t^x) \qquad (17) \\
& \phi_0^\cdot = \text{id} \\
& (\phi_1^\cdot)_\sharp \mu \in S_{F, \mathcal{L}}
\end{array}
$$

where $\|v_t\|_{L^p((\phi_t)_\sharp \mu)}^p = \int_{\mathbb{R}^d} \|v_t\|^p d(\phi_t)_\sharp \mu$. The result below shows that Problems (16) and (17) are equivalent and that the infima are realized as minima when $p = 2$:

**Theorem 2.** The infima of (16) and (17) are finite and are realized through a map $T$ which is (or a velocity field $v$ which induces) an optimal transportation map. When $c(x, y) = \|x - y\|^p$, then (16) and (17) are equivalent.

*Proof.* From the hypothesis above, there exists $\nu \in S_{F, \mathcal{L}}$ at a finite distance from $\mu$. Taking any transport map between $\mu$ and $\nu$, we see that the infima are finite (see Appendix B.1 for existence of optimal transport maps).

Consider (16) and take a minimizing sequence $(T_i, F_i)_i$. Set $\beta_i = (T_i)_\sharp \mu$. Then $(C(T_i))_i$ converges to the infimum which is strictly bounded by $M > 0$. Then, by definition, for $i$ large enough, $W_p^p(\mu, \nu_i) \le C(T_i) \le M$. So that $(\nu_i)_i$ is a bounded sequence in $\cup_F S_{F, \mathcal{L}}$. By the hypothesized total boundedness of bounded subsets and as $P_p(\mathbb{R}^d)$ endowed with $W_p$ is a complete metric space (see [Bolley, 2008] for a proof), up to an extraction, $(\nu_i)_i$ converges to $\nu^*$ in the closure of $\cup_F S_{F, \mathcal{L}}$. Moreover, up to an extraction, $(F_i)_i$ also converges to $F^*$ by compactness of the class of classifiers. Taking $T^*$ the optimal transport map between $\mu$ and $\nu^*$, we then have, by continuity of $\mathcal{L}$,

$$ T^*_\sharp \mu = \nu^* \in S_{F^*, \mathcal{L}} $$

and $C(T^*) \le \lim_{i} C(T_i)$ by optimality of $T^*$, which means, since $(C(T_i))_i$ is a minimizing sequence, that $C(T^*)$ minimizes (16). So $(T^*, F^*)$ is a minimizer and $T^*$ is an optimal transport map.

Finally, there exists, (Appendix B.4), a velocity field $v^*_i$ inducing the optimal transport map between $\mu$ and $\nu^*$ which then gives a minimizer $(v^*, F^*)$ for (17). By the same reasoning, taking a minimizing sequence $(v^{(i)}, F_i)_i$ and the induced maps $T_i$ shows that both problems are equivalent. $\square$
Note that uniqueness doesn't hold anymore, as the constraint $T_{\#} \mu \in S_{F,\mathcal{L}}$ in (17) is looser than in standard OT. However, as we show in the following section, the fact that the optimization problems are solved by OT maps will give regularity properties for the models induced by these optimization problems.

### 3.3.3 Regularity

Intuitively, the fact that we minimize the energy of the transport map transforming the data is akin to the idea of Occam's razor: among all the possible networks that correctly solve the task, the one transforming the data in the simplest way is selected. Moreover, it is possible to show that this optimal transformation is regular: our formulation provides an alternate view on generalization for modern deep learning architectures in the overparametrized regime.

Optimal maps can be as irregular as needed in order to fit the target distribution, however in much the same way as successfully trained DNNs, optimal maps are still surprisingly regular. In a way, they are as regular as possible given the constraints which is exactly the type of flexibility needed. However, the constraints in (16) and (17) are looser than in the standard definitions of Optimal Transport. Still, supposing that the input data distribution has a nicely behaved density, namely bounded and of compact support, with the same hypothesis as above, we have the following, which is mainly a corollary of Theorem 2:

**Proposition 1.** Consider $T^*$ the OT map induced by (16) (or (17)) given by Theorem 2. Take $X$, respectively $Y$, an open neighborhood of the support of $\mu$, respectively of $T_{\#}^* \mu$, then $T^*$ is differentiable, except on a set of null $\mu$ measure.

Additionally, if $X$ is convex and $T^*$ doesn't have singularities, then there exists $A$, respectively $B$, relatively closed in $X$, respectively $Y$, both of null measure, such that $T^*$ is $\eta$-Hölder continuous from $X \setminus A$ to $Y \setminus B$, i.e. for constants $0 < \eta \le 1$ and $C > 0$ and $\forall x, y \in X \setminus A$ we have

$$ \|T(x) - T(y)\| \le C \|x - y\|^{\eta} $$

*Proof.* This is a consequence of Theorem 2, the hypothesis made in this section and the regularity theorems in Appendix B.5. $\square$

**Remark 3.** If $Y$ in Proposition 1 is convex then $T^*$ is $\eta$-Hölder on all of $X$, i.e. $T \in C^{0,\eta}(X)$ (see Theorem 1).

**Remark 4.** If the distributions $\mu$ and $\nu = T_{\#}^* \mu$ in Proposition 1 are $C^{k,\eta}$ (i.e all derivatives up to the $k$-th derivative are $\eta$-Hölder), then the optimal transport map $T$ is $C^{k+1,\eta}$ (see Theorem 1). This means that the more regular the data, the more regular the network we find.

There are two results in Proposition 1. The first gives $\alpha$-a.e. differentiability. This is already as strong as might be expected from a classifier: there are necessarily discontinuities at the frontiers between different classes. The second is even more interesting: it gives $\eta$-Hölder continuity over as large a domain
as possible, and even a diffeomorphism if the data distribution is well-behaved enough. Here $\eta$ measures smoothness, the higher its value, the better. In particular, in the case of classification, this means that the Hausdorff dimension along the frontiers between the different classes is scaled by less than a factor of $1/\eta$ in the transported domain. If the densities are smooth, the dimension even becomes provably smaller by this result.

Intuitively, this means that the data is transported in a way that preserves and simplifies the patterns in the input distribution. In the following, we propose a practical algorithm implementing these models and use it for standard classification tasks, showing an improvement over standard models.

## 3.4 Practical Implementation

We propose here an algorithm for training ResNets using the least action principle by minimizing the kinetic energy. Starting from problem (17) with $p = 2$ and the Euclidean norm, we first discretize the differential equation via a forward Euler scheme, which yields $\phi_{m+1}^x = \phi_m^x + r_m(\phi_m^x)$. The discretized flow $r_m$ approximates $v_t$ at time $t = m/M$ and is parameterized by a residual function, giving a standard residual architecture. The residual blocks, along with a classifier $F$, are parametrized by $\theta$. Next, the constraint $(\phi_1)_\#\mu \in S_F, \mathcal{L}$ is rewritten as $\mathcal{L}(F, (\phi_1)_\#\mu) = 0$, denoted $\mathcal{L}(\theta) = 0$ below. Finally, as we only have access to a finite set $\mathcal{D}$ of samples $x$ from $\mu$, we use a Monte-Carlo approximation of the integral with respect to the distributions $(\phi_t)_\#\mu$. This gives us

$$ \begin{array}{ll} \min_{\theta} & \mathcal{C}(\theta) = \sum_{x \in \mathcal{D}} \sum_{m=0}^{M-1} \|r_m(\phi_m^x)\|^p \\ \text{subject to} & \phi_{m+1}^x = \phi_m^x + r_m(\phi_m^x) \text{ and } \phi_0^x = x \text{ for all } x \in \mathcal{D} \\ & \mathcal{L}(\theta) = 0 \end{array} \qquad (18) $$

Is it easy to see that the min-max problem $\min_{\theta} \max_{\lambda>0} \mathcal{C}(\theta) + \lambda \mathcal{L}(\theta)$ yields the same solution as Problem 18, as the first two constraints are satisfied trivially. If the constraint $\mathcal{L}(\theta) = 0$ corresponding to solving the task, which includes the classifier $F$, is not verified, this will cause the second term to grow unbounded, and the solution will thus be avoided by the minimization. This min-max problem can be solved using an iterative approach, starting from some initial $\lambda_0$ and $\theta_0$:

$$ \left\{ \begin{array}{l} \theta_{i+1} = \arg \min_{\theta} \mathcal{C}(\theta) + \lambda_i \mathcal{L}(\theta) \\ \lambda_{i+1} = \lambda_i + \tau \mathcal{L}(\theta_{i+1}) \end{array} \right. \qquad (19) $$

The minimization is done via stochastic gradient descent (SGD) for a number of steps $s$, where a step means a batch, starting from the previous parameter value $\theta_i$. This algorithm is similar to Uzawa's algorithm (also called the method
of multipliers) used in convex optimization [Santambrogio, 2015]. In practice, it is much more stable to divide the minimization objective in (19) by $\lambda_i$, yielding the following algorithm:

**Algorithm:** Training with Least Action Principle (LAP-Net)

**Input:** Training samples, step size $\tau$, number of steps $s$, initial weight $\lambda_0$

**Initialization:** Initialize the parameters $\theta_0$ and set $i = 0$

while not converged do

1. Starting from $\theta_i$, perform $s$ steps of stochastic gradient descent:

1.1. $\theta_{i+1}^0 = \theta_i$

1.2. $\theta_{i+1}^l = \theta_{i+1}^{l-1} - \epsilon(\nabla C(\theta_{i+1}^{l-1})/\lambda_i + \nabla L(\theta_{i+1}^{l-1}))$ for $l$ from 1 to $s$

1.3. $\theta_{i+1}^s = \theta_{i+1}^0$

2. Update the weight $\lambda_{i+1} = \lambda_i + \tau L(\theta_{i+1})$ and increment $i \leftarrow i + 1$

**Output:** Learned parameters $\theta$

While the high non-convexity makes it difficult to ensure exact optimality, we can still have some induced regularity when reaching a good local minimum:

**Proposition 2.** Suppose the pair $(F^{\theta^*}, T^{\theta^*})$ is reached by the optimization algorithm such that $T^{\theta^*}$ is an $\epsilon$-OT map between $\mu$ and its push-forward (meaning $\|T^{\theta^*} - T^*\|_{\infty} \le \epsilon$ where $T^*$ is the OT map). Then we have, with the same notations as in Proposition 1,

$$ \forall x, y \in X \setminus A, \quad \|T^{\theta^*}(x) - T^{\theta^*}(y)\| \le O(\epsilon + \|x - y\|^\eta) $$

*Proof.* We simply write the decomposition:

$$ T^{\theta^*}(x) - T^{\theta^*}(y) = T^{\theta^*}(x) - T^*(x) + T^*(x) - T^*(y) + T^*(y) - T^{\theta^*}(y) $$

and use the triangular inequality: the first and third terms are smaller than $\epsilon$ by hypothesis while Hölder continuity applies for the second by Proposition 1. $\square$

This shows that minimizing the transport cost still endows the model with some regularity, even in situations where the global minimum is not reached.

We also remark is that we can also consider a relaxation of the optimization problem 18 by assigning a fixed weight $\lambda$ to $\mathcal{L}$ and minimizing $\mathcal{C}(\theta) + \lambda\mathcal{L}(\theta)$. This provides a simpler and quite competitive benchmark (see Appendix F.3).

## 3.5 Experiments

### 3.5.1 Generalization

**MNIST Experiments.** The base model is a ResNet with 9 residual blocks. Two convolutional layers first encode the image of shape $1 \times 28 \times 28$ into shape $32 \times 14 \times 14$. A residual block contains two convolutional layers, each preceded by a ReLU activation and batch normalization. The classifier is made up of two fully connected layers separated by batch normalization and a ReLU activation.
We use an orthogonal initialization [Saxe et al., 2014] with gain 0.01. This and all vanilla models and their training regimes are implemented by following closely the cited papers that first introduced them, and our method is added over these training regimes. More implementation details are in Appendix F.1.

From the experiments in two dimensions, we suspect that adding the transport cost helps when the training set is small. For performance comparisons, we average the highest test accuracy achieved over 30 training epochs (over random orthogonal weight initializations and random subsets of the complete training set). We find that adding the transport cost improves generalization when the training set is very small (Table 1). We see that the improvement becomes more important as the training set becomes smaller and reaches an increase of almost 14 percentage points in the average test accuracy.

<table><caption>Table 1: Average highest test accuracy and 95% confidence interval of ResNet-9 over 50 instances on MNIST with training sets of different sizes (in %)</caption><thead><tr><th>Train size</th><th>ResNet</th><th>LAP-ResNet</th></tr></thead><tbody><tr><td>500</td><td>90.8 ± 0.4</td><td><strong>90.9</strong> ± 0.2</td></tr><tr><td>400</td><td>88.4 ± 0.5</td><td><strong>88.4</strong> ± 0.4</td></tr><tr><td>300</td><td>83.5 ± 0.5</td><td><strong>86.2</strong> ± 0.4</td></tr><tr><td>200</td><td>74.9 ± 1.0</td><td><strong>82.0</strong> ± 0.5</td></tr><tr><td>100</td><td>56.4 ± 1.5</td><td><strong>70.0</strong> ± 1.0</td></tr></tbody></table>

**CIFAR10 Experiments.** We run the same experiments on CIFAR10. The architecture is the same except that the encoder transforms the input which is of shape 3 × 32 × 32 into shape 100 × 16 × 16. We use $\lambda_0 = 0.1$, $\tau = 0.1$ and $s = 50$. We average the highest test accuracy achieved over 200 training epochs over random orthogonal weight initializations and random subsets of the complete train set. Here, we find that adding the transport cost helps for all sizes of the train set (which has 50 000 images in total). The increase in average precision becomes more important as the train set becomes smaller (Table 2).

<table><caption>Table 2: Average highest test accuracy and 95% confidence interval of ResNet-9 over 20 instances on CIFAR10 with training sets of different sizes (in %)</caption><thead><tr><th>Train size</th><th>ResNet</th><th>LAP-ResNet</th></tr></thead><tbody><tr><td>50 000</td><td>91.49 ± 0.10</td><td><strong>91.94</strong> ± 0.10</td></tr><tr><td>30 000</td><td>88.61 ± 0.14</td><td><strong>89.41</strong> ± 0.10</td></tr><tr><td>20 000</td><td>85.73 ± 0.14</td><td><strong>86.74</strong> ± 0.13</td></tr><tr><td>10 000</td><td>79.25 ± 0.25</td><td><strong>80.90</strong> ± 0.16</td></tr><tr><td>5 000</td><td>70.32 ± 0.32</td><td><strong>72.58</strong> ± 0.21</td></tr><tr><td>4 000</td><td>67.80 ± 0.26</td><td><strong>70.12</strong> ± 0.30</td></tr></tbody></table>
**CIFAR100 experiments.** On CIFAR100 [Krizhevsky, 2009], results using a ResNet are in Appendix F.3. We also used the ResNeXt [Xie et al., 2017] architecture. The residual block of a ResNeXt applies $x + \sum_{i} w_{i}(x)$ with the functions $w_{i}$ having the same architecture but independent weights, followed by a ReLU activation. We used the ResNeXt-50-32×4d architecture detailed in [Xie et al., 2017]. This is a much bigger and state-of-the-art network, as compared with the single representation ResNets used so far. It also extends the experimental results beyond the theoretical framework in three ways: the embedding dimension changes between the residual blocks, a block applies $x_{m+1} = \text{ReLU}(x_{m} + \sum_{i} w_{m,i}(x_{m}))$ and the encoder is no longer fixed. We found that penalizing $\sum_{i} w_{m,i}(x_{m})$ or the true residue $x_{m+1} - x_{m}$ is essentially equivalent. Table 3 shows consistent accuracy gains as our method (with $\lambda_{0} = 1$, $\tau = 0.1$ and $s = 5$) corrects the overfitting of the bigger ResNeXt network compared to ResNet.

<table><caption>Table 3: Average highest test accuracy and 95% confidence interval of ResNeXt50 over 10 instances on CIFAR100 with training sets of different sizes (in %)</caption><thead><tr><th>Train size</th><th>ResNeXt50</th><th>LAP-ResNeXt50</th></tr></thead><tbody><tr><td>50 000</td><td>72.97 ± 1.18</td><td><b>76.11</b> ± 0.78</td></tr><tr><td>25 000</td><td>62.55 ± 2.37</td><td><b>64.11</b> ± 1.85</td></tr><tr><td>12 500</td><td>45.90 ± 2.74</td><td><b>48.23</b> ± 1.84</td></tr></tbody></table>

### 3.5.2 Transport Visualization

We verify in Appendix F.2 using a pretrained autoencoder for visualization that the embeddings learned through LAP training on image datasets do indeed remain closer than the embeddings learned through vanilla training to the initial embedding that the first residual block receives.

More precisely, if we pretrain an autoencoder on an image dataset, we can use its encoder as the encoder of the residual network and freeze it during training. This makes it possible to visualize, in the same space as the input data (so as images), how the embeddings evolve by decoding, using the pretrained decoder, the outputs of the residual blocks. And indeed we find on MNIST, CIFAR10 and CIFAR100 that the decodings of the network's embeddings remain closer to the initial reconstruction by the autoencoder when using LAP training.

### 3.5.3 Stability

The stability of networks trained with our method is hinted at by the regularity results of Propositions 1 and 2 and is also seen in the lower variance in the performances of LAP-ResNets in Tables 1, 2 and 3 above, and in the tables in Appendices F.5 and F.3. This is expected as the model becomes more constrained and can be seen as an advantage, especially in cases where the results vary more with the initialization.
To explore this increased training stability, we try training ResNets that get
deeper without batch-normalization on CIFAR10. We see in Figure 14 that
training deep ResNets without batch-normalization is near impossible, whereas
LAP-ResNets maintain the same performance and stability without batch-
normalization for up to 50 blocks.

LAP-ResNets are also compared in this regard to the small step method of
[Zhang and Wynter, 2018, Zhang et al., 2019], which simply adds a small weight
$h$ of around 0.1 in front of the residue function to make ResNets more stable.
The Least Action Principle has the same stable test accuracy when training
without batch-normalization in Figure 14 as this method, while also improving
the test accuracy when using batch-normalization.

<figure><img src="image_18.png" /><figcaption>Figure 14: Test accuracy of ResNets of various depths without batch-normalization on CIFAR10.</figcaption></figure>

3.5.4 Training Dynamics

We now look at the evolution of the test accuracy during training in Figure 15 below. We see that the Least Action Principle acts by speeding up training in the first epochs and more generally when a jump in test accuracy happens.

The orange curve in Figure 15 is that of a network trained with the transport
cost simply as a regularizer. That means that instead of using the adaptive
optimization algorithm (19), we solve $\mathcal{L}(\theta) + \lambda C(\theta)$ with a fixed $\lambda$. Results using this training are in Appendix F.3. While it does not perform as well as LAP training, it has fewer hyperparameters to tune and still leads to better test accuracy than vanilla training.

<figure><img src="image_19.png" /><figcaption>Figure 15: Test accuracy during training of ResNeXt50 models on CIFAR100.</figcaption></figure>

## 3.6 Discussion and Conclusion

We can summarize this section as follows. We have empirically motivated and then formulated the following learning problem: among all the neural networks that correctly solve the task, select the one that transforms the data with the lowest transport cost. We have proven general results of existence and regularity for models trained this way, studied their behavior in low-dimensional settings when compared to vanilla models and shown their efficiency on standard classification tasks. We also found that the training is stabilized in an adaptive fashion without being slowed down.
Note that even though we have only considered residual networks, this framework
can be applied to any architecture by considering the static formulation (16)
of the problem. The improved stability of LAP-Nets suggests also looking into
their robustness to adversarial attacks. Another avenue of research would be to
experiment with alternative ground transport costs.
<figure><img src="image_20.png" /></figure>
# 4 Module-wise Training via the Minimizing Movement Scheme

## 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 trained networks 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. The work presented in this section has led to the following publication [Karkar et al., 2023a] at NeurIPS 2023.

## 4.1 Introduction

End-to-end backpropagation is the standard training method of neural nets. But there are reasons to look for alternatives. First, it is deemed biologically implausible [Gupta, 2020, Mostafa et al., 2018, Liao et al., 2016]. Second, it requires loading the whole model during training which can be impossible in constrained settings such as on mobile devices [Teng et al., 2020, Tang et al., 2021]. Thirdly, it forces the training of systems of cooperative networks to be sequential and synchronous, which is not flexible when the networks are distributed between a central agent and clients that operate at different rates [Jaderberg et al., 2017]. Fourthly, it prohibits training the layers in parallel. These limitations follow from the locking problems [Jaderberg et al., 2017] that end-to-end backpropagation suffers from: forward locking (each layer must wait for the previous layers to process its input), update locking (each layer must wait for the end of the forward pass to be updated) and backward locking (each layer must wait for errors to backpropagate from the last layer to be updated).

Dividing the network into modules, a module being made up of one or more layers, and greedily solving module-wise optimization problems sequentially (i.e. one after the other fully) or in parallel (i.e. batch-wise, see Section 4.3.4), solves update locking (and so also backward locking). When combined with batch buffers, parallel module-wise training solves all three problems [Belilovsky et al., 2020] and allows parallel training of the modules. Module-wise training is appealing in memory-constrained settings as it works without storing all activations at the same time, and when done sequentially, only requires loading and training one module at a time.
Despite its simplicity, module-wise training has been shown to scale well [Belilovsky et al., 2020, Pyeon et al., 2021], outperforming more complicated methods addressing the locking problems e.g. synthetic [Jaderberg et al., 2017, Czarnecki et al., 2017] and delayed [Huo et al., 2018b, Huo et al., 2018a] gradients. We can also deduce theoretical results about networks of greedily-trained shallow modules [Belilovsky et al., 2019, Belilovsky et al., 2020, Arora et al., 2014, Malach and Shalev-Shwartz, 2018, Huang et al., 2018] from the existing results about the convergence of the training of shallow networks [Arora et al., 2018, Bach, 2017, Janzamin et al., 2016, Ge et al., 2018, Du and Goel, 2018].

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 \le k \le K$, Problems

$$ (T_k, F_k) \in \arg \min_{T,F} \sum_{x \in \mathcal{D}} L(F, T \circ G_{k-1}(x)) \quad (20) $$

where $G_k = T_k \circ \dots \circ T_1$ for $1 \le k \le K$ and $G_0$=id. Here, $T_k$ is the module (one or many layers) and it receives the output of module $T_{k-1}$, and $F_k$ is an auxiliary classifier that processes the outputs of $T_k$ so that 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 or mean squared error. See Figure 16 for a representation of this training. The final network trained this way is $F_K \circ G_K$, but we can stop the forward pass at any depth $k$ and keep the network $F_k \circ G_k$ for inference if it performs better.

<figure><img src="image_21.png" /><figcaption>Figure 16: Module-wise training.</figcaption></figure>

Indeed, module-wise training suffers often from a well-documented stagnation problem [Marquez et al., 2018, Belilovsky et al., 2019, Wang et al., 2021, Pyeon
et al., 2021], whereby greedy early modules overfit and learn more discriminative features than end-to-end training, and deeper modules don't improve the test accuracy, or even degrade it. [Wang et al., 2021] verify this and also find (see Figure 17 below) that early modules destroy too much information between the features they learn and the inputs, and between the features and the labels. Later in Section 4.5.6, we find a large degradation in the accuracy of deeper modules when training on small datasets. Our methods avoids this degradation.

<figure><img src="image_22.png" /><figcaption>Figure 17: Linear separability (left), mutual information with input (center), and with output (right) of the features at different depths learned with module-wise training with *K* modules. The experiment is training a ResNet-32 on CIFAR10. *K*=1 is end-to-end training. *K*=2 means training two modules with 16 blocks each. From [Wang et al., 2021].</figcaption></figure>

To solve this issue, InfoPro [Wang et al., 2021] maximizes the mutual information that each module keeps with the input, in addition to minimizing the loss. While this method shows good results in practice, the intuition behind it was criticized in [Du et al., 2021]. To avoid early overfitting, we propose a different, transport-based perspective, leveraging the analogy between ResNets and the Euler scheme for ODEs [Weinan, 2017]. To preserve input information, we minimize the kinetic energy of the modules along with the training loss. This has tight connections with the theories of gradient flows in distribution space and optimal transport. Our approach is particularly well-adapted to ResNets [He et al., 2016a, He et al., 2016b], which remain competitive to this day [Wightman et al., 2021], and similar models such as ResNeXts [Xie et al., 2017], Wide ResNets [Zagoruyko and Komodakis, 2016] and vision transformers that are made up essentially of residual connections [Liu et al., 2021, Dosovitskiy et al., 2021], but is immediately usable on other architectures where many layers have the same input and output dimension (which allows the computation of residues) such as VGG [Simonyan and Zisserman, 2014]. Our contributions are the following:

- We propose a new method for module-wise training. Being a regularization, it is easier to implement and lighter than many recent methods such as PredSim [Nøkland and Eidnes, 2019] and InfoPro [Wang et al., 2021], that train another auxiliary network besides the auxiliary classifier for each
module (similarity matching loss for PredSim and reconstruction loss for InfoPro to estimate the mutual information with the input).

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

- Additionally, we propose *multi-lap sequential* training, a variant of sequential module-wise training, that for the same time and number of epochs, and the same memory usage, performs better in many cases.

- Experimentally, we improve the test accuracy of module-wise trained networks (ResNets, VGG and Swin Transformer), both in sequential and parallel module-wise training, beating 8 other methods. Our method also works when the modules are few and deep, and when they are numerous and shallow, when many methods focus on the first setting. Our regularization makes parallel module-wise training superior or comparable in accuracy to end-to-end training, while consuming 10% to 60% less memory.

## 4.2 Related Work

**Module-wise training methods.** Local layer-wise training was initially considered as an early training algorithm for unsupervised learning [Hinton et al., 1995], as a pre-training and initialization method [Bengio et al., 2006, Marquez et al., 2018], and as a complement to end-to-end training to avoid vanishing gradients in the Inception architecture [Szegedy et al., 2017].

However, it was recently shown to be competitive with end-to-end training on supervised image classification tasks [Belilovsky et al., 2019, Nøkland and Eidnes, 2019]. This has led to it being considered in practical settings with limited resources such as embedded training on mobile devices [Teng et al., 2020, Tang et al., 2021] and dealing with very large whole slide images in histopathology [Zhang et al., 2022].

Many papers consider using a different auxiliary loss, instead of or in addition to the classification loss: kernel similarity [Mandar Kulkarni, 2016], pairwise losses [Duan et al., 2022], class separability [Lengellé and Denœux, 1996], information-theory-inspired losses [Xu and Principe, 1999, Löwe et al., 2019, Nguyen and Choi, 2019, Ma et al., 2020, Wang et al., 2021] and biologically plausible losses [Löwe et al., 2019, Nøkland and Eidnes, 2019, Gupta, 2020, Bernd Illing, 2020, Mostafa et al., 2018]. Paper [Belilovsky et al., 2019] reports the best experimental results when solving the layer-wise problems sequentially. Methods PredSim [Nøkland and Eidnes, 2019], DGL [Belilovsky et al., 2020], Sedona [Pyeon et al., 2021] and InfoPro [Wang et al., 2021] report the best results when solving the module-wise problems in parallel. [Belilovsky et al., 2019, Belilovsky et al., 2020] do it simply
through architectural considerations regarding the auxiliary networks. Sedona [Pyeon et al., 2021] uses architecture search to choose where to split the network into modules and what auxiliary classifier to use. We describe how some of these methods work in Appendix D.

However, [Belilovsky et al., 2019] do not consider ResNets and PredSim state
that their method does not perform well on them. Sedona applies architecture
search to decide on where to split the network into up to 16 modules and
what auxiliary classifier to use before module-wise training. Only BoostResNet
[Huang et al., 2018] also proposes a block-wise training idea geared for ResNets.
However, their results only show better early performance on limited experiments
and end-to-end fine-tuning is required to be competitive. These methods can all
be combined with our regularization, and we use the auxiliary architecture from
[Belilovsky et al., 2019, Belilovsky et al., 2020].

In more recent developments, [Patel et al., 2022] attempts to also split the layers width-wise into groups of neurons and [Ren et al., 2023, Fournier et al., 2023] combine layer-wise training with forward gradient learning [Baydin et al., 2022], another alternative to end-to-end backpropagation. Finally, we mention that module-wise training has been extended to the training of spiking neural networks [Guo et al., 2022b] and of neural search networks [Chen et al., 2022].

**Reduced backpropagation methods.** Some methods allow some backward communication between the modules to improve performance, but therefore limit the computational advantages of module-wise training. For example, [Guo et al., 2022a] allow some backpropagation from a module to the last few layers of the preceding module and [Xiong et al., 2020, Gomez et al., 2020, Laskin et al., 2021] make the last layer in a module and the first in the next module the same.

**Delayed and predicted gradient methods.** Besides module-wise training, methods such as DNI [Jaderberg et al., 2017, Czarnecki et al., 2017], DDG [Huo et al., 2018b], FR [Huo et al., 2018a] and ADL [Zhuang et al., 2021], 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 [Sun et al., 2021] do. But they only split networks into a small number of sub-modules (less than five) that don't backpropagate to each other and observe training issues with more sub-modules [Huo et al., 2018a]. This makes them compare unfavorably to module-wise training [Belilovsky et al., 2020]. The high dimension of the predicted gradient which scales with the size of the network renders [Jaderberg et al., 2017, Czarnecki et al., 2017] challenging in practice. Therefore, despite its simplicity, greedy module-wise training is more appealing when working in a constrained setting.

**Other alternatives to end-to-end backpropagation.** Other similar approaches include ResIST [Dun et al., 2021] and DPA [Ni et al., 2022]. Resist randomly assigns residual blocks to modules that are trained independently and reassembled before another random partition. More of a distributed training
method, it is only compared with local SGD [Stich, 2019]. DPA surrounds modules by fixed small pre-trained layers and trains them before assembling them. See [Duan and Príncipe, 2022] for a survey on alternatives to end-to-end training.

**Wasserstein gradient flows in the data space.** Wasserstein gradient flows operating in the data space recently appeared in deep learning. In [Alvarez-Melis and Fusi, 2021], 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, [Gao et al., 2019, Liutkus et al., 2019, Arbel et al., 2019, Ansari et al., 2021] use the continuous 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.

## 4.3 Method and Theory

In this section we state the module-wise optimization problems we solve and show that successively solving these problems means following a minimizing movement scheme in distribution space that maximizes the separability of the embedding distributions. We then show that the solution modules exist and have some regularity as they are optimal transport maps. Appendices B and C give the necessary background on optimal transport and gradient flows.

### 4.3.1 Formulation

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. If each module is a single ResBlock (i.e. a function $T = \text{id} + r$), its kinetic energy is simply the squared norm of its residue $r = T - \text{id}$, which we add to the loss $L$ in the target of the greedy problems (20). Given $\tau > 0$ used to weight the regularization, we now solve, for $1 \le k \le K$, Problems

$$ (T_k^\tau, F_k^\tau) \in \arg \min_{T,F} \sum_{x \in \mathcal{D}} L(F, T \circ G_{k-1}^\tau(x)) + \frac{1}{2\tau} \|T \circ G_{k-1}^\tau(x) - G_{k-1}^\tau(x)\|^2 \quad (21) $$

where $G_k^\tau = T_k^\tau \circ \dots \circ T_1^\tau$ for $1 \le k \le K$ and $G_0^\tau = \text{id}$. The final network is now $F_K^\tau \circ G_K^\tau$. Intuitively, we can think that this biases the modules towards moving the points as little as possible, thus at least keeping the performance of the previous module. We focus on ResNets because they are already biased towards small displacements and that this bias is desirable and should be encouraged [Jastrzebski et al., 2018, Zhang et al., 2019, Hauser, 2019, De and Smith, 2020], and because $T(x) - x$ can always be computed as both terms have the same dimension. But the method can be applied to any module where this quantity 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 (21) is equivalent in the discrete case to the following problem

$$ (T_k^T, F_k^T) \in \arg \min_{T,F} \mathcal{L}(F, T_\# \rho_k^T) + \frac{1}{2\tau} \int_\Omega \|T(x) - x\|^2 d\rho_k^T(x) \quad (22) $$

with $\rho_{k+1}^T=(T_k^T)_\#\rho_k^T$ and $\rho_1^T=\rho$, i.e data embedding distributions $\rho_k^T$ are pushed forward by maps $T_k^T$, and $\mathcal{L}(F, T_\#\rho_k^T) = \int L(F, T(x)) d\rho_k^T(x) = \int L(F, \cdot) dT_\#\rho_k^T$. The difference with Problem (16) in Section 3, is that now we have a regularized problem instead of a problem with a hard constraint. But we will show in Section 4.3.4 that this still forces the solution to be an optimal transport map.

### 4.3.2 Background on the Minimizing Movement Scheme

Our main result is that solving problems (22) successively means following a minimizing movement scheme (MMS) in distribution space for minimizing $\mathcal{Z}(\mu) = \min_F \mathcal{L}(F, \mu)$, which represents the loss of the best classifier on distribution $\mu$. If we restrict ourselves to linear classifiers, $\mathcal{Z}(\rho_k^T)$ represents the linear separability of representation $\rho_k^T$ at module $k$ of the data distribution $\rho$. When auxiliary networks are not necessary, for example in generative tasks where the output and the target have the same shape, $\mathcal{Z}$ is simply $\mathcal{L}$. When the auxiliary network $F$ is pre-trained and fixed, $\mathcal{Z} = \mathcal{L}(F, \mu)$. The MMS was introduced in [Giorgi et al., 1980, Gianazza et al., 1994] as a metric counterpart to Euclidean gradient descent for minimizing functionals defined on metric spaces. The task we want to solve is minimizing $\mathcal{Z}$ in the Wasserstein space. We define the minimizing movement scheme in Definition 3. See Appendix C for more details.

We work in the metric Wasserstein space $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(\mu, \nu) = \min_{T \text{ s.t. } T_\# \mu = \nu} \int_{\Omega} \|T(x) - x\|^2_2 d\mu(x) \quad (23) $$

where we further assume that $\partial\Omega$ is negligible and that the distributions we are dealing with are absolutely continuous.

**Definition 3.** Given $\mathcal{Z}: W_2(\Omega) \to \mathbb{R}$, the minimizing movement scheme is a discretized gradient flow that is well-defined in non-Euclidean metric spaces and minimizes (under some conditions) $\mathcal{Z}$ starting from $\rho_1^T \in \mathcal{P}(\Omega)$. It is given by

$$ \rho_{k+1}^T \in \arg \min_{\rho \in \mathcal{P}(\Omega)} \mathcal{Z}(\rho) + \frac{1}{2\tau} W_2^2(\rho, \rho_k^T) \quad (24) $$

Equation (24) can be seen as a proximal step in the Wasserstein distance for minimizing $\mathcal{Z}$. Equation (24) can also be interpreted as a non-Euclidean version
of the implicit Euler scheme for following the gradient flow of $\mathcal{Z}$. Indeed, the Euclidean version of (24) for $F: \mathbb{R}^d \to \mathbb{R}$ is

$$x_{k+1}^{\tau} \in \arg \min_{x \in \mathbb{R}^d} F(x) + \frac{1}{2\tau} \|x - x_k^{\tau}\|^2 \quad (25)$$

And the optimality condition of problem (25) is

$$\frac{x_{k+1}^{\tau} - x_{k}^{\tau}}{\tau} = -\nabla F(x_{k+1}^{\tau}) \quad (26)$$

which is exactly the formula of the implicit Euler scheme (see Appendix A.1) with step size $\tau$ for the Cauchy problem of continuous-time gradient descent:

$$x'(t) = -\nabla F(x(t)) \quad (27)$$

A curve $x$ solution to (27) is called a *gradient flow* of $F$ and under convexity assumptions on $F$ tends to a minimizer of $F$ as $t$ grows. Under more conditions on $F$, the minimizing movement scheme (equivalently the implicit Euler scheme in the Euclidean case), converges as $\tau$ goes to 0 to a gradient flow $x$ of $F$ (i.e. to a solution of $x'(t) = -\nabla F(x(t)))$. Therefore if both sets of assumptions on $F$ are satisfied, $x_k^\tau$ converges to a minimizer of $F$ as $t$ tens to infinity and $\tau$ tends to 0. The advantage of the minimizing movement scheme formulation (25) compared to the implicit Euler scheme (26) is that it does not require the computation of a gradient and therefore can be more easily extended to non-Euclidean metric spaces such as the Wasserstein space in (24).

In the Wasserstein case, if $\mathcal{Z}$ is lower-semi continuous then problems (24) 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 C for details). 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 along the depth $k$, 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 $T_k^\tau$ in practice.

### 4.3.3 Link with the Minimizing Movement Scheme

So under the mentioned assumptions on $\Omega$ in the second paragraph of Section 4.3.2 and absolute continuity of the distributions, we have:

**Proposition 3.** The distributions $\rho_{k+1}^\tau = (T_k^\tau)_\sharp \rho_k^\tau$, where the functions $T_k^\tau$ are found by solving (22) and $\rho_1^\tau = \rho$ is the data distribution, coincide with the minimizing movement scheme (24) for $\mathcal{Z} = \min_F \mathcal{L}(F, \cdot)$.

*Proof.* The minimizing movement scheme (24) is equivalent to taking $\rho_{k+1}^\tau = (T_k^\tau)_\sharp \rho_k^\tau$ where

$$T_k^\tau \in \arg \min_{T: \Omega \to \Omega} \mathcal{Z}(T_k^\tau \rho_k^\tau) + \frac{1}{2\tau} W_2^2(T_k^\tau \rho_k^\tau, \rho_k^\tau) \quad (28)$$
under conditions that guarantee the existence of a transport map between $\rho_k^\tau$ and any other measure. Absolute continuity of $\rho_k^\tau$ suffices for this, and the loss can ensure that $\rho_{k+1}^\tau$ is absolutely continuous. Among the functions $T_k^\tau$ that solve problem (28), is the optimal transport map from $\rho_k^\tau$ to $\rho_{k+1}^\tau$. To solve specifically for this optimal transport map, we have to solve the equivalent Problem

$$T_k^\tau \in \arg\min_T \mathcal{Z}(T_\sharp \rho_k^\tau) + \frac{1}{2\tau} \int_\Omega \|T(x) - x\|^2 d\rho_k^\tau \quad (29)$$

Problems (28) and (29) have the same minimum value, but the minimizer of (29) is now an optimal transport map between $\rho_k^\tau$ and $\rho_{k+1}^\tau$. This is immediate from the definition (23) of the $W_2$ distance. Equivalently minimizing first over the auxiliary $F$ in (22), and from the definition of $\mathcal{Z}$, Problems (22) and (29) are equivalent, which concludes. □

Since we solve Problems (22) over neural networks, their representation power shown by universal approximation theorems [Cybenko, 1989, Hornik et al., 1989] is important to get close to equivalence between (24) and (22), as we need to approximate an optimal transport map. We also know that the training of each module, if it is shallow, converges [Arora et al., 2018, Bach, 2017, Janzamin et al., 2016, Ge et al., 2018, Du and Goel, 2018].

### 4.3.4 Regularity Result

As in Section 3, we can show, by adapting the proof of Theorem 2, that Problem (22) 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. The difference is that we have a regularized problem instead of a hard constraint. 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. We then have

**Theorem 5.** Problem (22) has a minimizer $(T_k^\tau, F_k^\tau)$ such that $T_k^\tau$ is an optimal transport map. For any minimizer $(T_k^\tau, F_k^\tau)$, $T_k^\tau$ is an optimal transport map.

*Proof.* Take a minimizing sequence $(\tilde{F}_i, \tilde{T}_i)$, i.e. such that $C(\tilde{F}_i, \tilde{T}_i) \to \min C$, where $C \ge 0$ is the target function in (22) and denote $\beta_i = \tilde{T}_{i\sharp} \rho_k^\tau$. Then by compacity $\tilde{F}_i \to F^*$ and $\beta_i \to \beta^*$ in duality with $C_b(\Omega)$ by Banach-Alaoglu. There exists $T^*$ an optimal transport map between $\rho_k^\tau$ and $\beta^*$ and we have that $C(F^*, T^*) \le \lim_{i \to \infty} C(\tilde{F}_i, \tilde{T}_i) = \min C$ by continuity of $\mathcal{L}$ and because

$$\int_{\Omega} \|T^*(x) - x\|^2 d\rho_k^\tau(x) = W_2^2(\rho_k^\tau, \beta^*) = \lim_{i \to \infty} W_2^2(\rho_k^\tau, \beta_i) \\ \le \lim_{i \to \infty} \int_{\Omega} \|\tilde{T}_i(x) - x\|^2 d\rho_k^\tau(x)$$

as $W_2$ metrizes weak convergence of measures. We take $(F_k^\tau, T_k^\tau)=(F^*, T^*)$. It is also immediate that for any minimizing pair, the transport map is optimal. □
As in Section 3, regularity properties follow for the solution modules. Given Theorem 1 in Appendix B and taken from [Figalli, 2017], $T_k^\tau$ is $\eta$-Hölder continuous almost everywhere and if the optimization algorithm we use to solve the discretized problem (21) 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 \le \epsilon$, then we have (using the triangle inequality) the following stability property of the module $\tilde{T}_k^\tau$:

$$ \|\tilde{T}_k^\tau(x) - \tilde{T}_k^\tau(y)\| \le 2\epsilon + C\|x - y\|^\eta \quad (30) $$

for almost every $x, y \in \text{supp}(\rho_k^\tau)$ and a constant $C > 0$. As we have seen in Section 3, these networks generalize better and overfit less in practice. Naively composing these stability bounds on $T_k^\tau$ and $\tilde{T}_k^\tau$ allows to get stability 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 4.3, 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 in module-wise training.

## 4.4 Practical implementation

### 4.4.1 Multi-block Modules

For simplicity, we presented in (21) the case where each module is a single ResBlock. However, in practice, we often split the network into modules that contain many ResBlocks each. We show here that regularizing the kinetic energy of such modules still amounts to a transport regularization, which means that Theorem 5, the regularity bound (30) and the link with gradient flows still apply.

If each module $T_k$ is made up of $M$ ResBlocks, i.e. applies $x_{m+1} = x_m + r_m(x_m)$ for $0 \le m < M$, we have seen in Section 3 that $\sum \|r_m(x_m)\|^2$ is its discrete kinetic energy. If $\phi_m^x$ denotes the position of a point $x$ after $m$ ResBlocks, then regularizing the kinetic energy of multi-block modules means solving

$$ (T_k^\tau, F_k^\tau) \in \arg \min_{T,F} \sum_{x \in D} (L(F, T(G_{k-1}^\tau}(x)) + \frac{1}{2\tau} \sum_{m=0}^{M-1} \|r_m(\phi_m^x)\|^2) \quad (31) \\ \text{s.t. } T = (\text{id} + r_{M-1}) \circ \dots \circ (\text{id} + r_0), \quad \phi_0^x = G_{k-1}^\tau(x) \\ \phi_{m+1}^x = \phi_m^x + r_m(\phi_m^x) $$

We also want to 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 better and more localized control of the transport. As expressed in (6), a ResNet can be seen as an Euler discretization of a differential equation and this new Problem (31) is then the discretization of Problem

$$ (T_k^\tau, F_k^\tau) \in \arg \min_{T,F} \mathcal{L}(F, T_{\#} \rho_k^\tau) + \frac{1}{2\tau} \int_0^1 \|v_t\|_{L^2((\phi_t)_{\#} \rho_k^\tau)}^2 dt \quad (32) \\ \text{s.t. } T = \phi_1, \partial_t \phi_t^x = v_t(\phi_t^x), \phi_0 = \text{id} $$
where $\rho_{k+1}^{\tau} = (T_k^{\tau})_{\#} \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 (32) the optimal transport problem in its dynamic formulation [Benamou and Brenier, 2000], and given the equivalence between the Monge OT problem (23) and the dynamic OT problem (13) in Appendix B, Problem (32) is in fact equivalent to the original continuous formulation (22), and the theoretical results in Section 4.3 follow (Proposition 3 follows immediately and take a minimizing sequence $(\tilde{F}_i, \tilde{v}^i)$ and the corresponding maps $\tilde{T}_i$ as in the proof of Theorem 5 to get the result).

### 4.4.2 Solving the Module-wise Problems

The module-wise problems can be solved in one of two ways. In sequential module-wise training, we 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. In parallel module-wise training, the modules are trained synchronously batch-wise, i.e. we 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. This is called decoupled greedy training in [Belilovsky et al., 2020], 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. Since it is observed that parallel module-wise training performs better than sequential training (see [Wang et al., 2021] and Section 4.5.5 for example), we propose a variant of sequential module-wise training that we call multi-lap sequential module-wise training that introduces more synchronicity in sequential training. Here, 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.5.5. Despite our theoretical framework being that of sequential module-wise training, our method improves the test accuracy of all three module-wise training regimes.

### 4.4.3 Varying the Regularization Weight

The discussion in Section 4.3.3 suggests taking a fixed $\tau$ 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 networks 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 as in Section 3.4 to solve Problems (21) and (31). 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 $C(\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}\ge 0$, learning rates ($\eta_i$) and batches ($x_i$), we apply for module $k$ and $i\ge 1$:

$$ \left\{ \begin{array}{l} \theta_{k,i+1} = \theta_{k,i} - \eta_i \nabla_{\theta}(\lambda_{k,i} L(\theta_{k,i}, x_i) + C(\theta_{k,i}, x_i)) \\ \lambda_{k,i+1} = \lambda_{k,i} + hL(\theta_{k,i+1}, x_{i+1}) \text{ if } i \bmod s = 0 \text{ else } \lambda_{k,i} \end{array} \right. \qquad (33) $$

The weights $\lambda_{k,i}$ will vary along modules $k$ even if we use the same initial weights $\lambda_{k,1} = \lambda_1$ 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 Theorem 5 or from Section 3 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$), and the number of hyper-parameters does not go up as the number of modules increases.

In practice (see Section 4.5.1 and Appendix G.1), this algorithm works best in only one experiment, but its behavior suggested the following heuristic which we use in the other experiments: 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 has then to be chosen.

## 4.5 Experiments

We test our method on classification tasks, $L$ being cross-entropy. We call our method TRGL for Transport-Regularized Greedy Learning. For the auxiliary classifiers, we use the architecture from DGL [Belilovsky et al., 2019, Belilovsky et al., 2020], that is a convolutional layer followed by an average pooling layer and a fully connected layer. This auxiliary classifier architecture is also very similar to that used by InfoPro [Wang et al., 2021]. The code is available at github.com/block-wise/module-wise and implementation details are in Appendix G.1.

We focus first on parallel module-wise training in Section 4.5.1 as it performs better in practice and has therefore been more explored in recent works. The methods we compare to are:
- VanGL is simply vanilla greedy module-wise training with the same architecture but without our regularization. We include its results for ablation study purposes.

- InfoPro [Wang et al., 2021] which maximize the mutual information each module keeps with the input to the network in addition to the classification loss. This therefore requires an additional auxiliary network besides the classifier. This method has two variants, depending on whether the cross-entropy loss (InfoPro Softmax) is used for the auxiliary classifiers, or a contrastive loss (InfoPro Softmax) inspired by [Chen et al., 2020b, Khosla et al., 2020, He et al., 2020].

- InfoProL [Pathak et al., 2022], which is a variant of InfoPro that maximize the mutual information each module keeps with its own input, instead of the input to the network. This requires slightly less memory.

- Sedona [Pyeon et al., 2021], which uses an architecture search phase to decide where to split the network into modules and which auxiliary classifier to use. This leads to larger early modules, reducing the memory savings of module-wise training.

- DGL [Belilovsky et al., 2020], which only focuses on the architecture of the auxiliary classifier, and whose auxiliary classifier we use.

- PredSim [Nøkland and Eidnes, 2019], which adds a similarity matching loss to the loss of every module, requiring a second auxiliary network.

- DDG [Huo et al., 2018b] and FR [Huo et al., 2018a], which are delayed gradient methods that aim to break the locking problems for parallelization and not for saving memory.

We describe how some of these methods work in Appendix D. We find that TRGL always has a better test accuracy than the other methods. Except on Transformers, is has a better test accuracy than end-to-end training, with as much as almost 60% less memory usage.

We then run experiments on sequential module-wise training in Section 4.5.5 with each module being a single residual block, which allows for the largest memory savings, as only one block and its classifier have to be loaded at a time. We reach comparable performances to [Belilovsky et al., 2019].

Finally, we verify in Section 4.5.6 that our method does avoid the early overfitting and subsequent stagnation or collapse in test accuracy.

### 4.5.1 Parallel Module-wise Training with Few Modules

The first experiment is training in parallel 3 residual architectures and a VGG-19 [Simonyan and Zisserman, 2014] divided into 4 modules on TinyImageNet. We compare in Table 4 below our results in this setting to three of the best recent parallel module-wise training methods: DGL [Belilovsky et al., 2020], PredSim
[Nøkland and Eidnes, 2019] and Sedona [Pyeon et al., 2021], from Table 2 in [Pyeon et al., 2021]. We find that our method 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.

<table><caption>Table 4: Accuracy of Parallel TRGL with 4 modules (average and 95% confidence over 5 runs) on TinyImageNet, compared to DGL, PredSim, Sedona and E2E from Table 2 in [Pyeon et al., 2021], with memory saved compared to E2E as a percentage of E2E memory in red.</caption><thead><tr><th>Dataset</th><th>Architecture</th><th>K</th><th>Parallel VanGL</th><th>Parallel TRGL (ours)</th><th>PredSim</th><th>DGL</th><th>Sedona</th><th>E2E</th></tr></thead><tbody><tr><td rowspan="4">TinyImageNet</td><td>VGG-19</td><td>4</td><td>56.17 ± 0.29 (↓ 27%)</td><td><b>57.28</b> ± 0.20 (↓ 21%)</td><td>44.70</td><td>51.40</td><td>56.56</td><td>58.74</td></tr><tr><td>ResNet-50</td><td>4</td><td>58.43 ± 0.45 (↓ 26%)</td><td><b>60.30</b> ± 0.58 (↓ 20%)</td><td>47.48</td><td>53.96</td><td>54.40</td><td>58.10</td></tr><tr><td>ResNet-101</td><td>4</td><td>63.64 ± 0.30 (↓ 24%)</td><td><b>63.71</b> ± 0.40 (↓ 11%)</td><td>53.92</td><td>53.80</td><td>59.12</td><td>62.01</td></tr><tr><td>ResNet-152</td><td>4</td><td>63.87 ± 0.16 (↓ 21%)</td><td><b>64.23</b> ± 0.14 (↓ 10%)</td><td>51.76</td><td>57.64</td><td>64.10</td><td>62.32</td></tr></tbody></table>

The second experiment is training in parallel two residual architectures divided into 2 modules on CIFAR100 [Krizhevsky, 2009]. We compare in Table 5 our results in this setting to the two delayed gradient methods DDG [Huo et al., 2018b] and FR [Huo et al., 2018a] from Table 2 in [Huo et al., 2018a]. 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 the memory savings compared to end-to-end training are negligible, but TRGL has a better accuracy by up to almost 2 percentage points.

<table><caption>Table 5: Accuracy of Parallel TRGL with 2 modules (average and 95% confidence over 3 runs) on CIFAR100, compared to DDG, FR and E2E from Table 2 in [Huo et al., 2018a].</caption><thead><tr><th>Dataset</th><th>Architecture</th><th>K</th><th>Par VanGL</th><th>Par TRGL (ours)</th><th>DDG</th><th>FR</th><th>E2E</th></tr></thead><tbody><tr><td rowspan="2">CIFAR100</td><td>ResNet-101</td><td>2</td><td>77.31 ± 0.27</td><td><b>77.87</b> ± 0.44</td><td>75.75</td><td>76.90</td><td>76.52</td></tr><tr><td>ResNet-152</td><td>2</td><td>75.40 ± 0.75</td><td><b>76.55</b> ± 1.90</td><td>73.61</td><td>76.39</td><td>74.80</td></tr></tbody></table>

The third experiment is training in parallel a ResNet-110 divided into two, four, eight and sixteen modules on STL10 [Coate et al., 2011]. We compare in Table 6 our results in this setting to the recent methods InfoPro [Wang et al., 2021] and DGL [Belilovsky et al., 2020] from Table 2 in [Wang et al., 2021]. Our TRGL clearly outperforms the other methods. It also outperforms end-to-end training in all but one case (that with 16 modules). Memory savings of parallel TRGL compared to end-to-end training (with a batch size of 64) reach around 48% and 58.5% with 8 and 16 modules respectively, with comparable test accuracy. With four modules, the 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.5.3 below for more details and comparison to the memory saving of InfoPro).
<table><caption>Table 6: Accuracy of Parallel TRGL with *K* modules (average and 95% interval over 5 runs) on STL10, compared to DGL, InfoPro and E2E from Table 2 in [Wang et al., 2021].</caption><thead><tr><th>Dataset</th><th>Architecture</th><th><i>K</i></th><th>Par VanGL</th><th>Par TRGL (ours)</th><th>DGL</th><th>InfoPro Softmax</th><th>InfoPro Contrast</th><th>E2E</th></tr></thead><tbody><tr><td rowspan="4">STL10</td><td rowspan="4">ResNet-110</td><td>2</td><td>79.85 ± 0.93</td><td><b>80.04</b> ± 0.85</td><td>75.03 ± 1.18</td><td>78.98 ± 0.51</td><td>79.01 ± 0.64</td><td>77.73 ± 1.61</td></tr><tr><td>4</td><td>77.11 ± 2.31</td><td><b>79.72</b> ± 0.81</td><td>73.23 ± 0.64</td><td>78.72 ± 0.27</td><td>77.27 ± 0.40</td><td>77.73 ± 1.61</td></tr><tr><td>8</td><td>75.71 ± 0.55</td><td><b>77.82</b> ± 0.73</td><td>72.67 ± 0.24</td><td>76.40 ± 0.49</td><td>74.85 ± 0.52</td><td>77.73 ± 1.61</td></tr><tr><td>16</td><td>73.57 ± 0.95</td><td><b>77.22</b> ± 1.20</td><td>72.27 ± 0.58</td><td>73.95 ± 0.71</td><td>73.73 ± 0.48</td><td>77.73 ± 1.61</td></tr></tbody></table>

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

<table><caption>Table 7: 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 [Pathak et al., 2022], with memory saved compared to E2E as a percentage of E2E memory consumption in red.</caption><thead><tr><th>Architecture</th><th>Dataset</th><th><i>K</i></th><th>Parallel VanGL</th><th>Parallel TRGL (ours)</th><th>InfoPro</th><th>InfoProL</th><th>E2E</th></tr></thead><tbody><tr><td rowspan="3">Swin-Tiny</td><td>STL10</td><td>4</td><td>67.00 ± 1.36 (<span style="color:red;">↓ 55%</span>)</td><td><b>67.92</b> ± 1.12 (<span style="color:red;">↓ 50%</span>)</td><td>64.61 (<span style="color:red;">↓ 38%</span>)</td><td>66.89 (<span style="color:red;">↓ 45%</span>)</td><td>72.19</td></tr><tr><td>CIFAR10</td><td>4</td><td>83.94 ± 0.42 (<span style="color:red;">↓ 33%</span>)</td><td><b>86.48</b> ± 0.54 (<span style="color:red;">↓ 29%</span>)</td><td>83.38 (<span style="color:red;">↓ 38%</span>)</td><td>86.28 (<span style="color:red;">↓ 45%</span>)</td><td>91.37</td></tr><tr><td>CIFAR100</td><td>4</td><td>69.34 ± 0.91 (<span style="color:red;">↓ 33%</span>)</td><td><b>74.11</b> ± 0.31 (<span style="color:red;">↓ 29%</span>)</td><td>68.36 (<span style="color:red;">↓ 38%</span>)</td><td>73.00 (<span style="color:red;">↓ 45%</span>)</td><td>75.03</td></tr></tbody></table>

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. In Table 8, 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, this requires only around half the memory end-to-end training requires.

<table><caption>Table 8: 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).</caption><thead><tr><th>Seq VanGL</th><th>Seq TRGL</th><th>MLS VanGL</th><th>MLS TRGL</th></tr></thead><tbody><tr><td>73.31</td><td>73.61</td><td>74.34</td><td><b>74.78</b></td></tr></tbody></table>

### 4.5.2 Ablation Study and Sensitivity to Hyperparameter $\tau$

We see in the Tables 4, 5 and 6 in Section 4.5.1 that the improvement in accuracy from the regularization compared to vanilla module-wise training (VanGL) tends
to increase as the number of modules increases and reaches almost 4 percentage points with 16 modules. We further show in Figure 43 in Appendix G.4 that our method is not very sensitive to the choice of $\tau$ over a large scale of values.

### 4.5.3 Memory Usage of Parallel TRGL

As seen above, parallel TRGL is lighter than end-to-end training by up to more than 50%. We compare here the memory needs of our method to that of InfoPro [Wang et al., 2021] 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 6).

InfoPro [Wang et al., 2021] also proposes to split the network into *K* modules that have the same weight but not necessarily the same number of layers, which leads to even less memory usage. Indeed, the earlier layers of a neural networks usually have more parameters and are heavier. Therefore, by counting the number of parameters and making the earlier modules shallower than the deeper ones we can split the network into modules that are equally heavy instead of equally deep. InfoPro only implement this for $K \le 4$ modules. When the modules are even in weight and not in depth, keeping the notations of [Wang et al., 2021], we call the training methods VanGL*, TRGL* and InfoPro*. As said above, this leads to shallower early modules which hurts performance according to [Pyeon et al., 2021]. We verify that this degradation is slight with our regularization in Table 30 in Appendix G.3 (to be compared with Table 6), and that TRGL* still outperforms VanGL*, InfoPro* and end-to-end training.

We report in Table 9 the memory saved as a percentage of the 6230 MiB memory required by end-to-end training with the same batch size. We see in Table 9 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 (VanGL) is small.

<table><caption>Table 9: 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%.</caption><thead><tr><th rowspan="2">K</th><th colspan="3">Equally deep modules</th><th colspan="3">Equally heavy modules</th></tr><tr><th>Par VanGL</th><th>Par TRGL (ours)</th><th>InfoPro</th><th>Par VanGL*</th><th>Par TRGL* (ours)</th><th>InfoPro*</th></tr></thead><tbody><tr><td>4</td><td>27% (77.11)</td><td>24% (<strong>79.72</strong>)</td><td>18% (78.72)</td><td>41% (77.14)</td><td>39% (<strong>78.94</strong>)</td><td>33% (78.78)</td></tr><tr><td>8</td><td>50% (75.71)</td><td>48% (<strong>77.82</strong>)</td><td>37% (76.40)</td><td></td><td></td><td></td></tr><tr><td>16</td><td>61% (73.57)</td><td>58% (<strong>77.22</strong>)</td><td>59% (73.95)</td><td></td><td></td><td></td></tr></tbody></table>

We summarize in Table 31 in Appendix G.3, memory savings of parallel TRGL with 4 modules that range between 8 and 30% of E2E memory, on the 4 networks trained on TinyImageNet in Table 4 and a the Swin Transformer from Table 7. Note that methods DDG [Huo et al., 2018b] and FR [Huo et al., 2018a],
being delayed gradient methods and not module-wise training methods, do
no save memory (they actually increase memory usage, see FR [Huo et al.,
2018a]). Sedona [Pyeon et al., 2021] 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 [Belilovsky et al., 2020] is architecture-wise
essentially identical to VanGL and consumes the same amount of memory.

4.5.4 Training Time of Parallel TRGL

Since we do not implement forward unlocking with batch buffers as in DGL (i.e.
only the backward passes of the modules happen in parallel), parallel module-
wise training does slightly slow down training in this case. Epoch time increases
by 6% with 2 modules and by 16% with 16 modules compared to end-to-end
training. 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.

4.5.5 Sequential Full Block-wise Training

Full block-wise sequential training, meaning that each module is a single residual
block and that the blocks are trained sequentially, allows to load and train only
one block at a time. The memory needed is therefore divided by approximately
the number of modules. The minimal memory needed to train the model is
now the size of the heaviest module, which is the first one if the modules are
equally deep, but it is possible to make the modules equally heavy (see Section
4.5.3 above). So even though it has been less explored in recent module-wise
training methods because of its inferior accuracy compared to parallel module-
wise training, it has been used in practice in very constrained settings such as
on-device training of sensors [Teng et al., 2020, Tang et al., 2021]. We test our
regularization therefore in this section on sequential block-wise training.

We propose in this setting 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 [He et al., 2016a] 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 [Belilovsky et al., 2019].

We check in Table 24 in Appendix G that this architecture works well with
parallel module-wise training by comparing on CIFAR10 a 2-7 ResNet with
DGL, InfoPro and DDG. The 2-7 ResNet has 45 millions parameters, which is
about the same as the ResNet-110 divided in two used by the other methods,
but vanilla parallel module-wise training with this architecture already performs
better than the other methods.

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, or at worst does not fail, in all settings (parallel or sequential with many or few splits), which other methods so far don't do. For sequential training, block *k* is trained for 50+10k epochs where 0≤*k*≤10, block 0 being the encoder. This idea of increasing the number of epochs per layer along with the depth is found in [Marquez et al., 2018]. For multi-lap sequential 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 10, we see that multi-lap sequential (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, which already performs quite close to end-to-end training in the full data regime and much better in the small data regime.

<table><caption>Table 10: 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 training.</caption><thead><tr><th>Train size</th><th>Seq VanGL</th><th>Seq TRGL</th><th>MLS VanGL</th><th>MLS TRGL</th><th>Par VanGL</th><th>Par TRGL</th><th>E2E</th></tr></thead><tbody><tr><td>50000</td><td>68.74 ± 0.45</td><td><b>68.79</b> ± 0.56</td><td>69.48 ± 0.53</td><td><b>69.95</b> ± 0.50</td><td>72.59 ± 0.40</td><td><b>72.63</b> ± 0.40</td><td>75.85 ± 0.70</td></tr><tr><td>25000</td><td>60.48 ± 0.15</td><td><b>60.59</b> ± 0.14</td><td>61.33 ± 0.23</td><td><b>61.71</b> ± 0.32</td><td>64.84 ± 0.19</td><td><b>65.01</b> ± 0.27</td><td>65.36 ± 0.31</td></tr><tr><td>12500</td><td>51.64 ± 0.33</td><td><b>51.74</b> ± 0.26</td><td>51.30 ± 0.22</td><td><b>51.89</b> ± 0.30</td><td>55.13 ± 0.24</td><td><b>55.40</b> ± 0.35</td><td>52.39 ± 0.97</td></tr><tr><td>5000</td><td>36.37 ± 0.33</td><td><b>36.40</b> ± 0.40</td><td>33.68 ± 0.48</td><td><b>34.61</b> ± 0.59</td><td>39.45 ± 0.23</td><td><b>40.36</b> ± 0.23</td><td>36.38 ± 0.31</td></tr></tbody></table>

We report further results for block-wise training on MNIST [LeCun et al., 2010] and CIFAR10 [Krizhevsky, 2009], but now we report the last block's accuracy. We see again greater improvement from the regularization as the training set gets smaller, gaining as much as 6 percentage points (Table 11 below, and Tables 25 and 26 in Appendix G.2).

<table><caption>Table 11: Average last block test accuracy and 95% confidence interval of 10-1 ResNet over 10 runs on CIFAR10 with different train sizes and sequential (Seq) TRGL and VanGL.</caption><thead><tr><th>Train</th><th>Seq VanGL</th><th>Seq TRGL</th><th>E2E</th></tr></thead><tbody><tr><td>50000</td><td>88.02 ± .18</td><td><b>88.20</b> ± .24</td><td>91.88 ± .18</td></tr><tr><td>25000</td><td>83.95 ± .13</td><td><b>84.28</b> ± .22</td><td>88.75 ± .27</td></tr><tr><td>10000</td><td>76.00 ± .39</td><td><b>77.18</b> ± .34</td><td>82.61 ± .35</td></tr><tr><td>5000</td><td>67.74 ± .49</td><td><b>69.67</b> ± .44</td><td>73.93 ± .67</td></tr><tr><td>1000</td><td>45.67 ± .88</td><td><b>51.34</b> ± .90</td><td>50.63 ± .98</td></tr></tbody></table>

The 88% accuracy of sequential training on CIFAR10 in Table 11 is the same in table 2 of [Belilovsky et al., 2019], which is the best method for layer-wise sequential training available, with VGG [Simonyan and Zisserman, 2014] networks of comparable depth and width.
The results here show a few limitations of our method, as the improvements from the regularization are often minimal on sequential training. However, the results show that our approach works, or at least does not hurt performance, 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 different settings from the original one in subsequent papers (e.g. delayed gradients methods when the number of modules increases [Huo et al., 2018a], and PredSim in [Pyeon et al., 2021]).

### 4.5.6 Avoiding Early Overfitting

We verify that our method avoids the stagnation in test accuracy with increasing depth. In Figure 18 below, we look at the test accuracy of each module after module-wise training with and without the regularization.

On the left, from the experiment with parallel module-wise training with 16 modules on STL10 from Table 6 in Section 4.5.1, we see that TRGL performs worse than vanilla module-wise training in the early modules, but surpasses it greatly in later modules. This indicates that it does avoid early overfitting and the subsequent stagnation in test accuracy. On the right, from experiments with sequential block-wise training from Table 11 in Appendix G, we see a large decline in the test accuracy of vanilla block-wise training after the first block that adding our transport regularization completely avoids. We see similar patterns in Figure 42 in Appendix G with parallel and multi-lap sequential block-wise training.

<figure><img src="image_23.png" /><figcaption>Figure 18: 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. Right: sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training of a 10-1 ResNet on 2% of CIFAR10.</figcaption></figure>
## 4.6 Discussion and Conclusion

We introduced a transport regularization for module-wise training that theoretically links it to gradient flows of the loss in distribution space. Our method provably leads to more regular modules and experimentally consistently improves the test accuracy of module-wise and block-wise sequential, parallel and multi-lap sequential (a variant of sequential training that we introduce) training. Through this simple method that does not complexify the architecture, we aim at making module-wise training competitive with end-to-end training while benefiting from its computational advantages: reduced memory usage and parallelism that is complementary to model and data parallelism in the case of parallel module-wise training, and training only one module at a time in constrained settings such as training on mobile devices in the case of sequential module-wise training.

The method can easily be combined with other methods of layer-wise training. Future work can also experiment with working in Wasserstein space $W_p$ for $p \neq 2$, i.e. regularizing with a norm $\|\cdot\|_p$ with $p \neq 2$. One can also ask how far the obtained composition network $G_K$ is from being an OT map itself, which could provide a better stability bound than the one obtained by naively chaining the stability bounds (30) that follow from each module $T_k$ being an OT map.
# 5 Adversarial Detection Through Transport Dynamics

## Abstract

We propose a detector of adversarial samples that is based on the view of residual networks as discrete dynamical systems. The detector tells clean inputs from abnormal ones by comparing the discrete vector fields they follow throughout the network's layers. We also show that regularizing this vector field during training makes the network more regular on the data distribution's support, thus making the network's activations on clean samples more distinguishable from those on abnormal samples. Experimentally, we compare our detector favorably to other detectors using seen and unseen attacks, and show that the regularization of the network's dynamics improves the performance of adversarial detectors that use the internal embeddings as inputs, while also improving the network's test accuracy. The work presented in this section has led to the following publication [Karkar et al., 2023b] at ECML 2023.

## 5.1 Introduction

Neural networks have improved performances on many learning tasks, including image classification. They are however vulnerable to adversarial attacks [Szegedy et al., 2013]. These attacks modify an image in a way that is imperceptible to the human eye but that fools the network into wrongly classifying the modified image. These attacks can transfer to and fool other networks [Moosavi-Dezfooli et al., 2017a], can be carried out in the physical world, for example causing autonomous cars to misclassify road signs [Eykholt et al., 2018], and can be generated without having access to the network [Liu et al., 2017]. Developing networks that are robust to such attacks or accompanied by detectors that can detect them is therefore indispensable to deploying them safely in the real world [Amodei et al., 2016].

We focus here on the task of detecting adversarial samples. Neural networks trained with a softmax classifier produce overconfident predictions even for out-of-distribution and adversarial inputs [Nguyen et al., 2015]. This makes it difficult to detect such inputs via the softmax outputs. An adversarial detector is a system capable of predicting whether an input at test time has been adversarially modified. Adversarial detectors are trained on a dataset made up of clean and attacked inputs, after network training. While simply training the detector on the inputs has been tried, using their embeddings works better [Carlini and Wagner, 2017a]. Detection methods vary by which activations to use and how to process them to extract the features that are fed to the classifier that tells clean samples from attacked ones.

We make two contributions in this work. First, we propose an adversarial detector that is based on the view of neural networks as dynamical systems that move inputs in space, time represented by depth, to separate them before applying a
linear classifier. Given their resemblance to the Euler scheme for ODEs, ResNets
[He et al., 2016b, He et al., 2016a, Weinan, 2017] are particularly amenable to
this analysis. But the analysis and implementation extend immediately to any
network where most layers have the same input and output dimensions. Our
detector follows the trajectory of samples in space, through time, to differentiate
clean and adversarial images. The statistics that we extract are the positions of
the internal embeddings in space approximated by their norms and cosines to
a fixed vector. This means that the extraction time and size are very low and
allows to consider all layers, when many methods have to restrict themselves to
a subset of layers.

[Wu et al., 2020] show an increased vulnerability of residual architectures to transferable attacks, precisely because of the skip connections. Given that skip connections are present in almost all state-of-the-art architectures, this motivates the need for a detector that is well adapted to residual architectures.

We test our detector on adversarial samples generated by 8 attacks, on 3 datasets and networks, comparing it favorably to the Mahalanobis [Lee et al., 2018] and natural scene statistics [Kherchouche et al., 2020] detectors. We also show that our detector can be used effectively for out-of-distribution detection even after only training it on adversarial samples.

Building on this dynamical viewpoint, our second contribution is to use the
transport regularization during training proposed in Section 3 to make the
activations of adversarial samples more distinguishable from those of clean
samples, thus making adversarial detectors perform better, while also improving
generalization.

We prove that the regularization achieves this by making the network more
regular on the support of the data distribution. This does not make it more
robust to attacks, but it will make the activations of the clean samples closer to
each other and further from those of abnormal out-of-distribution samples, thus
making adversarial detection easier, since adversarial attacks tend to lie outside
the data manifold. We favorably compare the effect of our regularization in this
regard to RCE training [Pang et al., 2018].

## 5.2 Related Work

An adversarial attack aims at slightly modifying an input image at test time by
no more than $\epsilon$ in a certain metric (an $L_p$ norm usually) so that the perturbation
is undetectable to the human and that the perturbed image is misclassified by
the network. More formally, given a classifier $f$ in a multi-class classification
task and $\epsilon > 0$, an adversarial sample $y$ constructed from a clean sample $x$
is $y = x + \delta$, such that $\|\delta\| \le \epsilon$ and $f(y) \ne f(x)$. The maximal perturbation
size $\epsilon$ has to be so small as to be almost imperceptible to a human (see Figure
19 below). An adversarial attack is an algorithm that finds such adversarial
samples. We present below some popular adversarial attacks that are used to
find such adversarial samples. We present in more detail how these attacks work
in Appendix E.1.

<figure><img src="image_24.png" /><figcaption>Figure 19: Example of an adversarial attack from [Goodfellow et al., 2015a].</figcaption></figure>

White-box attacks are attacks where the adversary has access to the network's
architecture, weights and gradients. The Fast Gradient Method (FGM) [Good-
fellow et al., 2015b] takes a perturbation step in the direction of the gradient
that maximizes the loss (see Figure 19 above). Projected Gradient Descent
(PGD) [Madry et al., 2018] and the Basic Iterative Method (BIM) [Kurakin
et al., 2017] are iterative versions of FGM. Auto-PGD (APGD) [Croce and Hein,
2020b] is a variant of PGD that uses an adaptive step size. Two slower but more
effective attacks are DeepFool (DF) [Moosavi-Dezfooli et al., 2016], which itera-
tively perturbs an input in the direction of the closest decision boundary, and
Carlini-Wagner (CW) [Carlini and Wagner, 2017b], which solves an optimization
problem to find the perturbation.

Black-box attacks don't have any knowledge about the network and can
only query it. They include the Boundary Attack (BA) [Brendel et al., 2018],
which starts from a large adversarial sample and moves towards the boundary
decision to minimize the perturbation, and Hop-Skip-Jump (HSJ) [Chen et al.,
2020a], which estimates the gradient direction at the decision boundary.

AutoAttack [Croce and Hein, 2020b] is combination of 3 white-box attacks (two versions of Auto-PGD and the FAB attack [Croce and Hein, 2020a]), and of the black-box Square Attack (SA) [Andriushchenko et al., 2020].

Two main defense mechanisms exist against adversarial attacks: adversarial robustness and adversarial detection.

Adversarial robustness means training a network that is not easily fooled by
adversarial samples, for example through adversarial training (augmenting the
training set with adversarial samples [Szegedy et al., 2013, Goodfellow et al.,
2015b]), regularization [Gu and Rigazio, 2015], distillation [Papernot et al.,
2016c] and denoising [Meng and Chen, 2017]. In the context of adversarial
robustness, the view of ResNets as transport systems was used in [Wang et al.,
2019b] to relate injecting noise into residual networks and ensembling them
to approximating the Feynman-Kac formula for convection-diffusion, which
makes the level sets more regular and the network more robust to adversarial
attacks. These robustness-enhancing defense mechanisms can generally be
avoided by newer attacks [Carlini and Wagner, 2017a], and a theoretical debate
exists over a trade-off that is observed in practice between test accuracy and
adversarial robustness [Yang et al., 2020b, Raghunathan et al., 2020, Schmidt
et al., 2018, Dohmatob, 2019].

Given these limitations of adversarial robustness, a detector capable of detecting
adversarial samples is an important addition to a neural network in deployment.
An advantage to having an adversarial detector is that it can, at least in
principle, be used to also detect out-of-distribution samples, for which the notion
of robustness does not make sense. While simply training the detector on a
dataset of clean and adversarial inputs has been tried, using their embeddings as
they go through the network works better for most detection methods [Carlini
and Wagner, 2017a]. But using all the intermediate embeddings of an input
as features to train an adversarial detector on is impossible given the huge
memory and time it would require. Detection methods therefore vary by which
activations they use and how to process them to extract the features that are
fed to the classifier that tells clean samples from attacked ones.

An early idea for adversarial detection was to use a second network [Metzen
et al., 2017]. However, this network can be adversarially attacked as easily as the
first network. More recent popular statistical approaches include LID [Ma et al.,
2018], Mahalanobis [Lee et al., 2018, Ren et al., 2021], MMD [Grosse et al., 2017]
and KD [Feinman et al., 2017]. LID trains the detector on the local intrinsic
dimensionality (LID, [Houle, 2013, Amsaleg et al., 2015, Tempczyk et al., 2022]
of activations approximated over a batch. Mahalanobis train the detector on the
Mahalanobis distances between the activations and a Gaussian fitted to them
during network training, assuming they are normally distributed. MMD employs
a statistical test that uses the maximum mean discrepancy distance (MMD,
[Gretton et al., 2012]) between shuffled inputs. KD combines kernel density (KD,
[Parzen, 1962, Rosenblatt, 1956]) estimates and Bayesian uncertainty estimates
obtained via dropout [Srivastava et al., 2014] as detection features.

Our detector is not a statistical approach and does not need batch-level statistics, nor statistics from the training data of the network. [Carrara et al., 2019] is the most similar approach as they monitor the evolution of the positions of adversarial samples by comparing them to representative pivot positions of clean points from each class. This also requires the detector to have access to the training data of the network to compute the pivot positions as a centroid or a medoid, and makes the extraction time and size potentially much larger as it increases linearly with the number of classes.

Detectors trained in the Fourier domain [Harder et al., 2021] or on natural
scene statistics [Kherchouche et al., 2020, Kherchouche et al., 2022] have been
proposed. See Appendix E.2 for details and [Aldahdooh et al., 2022] for a review.
Our second contribution is to regularize the network in a way that makes it Hölder-continuous, but only on the input distribution's support. Estimations of the Lipschitz constant of a network have been used as estimates of its robustness to adversarial samples in [Weng et al., 2018, Szegedy et al., 2013, Virmaux and Scaman, 2018, Hein and Andriushchenko, 2017], and making the network more Lipschitz (e.g. by penalizing an upper bound on its Lipschitz constant) has been used to make it more robust to adversarial samples in [Hein and Andriushchenko, 2017, Cisse et al., 2017]. These regularizations [Cisse et al., 2017, Gu and Rigazio, 2015] often work directly on the weights of the neural network, therefore making the neural network more regular on the entire input space. The main difference with our method is that we only endue the network with regularity on the support of the clean samples. This won't make it more robust to adversarial samples who lie outside the support of the data distribution, but it will make the activations of the clean samples closer to each other and further from those of abnormal samples, thus making adversarial detection easier.

That adversarial samples tend to lie outside the data manifold, particularly in its co-dimensions, is a common observation and explanation for why adversarial attacks are easy to find in high dimensions [Gilmer et al., 2018, Tanay and Griffin, 2016, Song et al., 2018, Ma et al., 2018, Samangouei et al., 2018, Khoury and Hadfield-Menell, 2018, Alemany and Pissinou, 2022, Feinman et al., 2017, Chattopadhyay et al., 2019]. According to [Tanay and Griffin, 2016], adversarial samples are easy to find in high dimensions because the decision boundary intersects the data manifold but extends beyond it while remaining close to it (see Figure 20 below).

<figure><img src="image_25.png" /><figcaption>Figure 20: Boundary tilting perspective of adversarial samples from [Tanay and Griffin, 2016]. Adversarial samples lie on the other side of the boundary plane in the co-dimension of the data manifold, and not on other side of the boundary line that intersects the data manifold.</figcaption></figure>

To the best of our knowledge, [Pang et al., 2018] is the only other method that
also attempts to improve detection by encouraging the network during (normal not adversarial) training to learn latent representations that are more different between clean and adversarial samples. They do this by replacing the cross-entropy loss by reverse cross-entropy (RCE) that encourages uniform softmax outputs among the non-predicted classes. We find that our regularization leads to better classification accuracy and adversarial detection than RCE.

## 5.3 Method and Theory

We take the view that a residual network moves its inputs according to a discrete vector field to separate them. Points in each class follow similar trajectories. Heuristically, for a successful adversarial sample that lies close to clean samples, the vector field it follows has to be different at some step from that of the clean samples (in its original class and in other classes), so that it can join the trajectory of the points in another class. In Section 5.3.1, we present how to detect adversarial attacks by looking at these trajectories. In Section 5.3.2, we discuss the link between the regularity of neural networks and their robustness to adversarial attacks and use the transport regularization (18) to improve detectability of adversarial attacks.

### 5.3.1 Detection

Given a residual network that applies $x_{m+1} = x_m + hr_m(x_m)$ to an input $x_0$ for $0 \le m < M$, we consider the embeddings $x_m$ for $0 < m \le M$, or the residues $r_m(x_m)$ for $0 \le m < M$. To describe their positions in space, we take their norms and their cosine similarities with a fixed vector as features to train our adversarial detector on. Using only the norms already gave good detection accuracy. Cosines to other orthogonal vectors can be added to better locate the points at the price of increasing the number of features. We found that using only one vector already gives state-of-the-art detection, so we only use the norms and cosines to a fixed vector of ones. We train the detector (a random forest in practice) on these features. The embeddings $x_m$ and the residues $r_m(x_m)$ can equivalently describe the trajectory of $x_0$ in space through the blocks. In practice, we use the residues $r_m(x_m)$, with their norms squared and averaged. So the feature vector given to the random forest for each $x_0$ that goes through a network that applies $x_{m+1} = x_m + h r_m(x_m)$ is

$$ \left( \frac{1}{d_m} \|r_m(x_m)\|^2, \cos(r_m(x_m), \mathbf{1}_m) \right)_{m=0}^{M-1} \qquad (34) $$

and the label is 0 if $x_0$ is clean and 1 if it is adversarial. Here cos is the cosine similarity between two vectors and $\mathbf{1}_{d_m}$ is a vector of ones of size $d_m$ where $d_m$ is the size of $r_m(x_m)$. For any non-residual architecture $x_{m+1} = g_m(x_m)$, the vector $x_{m+1} - x_m$ can be used instead of $r_m(x_m)$ on layers that have the same input and output dimension, allowing to apply the method to any network with many such layers. And we do test the detector on a ResNeXt, which
does not fully satisfy the dynamic view, as the activation is applied after the skip-connection, i.e. $x_{m+1} = \text{ReLU}(x_m + h r_m(x_m))$.

The number of features is twice that of residual blocks (a norm and a cosine per block). This is of the same order as for other popular detectors such as Mahalanobis ([Lee et al., 2018]) and LID ([Ma et al., 2018]) that extract one feature per residual stage (a residual stage is a group of blocks that keep the same dimension). Even for common large architectures, twice the number of residual blocks is still a small number of features for training a binary classifier (ResNet152 has 50 blocks). More importantly, the features we extract (norms and cosines) are quick to calculate, whereas those of other methods require involved statistical computations on the activations. We include in Section 5.4.6 a favorable time comparison of our detector to the Mahalanobis detector. Another advantage is that our detector does not have a hyper-parameter to tune unlike the Mahalanobis and LID detectors.

### 5.3.2 Regularization

Regularity of neural networks (typically Lipschitz continuity) has been used as a measure of their robustness to adversarial attacks [Weng et al., 2018, Szegedy et al., 2013, Virmaux and Scaman, 2018, Hein and Andriushchenko, 2017, Cisse et al., 2017]. Indeed, the smaller the Lipschitz constant $L$ of a function $f$ satisfying $\|f(x) - f(y)\| \le L\|x - y\|$, the less $f$ changes its output $f(y)$ for a perturbation (adversarial or not) $y$ of $x$. Regularizing a network to make it more Lipschitz and more robust has therefore been tried in [Hein and Andriushchenko, 2017] and [Cisse et al., 2017]. For this to work, the regularization has to apply to adversarial points, i.e. outside the support of the clean data distribution. Indeed, the Lipschitz continuity obtained though most of these methods and analyses apply on the entire input space $\mathbb{R}^d$ as they penalize the network's weights directly. Likewise, choosing a small step size $0<h<1$ for the residual blocks as in [Zhang and Wynter, 2018, Zhang et al., 2019], while improving stability, will have the same effect on all inputs, adversarial or not.

We propose an alternative approach where we regularize the network only on the support of the input distribution, making it $\eta$-Hölder on it. Since this result does not apply outside the input distribution's support, particularly in the adversarial spaces, then this regularity can serve to make adversarial samples more distinguishable from clean ones, and therefore easier to detect. We show experimentally that the behavior of the network will be more distinguishable between clean and adversarial samples in practice in Section 5.4.1.

So as in Section 3.4, we train our network by solving

$$ 
\begin{array}{ll}
\min_{\theta} & \mathcal{C}(\theta) = \sum_{x \in \mathcal{D}} \sum_{m=0}^{M-1} \|r_m(\phi_m^x)\|^p \\
\text{subject to} & \phi_{m+1}^x = \phi_m^x + h r_m(\phi_m^x) \text{ and } \phi_0^x = x \text{ for all } x \in \mathcal{D} \\
& \mathcal{L}(\theta) = 0
\end{array}
\qquad (35)
$$
which is the discretization of

$$
\begin{array}{ll}
\inf_{v,F} & \displaystyle \int_0^1 \|v_t\|_{L^p((\phi_t)_\sharp \mu)}^p \, dt \\
\text{subject to} & \partial_t \phi_t^x = v_t(\phi_t^x) \qquad (36) \\
& \phi_0^\cdot = \text{id} \\
& (\phi_1^\cdot)_\sharp \mu \in S_{F,\mathcal{L}}
\end{array}
$$

which is equivalent to

$$
\begin{array}{ll}
\inf_{T,F} & \mathcal{C}(T) = \int_{\Omega} c(x, T(x)) \, d\mu(x) \qquad (37) \\
\text{subject to} & T_{\sharp}\mu \in S_{F,\mathcal{L}}
\end{array}
$$

We solve again (35) through the method of multipliers

$$
\left\{
\begin{array}{l}
\theta_{i+1} = \arg \min_{\theta} \mathcal{C}(\theta) + \lambda_i \mathcal{L}(\theta) \\
\lambda_{i+1} = \lambda_i + \tau \mathcal{L}(\theta_{i+1})
\end{array}
\right.
\qquad (38)
$$

We had a regularity result on the transport part $T$ of the network in Section 3.4, but we now look at individual residual blocks. We first prove a regularity result on each residual $r_m$ that applies on the clean data points moving according to $v$ solution to (17).

We take $\Omega \subset \mathbb{R}^d$ convex and compact and the data distribution $\alpha \in \mathcal{P}(\Omega)$ absolutely continuous and such that $\delta\Omega$ is $\alpha$-negligible. We suppose that there exists an open bounded convex set $X \subset \Omega$ such that $\alpha$ is bounded away from zero and infinity on $X$ and is zero on its complement $X^{\mathcal{C}}$. From Theorem 2 in Section 3.3.2, we know that (37) and (36) are equivalent and have solutions $(T, F)$ and $(v, F)$ such that $T$ is an optimal transport map between $\alpha$ and $\beta := T_{\sharp}\alpha$. We suppose that $\beta$ is absolutely continuous and that there exists an open bounded convex set $Y \subset \Omega$ such that $\beta$ is bounded away from zero and infinity on $Y$ and is zero on $Y^{\mathcal{C}}$. In the rest of this section, $v$ solves (36) and we suppose that we find a solution to the discretized problem (35) that is an $\varepsilon/2$-approximation of $v$, i.e. $\|r_m - v_{t_m}\|_{\infty} \le \varepsilon/2$ for all $0 \le m < M$, with $t_m = m/M$.

In Theorem 6, we show that the regularization makes the residual blocks of the network $\eta$-Hölder (with an error of $\varepsilon$) on the support of the input distribution as it moves according to the theoretical vector field solution $v$. The results hold for all norms on $\mathbb{R}^d$.

**Theorem 6.** For $a, b \in \text{support}(\alpha_{t_m})$ where $\alpha_t := (\phi_t)_\sharp \alpha$ with $\phi$ solving (36) along with $v$, we have

$$
\|r_m(a) - r_m(b)\| \le \varepsilon + K \|a - b\|^{\zeta_1} \text{ if } \|a - b\| \le 1
$$

$$
\|r_m(a) - r_m(b)\| \le \varepsilon + K \|a - b\|^{\zeta_2} \text{ if } \|a - b\| > 1
$$

for constants $K > 0$ and $0 < \zeta_1 \le \zeta_2 \le 1$,
*Proof.* A solution $v$ to (36) exists and is linked to an optimal transport map $T$ that is a solution to (37) through $v_t = (T - \mathrm{id}) \circ T_t^{-1}$ where $T_t := (1 - t)\mathrm{id} + tT$ which is invertible (see Appendix B). By Theorem 1 in Appendix B, being an optimal transport map, $T$ is $\eta$-Hölder on $X$. So for all $t \in [0, 1[$ and all $a, b \in \mathrm{support}(\alpha_t)$, where $\alpha_t = (\phi_t)_{\#}\alpha = (T_t)_{\#}\alpha$ with $\phi$ solving (36) together with $v$, we have

$$ \|v_t(a) - v_t(b)\| \le \|T_t^{-1}(a) - T_t^{-1}(b)\| + C \|T_t^{-1}(a) - T_t^{-1}(b)\|^{\eta} \quad (39) $$

Since $(\alpha_t)_{t=0}^1$ is a geodesic between $\alpha$ and $\beta = \alpha_1 = T_{\#}\alpha$, then $(\alpha_s)_{s=0}^t$ is a geodesic between $\alpha$ and $\alpha_t$ (modulo reparameterization to $[0, 1]$). And since $\alpha_s = (T_s)_{\#}\alpha$, the map $T_t$ is an optimal transport map between $\alpha$ and $\alpha_t$. Therefore its inverse $T_t^{-1}$ is an optimal transport map (see Theorem 25 in Appendix B) and is $\eta_t$-Hölder with $0<\eta_t\le 1$ (being a push-forward by $T_t$, the support of $\alpha_t$ satisfies the conditions of Theorem 1 in Appendix B). Therefore, for all $a, b \in \mathrm{support}(\alpha_t)$

$$ \|v_t(a) - v_t(b)\| \le C_t \|a - b\|^{\eta_t} + CC_t^{\eta} \|a - b\|^{\eta\eta_t} \quad (40) $$

and for all $a, b \in \mathrm{support}(\alpha_{t_m})$

$$ \|r_m(a) - r_m(b)\| \le \varepsilon + C_{t_m} \|a - b\|^{\eta_{t_m}} + CC_{t_m}^{\eta} \|a - b\|^{\eta\eta_{t_m}} \quad (41) $$

by the hypothesis on $r$ and the triangle inequality. Let $K := \max_m C_{t_m} + CC_{t_m}^{\eta}$, $\zeta_1 := \eta \min_m \eta_{t_m}$ and $\zeta_2 := \max_m \eta_{t_m}$. Then, we have the desired result immediately from (41). $\square$

We use Theorem 6 to now bound the distance between the residues at depth $m$ as a function of the distance between the network's inputs. For inputs $a_0$ and $b_0$ to the network, the intermediate embeddings are $a_{m+1} = a_m + hr_m(a_m)$ and $b_{m+1} = b_m + hr_m(b_m)$, and the residues used to compute features for adversarial detection are $r_m(a_m)$ and $r_m(b_m)$. So we want to bound $\|r_m(a_m) - r_m(b_m)\|$ as a function of $\|a_0 - b_0\|$. This is usually done by multiplying the Lipschitz constants of each block up to depth $m$, which leads to an overestimation [Huster et al., 2019], or through more complex estimation algorithms [Virmaux and Scaman, 2018, Latorre et al., 2020, Bhowmick et al., 2021]. Bound (39) allows through $T_t^{-1}$ to avoid multiplying the Hölder constants of the blocks. If $a_0$ and $b_0$ are on the clean data support $X$, we get Theorem 7, which we prove together with Theorem 8 below.

**Theorem 7.** For $a_0, b_0 \in X$ and constants $C, L>0$,

$$ \|r_m(a_m) - r_m(b_m)\| \le \varepsilon + \|a_0 - b_0\| + C\|a_0 - b_0\|^{\eta} + L(\|a_m - \phi_{t_m}^{a_0}\| + \|b_m - \phi_{t_m}^{b_0}\|) $$

The term $\mu(a_0):=\|a_m-\phi_{t_m}^{a_0}\|$ (and $\mu(b_0):=\|b_m-\phi_{t_m}^{b_0}\|$) is the distance between the point $a_m$ after $m$ residual blocks and the point $\phi_{t_m}^{a_0}$ we get by following the
theoretical solution vector field $v$ up to time $t_m$ starting from $a_0$. If $a_0$ and $b_0$ are not on the data support $X$, an extra term has to be introduced to use bound (39). Bounding the terms $\mu(a_0)$ and $\mu(b_0)$ is possible under more regularity assumptions on $v$. We assume then that $v$ is $C^1$ and Lipschitz in $x$, which is not stronger than the regularity we get on $v$ through our regularization, as it does not give a similar result to bound (39). We have for all inputs $a_0$ and $b_0$, whether they are clean or not, Theorem 8:

**Theorem 8.** For $a_0, b_0 \in \mathbb{R}^d$ and constants $R, S > 0$,

$$ \|r_m(a_m) - r_m(b_m)\| \le \varepsilon + LS\varepsilon + LSRh + \|a_0 - b_0\| + C\|a_0 - b_0\|^\eta + \\ + LS(\mathrm{dist}(a_0, X) + \mathrm{dist}(b_0, X)) $$

*Proof.* Since $T_t^{-1}(\phi_t^x) = x$, we have for any $a_0, b_0 \in X$ by the triangle inequality

$$ \begin{aligned} \|r_m(a_m) - r_m(b_m)\| &\le \|r_m(a_m) - r_m(\phi_{t_m}^{a_0})\| + \|r_m(\phi_{t_m}^{a_0}) - v_{t_m}(\phi_{t_m}^{a_0})\| + \\ &\quad + \|v_{t_m}(\phi_{t_m}^{a_0}) - v_{t_m}(\phi_{t_m}^{b_0})\| + \|r_m(\phi_{t_m}^{b_0}) - v_{t_m}(\phi_{t_m}^{b_0})\| + \\ &\quad + \|r_m(b_m) - r_m(\phi_{t_m}^{b_0})\| \\ &\le \varepsilon + \|a_0 - b_0\| + C\|a_0 - b_0\|^\eta + \\ &\quad + L(\|a_m - \phi_{t_m}^{a_0}\| + \|b_m - \phi_{t_m}^{b_0}\|) \end{aligned} $$

where $L = \max_m L_m$ and $L_m$ is the Lipschitz constant of $r_m$ (which is Lipschitz being a composition of matrix multiplications and activations such as ReLU). This the bound in Theorem 7.

In this bound, the term $\|a_m - \phi_{t_m}^{a_0}\|$ (and likewise $\|b_m - \phi_{t_m}^{b_0}\|$) represents the distance between the point $a_m$ we get after $m$ residual blocks (i.e. after $m$ Euler steps using approximation $r$ of $v$) and the point $\phi_{t_m}^{a_0}$ we get by following the solution vector field $v$ up to time $t_m$. By the properties of the Euler scheme (consistency and zero-stability, see Corollaries 11, 13 and 15 in Section A.1), under more conditions on $v$, it is possible to bound this term. Indeed, if $v$ is $C^1$ and $M$-Lipschitz in $x$ (this is not stronger than the regularity we get on $v$ through our regularization, because we still need to use (39)), we have for constants $R, S > 0$,

$$ \|\phi_{t_m}^{a_0} - a_m\| \le \|\phi_{t_m}^{a_0} - \tilde{a}_m\| + \|\tilde{a}_m - a_m\| \le S\varepsilon + SRh $$

where $\tilde{a}_m$ comes from the Euler scheme with access to $v$ ($\tilde{a}_{m+1}:=\tilde{a}_m+h v_{t_m}(\tilde{a}_m)$ and $\tilde{a}_0:=a_0$), $R$ is the consistency constant of the Euler method and $S$ is its zero-stability constant. If $a_0, b_0 \notin X$, we need to introduce $\hat{a}_0 := \mathrm{Proj}_X(a_0)$ and $\hat{b}_0 := \mathrm{Proj}_X(b_0)$ to apply (39). We now get

$$ \|r_m(a_m) - r_m(b_m)\| \le \varepsilon + \|a_0 - b_0\| + C\|a_0 - b_0\|^\eta + L(\|a_m - \phi_{t_m}^{\hat{a}_0}\| + \|b_m - \phi_{t_m}^{\hat{b}_0}\|) $$

Bounding the terms $\|a_m - \phi_{t_m}^{\hat{a}_0}\|$ and $\|b_m - \phi_{t_m}^{\hat{b}_0}\|$ now gives

$$ \|\phi_{t_m}^{\hat{a}_0} - a_m\| \le \|a_m - \tilde{a}_m\| + \|\tilde{a}_m - \phi_{t_m}^{\hat{a}_0}\| \le S(\|a_0 - \hat{a}_0\| + \varepsilon) + SRh $$
where $\tilde{a}_m$ now comes from the Euler scheme with access to $v$ that starts at $\hat{a}_0$ (meaning $\tilde{a}_{m+1}:=\tilde{a}_m+h\hat{v}_{t_m}(\tilde{a}_m)$ and $\tilde{a}_0:=\hat{a}_0$). Likewise, we get the same bound for $\|b_m - \phi_{t_m}^{\tilde{b}_0}\|$.

Since $\|a_0 - \hat{a}_0\| = \mathrm{dist}(a_0, X)$ and $\|b_0 - \hat{b}_0\| = \mathrm{dist}(b_0, X)$, we get the bound in Theorem 8. Note that if we use the stability of the differential equation instead of the stability of the Euler method to bound $\|a_m - \phi_{t_m}^{\hat{a}_0}\|$ we get the same result. Indeed, if $a_m$ again comes from the Euler scheme with access to $v$ that starts at $a_0$ (meaning $\tilde{a}_{m+1}:=\tilde{a}_m+h\hat{v}_{t_m}(\tilde{a}_m)$ and $\tilde{a}_0:=a_0$), we can write, for some constant $F > 0$

$$ \|\phi_{t_m}^{\hat{a}_0} - a_m\| \le \|a_m - \tilde{a}_m\| + \|\tilde{a}_m - \phi_{t_m}^{\hat{a}_0}\| + \|\phi_{t_m}^{\hat{a}_0} - \phi_{t_m}^{\hat{a}_0}\| \le S\epsilon + SRh + F\|a_0 - \hat{a}_0\| $$

since

$$ \begin{aligned} \|\phi_{t_m}^{\hat{a}_0} - \phi_{t_m}^{\hat{a}_0}\| &\le \|a_0 - \hat{a}_0\| + \int_0^{t_m} \|v_s(\phi_s^{\hat{a}_0}) - v_s(\phi_s^{\hat{a}_0})\| \mathrm{d}s \\ &\le \|a_0 - \hat{a}_0\| + M \int_0^{t_m} \|\phi_s^{\hat{a}_0} - \phi_s^{\hat{a}_0}\| \mathrm{d}s \\ &\le F \|a_0 - \hat{a}_0\| \end{aligned} $$

where we get the last line by Gronwall's lemma. □

The terms $\mathrm{dist}(a_0, X)$ and $\mathrm{dist}(b_0, X)$ in Theorem 8 show that the regularity guarantee is increased for inputs in $X$. The trajectories of clean points are then closer to each other and more different from those of abnormal samples outside $X$ as in Figure 22 in Section 5.4.1.

## 5.4 Experiments

We evaluate our method on adversarial samples found by 8 adversarial attacks. The threat model is as follows. We use 2 popular white-box attacks that have access to the network and to its weights and architecture but not to its training data: FGM [Goodfellow et al., 2015b], PGD [Madry et al., 2018] (specifically the Auto-PGD-CE (APGD) variant [Croce and Hein, 2020b]), BIM [Kurakin et al., 2017], DF [Moosavi-Dezfooli et al., 2016], CW [Carlini and Wagner, 2017b], and AA [Croce and Hein, 2020b], and 2 black-box attacks that can only query the network: HSJ [Chen et al., 2020a] and BA [Brendel et al., 2018]. We assume the attacker has no access to the detector and use the untargeted (i.e. not seeking to direct the mistake towards a particular class) versions of the attacks.

For FGM, APGD and BIM, we use a maximal perturbation size of $\epsilon = 0.03$. We use the $L_2$ norm for CW and HSJ and the $L_\infty$ norm for all the other attacks. We use the Adversarial Robustness Toolbox (ART) Python package [Nicolae et al., 2018] to generate the adversarial samples, except for AA for which we use the authors' original code [Croce and Hein, 2020b]. We use the default
values from the ART package and the code of [Croce and Hein, 2020b] for the
hyper-parameters of the attacks that are not mentioned here. The number of
iterations is 50 for HSJ, 5000 for BA, 10 for CW and 100 for APGD and DF.

We compare our detector (which we call the Transport detector or TR in the
tables below) to the Mahalanobis detector (MH in the tables below) of [Lee
et al., 2018] and to the detector of [Kherchouche et al., 2020, Kherchouche et al.,
2022] that uses natural scene statistics (NS in the tables below). We compare
our transport regularization to reverse cross entropy (RCE) training of [Pang
et al., 2018], which is also meant to improve detection of adversarial examples.

We use three architectures and datasets for our experiments: a ResNeXt50
trained on CIFAR100, a ResNet110 trained on CIFAR10 and a WideResNet
trained on TinyImageNet. For each architecture-dataset pair, the network is
trained normally with cross entropy, with our regularization added to cross
entropy (called a LAP-network) and with reverse cross entropy instead of
cross entropy (called an RCE-network). For our regularization, we use (19)
with $\tau=1$, $s=1$ and $\lambda_0=1$ for all networks. These hyper-parameters are cho-
sen to improve validation accuracy during training and not adversarial detec-
tion. More training details are in Appendix H.1. The code is available at
github.com/advadvadvadvadv/adv.

We find that our detector consistently outperforms the other detectors, also
on out-of-distribution detection, and especially on generalization to unseen
attacks. LAP training mainly improves, sometimes largely, the accuracy of the
Mahalanobis detector on seen attacks, and consistently does better than RCE.

5.4.1 Preliminary Experiments

We find again that LAP training improves generalization. The vanilla ResNeXt50
has a test accuracy of 74.38% on CIFAR100 while the LAP-ResNeXt50 has an
accuracy of 77.2%. The vanilla ResNet110 has a test accuracy of 92.52% on
CIFAR10 while the LAP-ResNet110 has an accuracy of 93.52% and the RCE-
ResNet110 of 93.1%. The vanilla WideResNet has a test accuracy of 65.14% on
TinyImageNet while the LAP-WideResNet has an accuracy of 65.34%.

We see in Figure 21 below that regularizing the transport cost $C$ through LAP
training makes it more distinguishable between clean and adversarial samples.
Indeed, the histograms of $C$ on adversarial samples move away from the histogram
of $C$ on clean samples when using LAP training instead of vanilla training. Simply
using the empirical quantiles of $C$ on clean samples to reject abnormal samples
with a transport cost $C$ that falls outside of these quantiles allows then to detect
samples from some simple attacks with near perfect recall and a chosen fixed
false positive rate, without training on or even seeing adversarial samples. In
Figure 21, using the 0.02 and 0.98 empirical quantiles of $C$ on the clean samples
allows to detect adversarial samples from the FGM attack with big enough
maximal perturbation $\epsilon$ almost perfectly on the LAP-ResNeXt50 trained on
CIFAR100, but not on the vanilla ResNeXt50.
<figure><img src="image_26.png" /><figcaption>Figure 21: Histogram of the transport cost $C$ for clean and FGM-attacked test samples with different values of $\epsilon$ on CIFAR100. The vertical lines represent the 0.02 and 0.98 empirical quantiles of the transport cost of the clean samples. Left: ResNeXt50. Right: LAP-ResNeXt50.</figcaption></figure>

We visualize this increased difference in behavior between clean and out-of-distribution (OOD) points with LAP training in 2 dimensions in Figure 22. The experiments is binary classification of the circles dataset from scikit-learn [Pedregosa et al., 2011], where we add OOD points at test time.

<figure><img src="image_27.png" /><figcaption>Figure 22: Transformed circles test set from scikit-learn (red and blue) and out-of-distribution points (green) after blocks 6 and 9 of a ResNet with 9 blocks. We use LAP training in the second row, which makes the movements of the clean points (red and blue) more similar to each other and more different from the movements of the green out-of-distribution points than when using the vanilla network in the first row. In particular, without the regularization, the green points are closer to the clean red points after blocks 6 and 9 which is undesirable.</figcaption></figure>
5.4.2 Detection of Seen Attacks

For detection training, the test set is split in 0.9/0.1 proportions into two
datasets, B1 and B2. For each image in B1 (respectively B2), an adversarial
sample is generated and a balanced detection training set (respectively detection
test set) is created. Since adversarial samples are created for a specific network,
this is done for the vanilla network and its LAP and RCE versions. We tried
augmenting the detection training dataset with a random perturbation of each
image, to be considered as clean, as in [Lee et al., 2018], but we found that this
usually does not improve detection accuracy. This dataset creation protocol
is standard and is depicted in Figure 23 below. We tried limiting the dataset
to successfully attacked images only as in [Lee et al., 2018], but did not find a
consistent improvement in detection, even only on successfully attacked images.

<figure><img src="image_28.png" /><figcaption>Figure 23: Adversarial detection dataset creation. A∪B1∪B2 is the original dataset, where A is the training set and B1∪B2 is the test set. We create a noisy version of B1∪B2 by adding random noise to each sample in B1∪B2 to get C1∪C2. Noisy samples are considered clean (i.e. not attacked) in adversarial detection training. We create an attacked version of B1∪B2 by creating an attacked image from each image in B1∪B2 to get D1∪D2. In the case of generalization to unseen attacks, Attack 2 used to create D2 from B2 is different from Attack 1 used to create D1 from B1. Otherwise, Attack 1 and Attack 2 are the same. B1∪C1∪D1 is the adversarial detection training set and B2∪C2∪D2 is the adversarial detection test set.</figcaption></figure>
Samples in the detection training set are fed through the network and the features of each detection method are extracted. We tried three classifiers (logistic regression, random forest and SVM from scikit-learn [Pedregosa et al., 2011]) trained on these features for each detection method, and kept the random forest as it always performs best.

We tried two methods to improve all detectors: class-conditioning and ensembling. In class-conditioning, the features are grouped by the class predicted by the network, and a detector is trained for every class. At test time, the detector trained on the features of the predicted class is used. A detector is also trained on all samples regardless of the predicted class and is used in case a certain class is never targeted by the attack. We also tried ensembling the class-conditional detector with the general all-class detector: an input is considered an attack if at least one detector says so. This ensemble of the class-conditional detector and the general detector performs best for all detection methods, and is the one we use.

We report the detection accuracy of each detector on the detection test set for both the vanilla version of the network and its LAP version in Table 12 below. In each cell, the first number corresponds to the vanilla network and the second to the regularized LAP-network. Since the NS detector takes the image and not its embeddings as input, the impact of LAP and RCE training on its performance is minimal and we report its performance on the vanilla network only.

These results are averaged over 5 runs and the standard deviations are in Tables 32, 33 and 34 in Appendix H.2, along with results on RCE-networks. Since they are much slower to generate, we test detection of black-box attacks HSJ and BA only on the network-data pair of ResNet110-CIFAR10 and of the DF and the CW attacks only on ResNet110-CIFAR10 and ResNeXt50-CIFAR100.

Results in Table 12 show two things. First, our detector performs better than both other detectors, with or without the regularization. Second, both the TR and MH detectors work better on the LAP-networks most times. The MH detector benefits more from the regularization, but on all attacks, the best detector is always the Transport detector. In Tables 32 and 33 in Appendix H.2, RCE often improves detection accuracy in this experiment, but clearly less than LAP training.

Results for ResNet110 on CIFAR10 are in the first row of Table 12. Our detector outperforms the Mahalanobis detector by 8 to 16 percentage points depending on the attack on the vanilla ResNet110, and the NS detector by up to 5 percentage points. LAP training improves the detection accuracy of our detector by 1.5 percentage points on average and that of the Mahalanobis detector by a substantial 8.2 percentage points on average.

Results for ResNeXt50 on CIFAR100 are in the second row of Table 12. Our detector outperforms the Mahalanobis detector by 1 to 5 percentage points
depending on the attack on the vanilla ResNeXt50, and the NS detector by up to 3 percentage points. LAP training improves the detection accuracy of both detectors by 1 percentage point on average.

Results for a WideResNet on TinyImageNet are in the third row of Table 12. Our detector outperforms the Mahalanobis detector by 3 to 15.5 percentage points depending on the attack on the vanilla WideResNet, and the NS detector slightly. LAP training improves the detection accuracy of the Mahalanobis detector by 0.6 percentage points on average.

Detection rates of successful adversarial samples (i.e. those that fool the network) are in Table 38 in Appendix H.4 and are higher than 95% on our detector. False positive rates are in Table 40 in Appendix H.5 and are always less than 5% on our detector. The AUROC is in Table 42 in Appendix H.6. On all these metrics, our detector outperforms the other detectors, and LAP-training greatly improves the performance of the Mahalanobis detector.

<table><caption>Table 12: Average detection accuracy of seen attacks on Network/LAP-Network.</caption><thead><tr><th rowspan="2">Network/Data</th><th rowspan="2">Det</th><th colspan="8">Attack</th></tr><tr><th>FGM</th><th>APGD</th><th>BIM</th><th>AA</th><th>DF</th><th>CW</th><th>HSJ</th><th>BA</th></tr></thead><tbody><tr><td rowspan="3">ResNet110 CIFAR10</td><td>TR</td><td>97.1/98.7</td><td>94.1/97.5</td><td>97.5/99.3</td><td>88.9/94.1</td><td>99.9/99.8</td><td>98.0/98.0</td><td>99.9/99.9</td><td>96.6/97.0</td></tr><tr><td>MH</td><td>87.8/95.6</td><td>82.1/90.7</td><td>86.8/95.4</td><td>80.5/90.0</td><td>91.5/96.7</td><td>85.6/93.4</td><td>85.5/94.6</td><td>80.2/89.6</td></tr><tr><td>NS</td><td>94.6</td><td>94.3</td><td>95.0</td><td>88.8</td><td>99.8</td><td>93.9</td><td>99.7</td><td>92.1</td></tr><tr><td rowspan="3">ResNeXt50 CIFAR100</td><td>TR</td><td>97.3/98.3</td><td>96.0/97.8</td><td>98.0/98.9</td><td>84.9/87.6</td><td>99.8/99.6</td><td>97.0/97.8</td><td></td><td></td></tr><tr><td>MH</td><td>95.8/96.8</td><td>93.9/94.6</td><td>96.1/97.8</td><td>83.9/86.6</td><td>97.3/97.1</td><td>95.4/96.4</td><td></td><td></td></tr><tr><td>NS</td><td>94.7</td><td>94.2</td><td>94.7</td><td>84.8</td><td>99.6</td><td>90.7</td><td></td><td></td></tr><tr><td rowspan="3">WideResNet TinyImageNet</td><td>TR</td><td>95.4/95.1</td><td>95.2/95.2</td><td>95.3/95.1</td><td>81.4/81.2</td><td></td><td></td><td></td><td></td></tr><tr><td>MH</td><td>81.1/82.3</td><td>79.7/80.6</td><td>81.2/82.5</td><td>78.4/78.4</td><td></td><td></td><td></td><td></td></tr><tr><td>NS</td><td>94.9</td><td>94.9</td><td>95.0</td><td>81.3</td><td></td><td></td><td></td><td></td></tr></tbody></table>

### 5.4.3 Detection of Unseen Attacks

An important setting is when we don't know which attack might be used at test time or only have time to train on one attack. We still want our detector to generalize well to other attacks. To test this, we use the same vanilla networks as above but only train the detectors on the detection training set created using the simplest and quickest attack (FGM) and test them on the detection test sets created using the six other attacks.

Results are in Table 13 below. We see that our detection method has very good detection generalization to unseen attacks, even those very different from FGM, comfortably better than the Mahalanobis detector, by up to 19 percentage points, while the NS detector only generalizes to variants of FGM (APGD and BIM), and fails on the other attacks. These results are averaged over 5 runs and the standard deviations are in Tables 35, 36 and 37 in Appendix H.3.