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

Introduction

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 [START_REF] Russell | Artificial Intelligence: A Modern Approach[END_REF]. 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 [START_REF] Goodfellow | Deep Learning[END_REF]. It has become the state-of-the-art in nearly all tasks, from classification to generation, in the fields of computer vision [START_REF] Krizhevsky | Imagenet classification with deep convolutional neural networks[END_REF], natural language [START_REF] Wu | Google's neural machine translation system: Bridging the gap between human and machine translation[END_REF], sound [START_REF] Purwins | Deep learning for audio signal processing[END_REF], speech [START_REF] Mehrish | A review of deep learning techniques for speech processing[END_REF] and video [START_REF] Sharma | Video processing using deep learning techniques: A systematic literature review[END_REF] processing, and has been applied successfully to other tasks such as dynamical system forecasting [de Bezenac et al., 2018], biology [START_REF] Jumper | Highly accurate protein structure prediction with alphafold[END_REF] and reinforcement learning [START_REF] Silver | Mastering the game of go with deep neural networks and tree search[END_REF].

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 [START_REF] Rumelhart | Learning representations by back-propagating errors[END_REF] and stochastic gradient descent [Bottou, 2012, Amari, 1967], or a variant of it [START_REF] Bottou | Optimization methods for large-scale machine learning[END_REF], for training the model, non-linear activations (ReLU [Fukushima, 1975] very often), normalization layers (most often batch normalization [START_REF] Ioffe | Batch normalization: Accelerating deep network training by reducing internal covariate shift[END_REF] or layer normalization [START_REF] Ba | Layer normalization[END_REF]) and the importance of a good initialization [START_REF] Saxe | Exact solutions to the nonlinear dynamics of learning in deep linear neural network[END_REF][START_REF] He | Delving deep into rectifiers: Surpassing human-level performance on imagenet classification[END_REF]. Among these ever-present elements is the architectural component known as a skip (or residual) connection [START_REF] He | Deep residual learning for image recognition[END_REF], 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.

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.

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

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 [START_REF] Karkar | A principle of least action for the training of neural networks[END_REF]. 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., 2023a], 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 [START_REF] Karkar | Adversarial sample detection through neural network transport dynamics[END_REF], 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.

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.

Residual Networks

Residual networks (ResNets) [START_REF] He | Deep residual learning for image recognition[END_REF], He et al., 2016a] and their variants such as ResNeXt [Xie et al., 2017] and WideResNet [START_REF] Zagoruyko | Wide residual networks[END_REF] 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 [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF], 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 [START_REF] Zhu | Unpaired image-to-image translation using cycle-consistent adversarial networks[END_REF], image segmentation [START_REF] Pohlen | Full-resolution residual networks for semantic segmentation in street scenes[END_REF], Chen et al., 2017a] and action recognition [START_REF] Tran | A closer look at spatiotemporal convolutions for action recognition[END_REF]. They have also been used with success in speech and natural language processing applications [START_REF]Google's neural machine translation system: Bridging the gap between human and machine translation[END_REF], Xiong et al., 2017, van den Oord et al., 2016].

Furthermore, residual connections are a basic building block in many state-ofthe-art architectures such as EfficientNet [START_REF] Tan | EfficientNet: Rethinking model scaling for convolutional neural networks[END_REF] and MobileNetV2 [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF], which use a so-called inverted residual block. Residual connections have also been added to prior architectures such as Inception to improve their performances [START_REF] Szegedy | Inception-v4, inception-resnet and the impact of residual connections on learning[END_REF], and are used in some U-Net architectures [START_REF] Janner | Planning with diffusion for flexible behavior synthesis[END_REF]. 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 [START_REF] Liu | Swin transformer: Hierarchical vision transformer using shifted windows[END_REF] and ViT [START_REF] Dosovitskiy | An image is worth 16x16 words: Transformers for image recognition at scale[END_REF], video Transformers such as TimeSFormer [START_REF] Bertasius | Is space-time attention all you need for video understanding[END_REF], and natural language processing Transformers such as GPT [START_REF] Radford | Improving language understanding with unsupervised learning[END_REF] 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 ) (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 •f 2 •f 3 •.
..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 [START_REF] Simonyan | Very deep convolutional networks for large-scale image recognition[END_REF]. 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.

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. The other common residual block architecture is called the bottleneck residual block [START_REF] He | Deep residual learning for image recognition[END_REF], 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 × 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. et al., 2017]. A convolutional layer is shown as (number of input channels, filter size, number of output channels).

Finally, we mention the inverted residual block. It was proposed initially for the MobileNetV2 [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF] 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. From [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF]. The thickness indicates the number of channels.

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.

7

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. 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, [START_REF] Dosovitskiy | An image is worth 16x16 words: Transformers for image recognition at scale[END_REF]) is very similar. [START_REF] Liu | Swin transformer: Hierarchical vision transformer using shifted windows[END_REF]. MLP: multi-layer perception. LN: layer normalization. W-MSA: window multi-head self-attention. SW-MSA: shifted window multi-head self-attention.

Interpretations of Residual Networks

The initial intuition behind the design of the ResBlock in [START_REF] He | Deep residual learning for image recognition[END_REF] 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 [START_REF] Balduzzi | The shattered gradients problem: If resnets are the answer, then what is the question[END_REF], 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 [START_REF] Veit | Residual networks behave like ensembles of relatively shallow networks[END_REF], 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 ))).

[ [START_REF] Huang | Learning deep resnet blocks sequentially using boosting theory[END_REF] combine this point of view with boosting theory [START_REF] Freund | A decisiontheoretic generalization of on-line learning and an application to boosting[END_REF] to propose an algorithm for sequentially training residual networks, as we do in Section 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 [START_REF] Greff | Highway and residual networks learn unrolled iterative estimation[END_REF], 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 [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF] 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:

L(x M ) = L(x M -1 + r M -1 (x M -1 )) = L(x M -1 ) + ∇L(x M -1 ) • r M -1 (x M -1 ) + O(r 2 M -1 (x M -1 )) = L(x m ) + m≤i<M ∇L(x i ) • r i (x i ) + O(r 2 i (x i ))
which holds for any 0 ≤ 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 ∇L(x i ) • r i (x i ), which can be done by making r i (x i ) point in the opposite half-space to ∇L(x i ). The block operations x i + r i (x i ) are therefore moving the input towards the half-space of -∇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 ∇L(x i ) • 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 [START_REF] Greff | Highway and residual networks learn unrolled iterative estimation[END_REF], Hauser, 2019, Jastrzebski et al., 2018[START_REF]Multi-level residual networks from dynamical systems view[END_REF]] 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 ∥ri(xi)∥ ∥xi∥ and the cosine loss ∇L(xi)•ri(xi) ∥∇L(xi)∥∥ri(xi)∥ on CIFAR10 of 4 ResNets.

Figure 6: Ratio of residue norm to input norm (first row) and cosine loss (second row) of 4 residual architectures on CIFAR10 from [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF].

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 andWynter, 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 [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF], 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[START_REF] Hayou | [END_REF], 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 [START_REF] Hayou | Regularization in resnet with stochastic depth[END_REF], and of courses reduces training time.

In [START_REF] Huang | Deep networks with stochastic depth[END_REF], 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, [START_REF] De | Batch normalization biases residual blocks towards the identity function in deep networks[END_REF]] 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.

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 : R d → R d , consider the initial value problem of finding a function x :

I → R d that satisfies x ′ (t) = v t (x(t)) (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 + t t0 v t (x(t)) dt (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≤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 tn (x n ) (4)
starting from the initial condition x 0 . This method simply comes from a firstorder 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 tn+1 (x n+1 ) (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 ∥ --------→ hmax→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.

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 ) ←→ x ′ (t) = v t (x(t)) (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 [START_REF] Lu | Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations[END_REF], implicit-explicit methods [START_REF] Haber | Imexnet -a forward stable deep neural network[END_REF], and Runge-Kutta methods [START_REF] Zhu | Convolutional neural networks combined with runge-kutta methods[END_REF]. We present some of these architectures below.

The IMEX architecture [START_REF] Haber | Imexnet -a forward stable deep neural network[END_REF] considers a parametrized differential equation x ′ (t) = v(x(t); θ(t)) where v is a neural network parametrized by θ and discretizes it using the implicit-explicit numerical method [START_REF] Ascher | Implicit-explicit methods for time-dependent partial differential equations[END_REF], Ascher et al., 1997, Ascher and Petzold, 1998]. This method rewrites the differential equation as follows x ′ (t) = v(x(t); θ(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); θ(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 = (Id + hA) -1 (x n + hAx n + hv(x n , θ n )) (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 [START_REF] Haber | Imexnet -a forward stable deep neural network[END_REF], 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 (Id + hA) -1 , since this matrix is dense. Experiments on tasks where the field of view is important such as image segmentation [START_REF] Minaee | Image segmentation using deep learning: A survey[END_REF] and depth estimation [START_REF] Masoumian | Monocular depth estimation using deep learning: A review[END_REF] confirm this second advantage over ResNets. [START_REF] Lu | Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations[END_REF] build architectures called LM-ResNet and LM-ResNeXt that mimic linear multi-step numerical schemes [Dahlquist, 1963], while FractalNet [START_REF] Larsson | Fractalnet: Ultra-deep neural networks without residuals[END_REF] and RKCNN [START_REF] Zhu | Convolutional neural networks combined with runge-kutta methods[END_REF]] mimic Runge-Kutta numerical methods (see Appendix A.2 for a definition of these methods). PolyNet [START_REF] Zhang | Polynet: A pursuit of structural diversity in very deep networks[END_REF] use a polynomial approximation to the implicit Euler method [START_REF]Swin Transformer architecture with 4 residual stages[END_REF] 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 [START_REF] Ruthotto | Deep neural networks motivated by partial differential equations[END_REF]. 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 [START_REF] Haber | Stable architectures for deep neural networks[END_REF], and to enforcing particular conditions on the weight matrices [START_REF] Ciccone | Nais-net: Stable deep networks from non-autonomous differential equations[END_REF].

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 [START_REF] Gomez | The reversible residual network: Backpropagation without storing activations[END_REF] 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 [START_REF] Jacobsen | i-revnet: Deep invertible networks[END_REF] are an extension of these RevNets.

13 Some of the architectures discussed above are illustrated in Figure 7 below.

Figure 7: Some architectures inspired by numerical schemes. From [START_REF] Lu | Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations[END_REF].

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 [START_REF] Chen | Neural ordinary differential equations[END_REF], 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 θ) 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 θ

L(x(1)) = L x 0 + 1 0 v(t, x(t); θ) dt ≈ L(Solver(v(., .; θ), x 0 )) (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 θ can be done in two ways. The first is to discretize then differentiate, i.e. to differentiate L(Solver(v(., .; θ), x 0 )) with respect to 14 θ 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 + 1 0 v(t, x(t); θ) dt) in θ.
The formula for this gradient includes an integral that is then also discretized through a numerical solver.

In [START_REF] Li | Maximum principle based algorithms for deep learning[END_REF], the parameter θ in ( 8) is allowed to depend on t and the resulting problem is looked at as the optimal control problem [START_REF] Athans | Optimal Control: An Introduction to the Theory and Its Applications[END_REF] of finding a control function θ(t):[0, 1] → Θ that is going to guide the solution curve x to the ordinary differential equation x ′ (t) = v(t, x(t); θ(t)) towards a 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, [START_REF] Chernousko | Method of successive approximations for solution of optimal control problems[END_REF]) 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.

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 → Y and a measure µ on X, the pushforward of µ by T is the measure

T ♯ µ on Y defined by T ♯ µ(A) = µ(T -1 (A)) for every measurable set A ⊂ Y .
The metric Wasserstein space W 2 (Ω) := (P(Ω), W 2 ), with Ω a convex and compact subset of R d , is the space P(Ω) of probability measures over Ω, equipped with the distance W 2 given by the solution to the Kantorovitch problem

W 2 2 (µ, ν) = inf γ∈Π(µ,ν) Ω×Ω ∥x -y∥ 2 dγ(x, y) (9)
where Π(µ, ν) is the set of probability distribution over Ω × Ω with first marginal equal to µ and second marginal equal to ν, i.e

Π(µ, ν) = {γ ∈ P(Ω × Ω) | π 1♯ γ = µ, π 2♯ γ = ν}
where π 1 (x, y) = x and π 2 (x, y) = y are projections. Here, γ ∈ Π(µ, ν) is called a transport plan between µ and ν, and γ(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). If µ is absolutely continuous and ∂Ω is µ-negligible then the minimization problem in (9) has a unique solution and is equivalent to the Monge problem

W 2 2 (µ, ν) = min T s.t. T ♯ µ=ν Ω ∥T (x) -x∥ 2 dµ(x) (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 ♯ µ = ν is called a transport map between µ and ν.

Figure 9: The Monge problem from µ to ν. From [START_REF] Kolouri | Optimal mass transport: Signal processing and machine-learning applications[END_REF].

The Monge problem has a unique solution T ⋆ linked to the solution γ ⋆ of (9) through γ ⋆ = (id, T ⋆ ) ♯ µ, where (id, T ⋆ ) is the function x → (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 π(x, y).

Another equivalent formulation of the optimal transport problem in this setting is the dynamical formulation [START_REF] Benamou | A computational fluid mechanics solution to the monge-kantorovich mass transfer problem[END_REF]. Here, instead of directly pushing samples of µ to ν using T , we can equivalently displace mass, according to a continuous flow with velocity v t : R d → R d . This implies that the density µ t at time t satisfies the continuity equation

∂ t µ t + ∇ • (µ t v t ) = 0 (11)
assuming that initial and final conditions are given by µ 0 = µ and µ 1 = ν 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 (µ, ν) = min v 1 0 ∥v t ∥ 2 L 2 (µt) dt (12) s.t. ∂ t µ t + ∇ • (µ t v t ) = 0, µ 0 = µ, µ 1 = ν
Instead of describing the evolution of the density through the continuity equation, we can describe the paths ϕ x t taken by particles at position x from µ when displaced along the flow v. Here ϕ x t is the position at time t of the particle that was at x ∼ µ at time 0. The continuity equation is then equivalent to ∂ t ϕ x t = v t (ϕ x t ). Rewriting the conditions as necessary, Problem (12) becomes

W 2 2 (µ, ν) = min v 1 0 ∥v t ∥ 2 L 2 ((ϕ • t ) ♯ µ) dt (13) s.t. ∂ t ϕ x t = v t (ϕ x t ), ϕ • 0 = id, (ϕ • 1 ) ♯ µ = ν and the optimal T ⋆ that solves (10) is in fact T ⋆ (x) = ϕ x 1 for ϕ that solves ∂ t ϕ x t = v ⋆ t (ϕ x t )
together with the optimal v ⋆ from ( 13). This formulation is the most useful to us, as we will discretize ∂ t ϕ x t = v ⋆ t (ϕ x t ) to get a residual architecture. 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 η = 1:

Definition 2. A function f is η-Hölder on a set X ⊂ R d if ∀ a, b ∈ X, we have ∥f (a) -f (b)∥ ≤ C∥a -b∥ η
for some constants C>0 and 0<η≤1. We denote this f ∈ C 0,η (X). If all derivatives of f up to the k-th order are η-Hölder on X, we write f ∈ C k,η (X).

We can now state the following regularity theorem for optimal transport maps: Theorem 1. Suppose there are X, Y , bounded open sets in R d , such that the densities of µ and ν are null in their respective complements and bounded away from zero and infinity over them respectively.

Then, if Y is convex, there exists 0 < η ≤ 1 such that the optimal transport map T between µ and ν is in C 0,η (X).

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

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

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 [START_REF] Gai | A mathematical principle of deep learning: Learn the geodesic curve in the wasserstein space[END_REF], which was released shortly after our first work [START_REF] Karkar | A principle of least action for the training of neural networks[END_REF], 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, [START_REF] Finlay | How to train your neural ode[END_REF] 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 [START_REF] Sonoda | Transport analysis of infinitely deep neural network[END_REF] to analyze deep gaussian denoising autoencoders (not necessarily implemented through 18 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 [START_REF] Zhu | Unpaired image-to-image translation using cycle-consistent adversarial networks[END_REF], 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 [START_REF] Alvarez-Melis | Dataset dynamics via gradient flows in probability space[END_REF], 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 andBach, 2018, Chizat, 2020] used Wasserstein gradient flows in the parameter space to prove convergence of the training of two-layer neural networks.

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 [START_REF] Karkar | A principle of least action for the training of neural networks[END_REF] which won the runner-up best student paper award at ECML 2020.

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 counterintuitive. 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 [START_REF] Belkin | Reconciling modern machine-learning practice and the classical bias-variance trade-off[END_REF], 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 [START_REF] Belkin | To understand deep learning we need to understand kernel learning[END_REF], 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.

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

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

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 [START_REF] Pedregosa | Scikit-learn: Machine learning in Python[END_REF]) 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. 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.

Method and Theory

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 •T •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 ρ on R n , a transformation E : R n → Ω ⊂ R d is applied, transforming it into µ = E ♯ ρ, a distribution on Ω. 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 µ is transformed by a mapping T : Ω → Ω, 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 : R d → Y is applied to T ♯ µ in order to compute the loss 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. 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 ♯ µ needs to be linearly separable if F is linear. This model is quite general, as many ResNet-like architectures [START_REF] He | Deep residual learning for image recognition[END_REF], 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 [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF], the transport modules have similar behavior. The model can then be considered as a simplified ResNet, sometimes called a single representation ResNet. Note that [START_REF] Sandler | Non-discriminative data or weak model? on the relative importance of data and model resolution[END_REF] 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 [START_REF] Belkin | Reconciling modern machine-learning practice and the classical bias-variance trade-off[END_REF], 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 ♯ µ must be equal to some prescribed distribution ν which is the target of the generation process. But in general, T is less strictly constrained and the condition depends on F and L. This leads us to define a set of admissible targets for the task:

S F,L = {ν ∈ P(Ω) | L(F, ν) = 0} (14)
In general, 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:

Find (T, F ) such that T ♯ µ ∈ S F,L (15)
An important observation is that, even when S F,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 L being the cross entropy loss. The task consists in separating N classes. We denote µ i the class distributions which are supposed to be distributions in R d of mutually disjoint supports, meaning that there is no ambiguity in the class of data points, and such that µ = w i µ i with w i = 1 and w i ∈]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,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.

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 µ to S F,L and thus solve the task, one with a minimal transport cost. This gives us the following optimization problem: inf

T,F C(T ) = Ω c(x, T (x)) dµ(x) subject to T ♯ µ ∈ S F,L (16)
We consider ground costs c(x, y) = ∥x -y∥ p with p > 1, and suppose Ω compact with negligible boundary and µ absolutely continuous. We assume that the space of classifiers is compact, that the loss L is continuous, that the set ∪ F ∈F S F,L is at a finite p-Wasserstein distance W p from µ (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 L and of a class of functions F for the classifier F . The equivalent dynamical version of ( 16) is then, as per Section 2.8,

inf v,F 1 0 ∥v t ∥ p L p ((ϕ • t ) ♯ µ) dt subject to ∂ t ϕ x t = v t (ϕ x t ) ϕ • 0 = id (ϕ • 1 ) ♯ µ ∈ S F,L (17)
where

∥v t ∥ p L p ((ϕ • t ) ♯ µ) = R d ∥v t ∥ p d(ϕ • t ) ♯ µ.
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 ν ∈ S F,L at a finite distance from µ. Taking any transport map between µ and ν, we see that the infima are finite (see Appendix B.1 for existence of optimal transport maps).

Consider [START_REF]Module-wise training[END_REF] and take a minimizing sequence (T i , F i ) i . Set β i = (T i ) ♯ µ. 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 (µ, ν i ) ≤ C(T i ) ≤ M . So that (ν i ) i is a bounded sequence in ∪ F S F,L . By the hypothesized total boundedness of bounded subsets and as P p (R d ) endowed with W p is a complete metric space (see [Bolley, 2008] for a proof), up to an extraction, (ν i ) i converges to ν ⋆ in the closure of ∪ F S F,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 µ and ν ⋆ , we then have, by continuity of L,

T ⋆ ♯ µ = ν ⋆ ∈ S F ⋆ ,L
and C(T ⋆ ) ≤ lim C(T i ) by optimality of T ⋆ , which means, since (C(T i )) i is a minimizing sequence, that C(T ⋆ ) minimizes [START_REF]Module-wise training[END_REF]. So (T ⋆ , F ⋆ ) is a minimizer and T ⋆ is an optimal transport map.

Finally, there exists, (Appendix B.4), a velocity field v ⋆ t inducing the optimal transport map between µ and ν ⋆ which then gives a minimizer (v ⋆ , F ⋆ ) for [START_REF]Early overfitting of vanilla module-wise training[END_REF]. 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.

Note that uniqueness doesn't hold anymore, as the constraint T ♯ µ ∈ S F,L in [START_REF]Early overfitting of vanilla module-wise training[END_REF] 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.

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 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 [START_REF]Early overfitting of vanilla module-wise training[END_REF] 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 [START_REF]Module-wise training[END_REF] (or ( 17)) given by Theorem 2. Take X, respectively Y , an open neighborhood of the support of µ, respectively of T ⋆ ♯ µ, then T ⋆ is differentiable, except on a set of null µ 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 η-Hölder continuous from X\A to Y \B, i.e. for constants 0 < η ≤ 1 and C > 0 and ∀ x, y ∈ X \ A we have

∥T (x) -T (y)∥ ≤ C∥x -y∥ η
Proof. This is a consequence of Theorem 2, the hypothesis made in this section and the regularity theorems in Appendix B.5.

Remark 3. If Y in Proposition 1 is convex then T ⋆ is η-Hölder on all of X, i.e. T ∈ C 0,η (X) (see Theorem 1).
Remark 4. If the distributions µ and ν = T ⋆ ♯ µ in Proposition 1 are C k,η (i.e all derivatives up to the k-th derivative are η-Hölder), then the optimal transport map T is C k+1,η (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 α-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 η-Hölder continuity over as large a domain as possible, and even a diffeomorphism if the data distribution is well-behaved enough. Here η 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/η 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.

Practical Implementation

We propose here an algorithm for training ResNets using the least action principle by minimizing the kinetic energy. Starting from problem [START_REF]Early overfitting of vanilla module-wise training[END_REF] with p = 2 and the Euclidean norm, we first discretize the differential equation via a forward Euler scheme, which yields ϕ x m+1 = ϕ x m + r m (ϕ x m ). 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 θ. Next, the constraint

(ϕ • 1 ) ♯ µ ∈ S F,L is rewritten as L(F, (ϕ • 1 ) ♯ µ) = 0, denoted L(θ) = 0 below.
Finally, as we only have access to a finite set D of samples x from µ, we use a Monte-Carlo approximation of the integral with respect to the distributions (ϕ • t ) ♯ µ. This gives us

min θ C(θ) = x∈D M -1 m=0 ∥r m (ϕ x m )∥ p subject to ϕ x m+1 = ϕ x m + r m (ϕ x m ) and ϕ x 0 = x for all x ∈ D L(θ) = 0 (18)
Is is easy to see that the min-max problem min θ max λ>0 C(θ) + λ L(θ) yields the same solution as Problem 18, as the first two constraints are satisfied trivially. If the constraint L(θ) = 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 λ 0 and θ 0 :

θ i+1 = arg min θ C(θ) + λ i L(θ) λ i+1 = λ i + τ L(θ i+1 ) (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 θ 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 λ i , yielding the following algorithm:

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

Input: Training samples, step size τ , number of steps s, initial weight λ 0 Initialization: Initialize the parameters θ 0 and set i = 0 while not converged do 1. Starting from θ i , perform s steps of stochastic gradient descent:

1.1. θ 0 i+1 = θ i 1.2. θ l i+1 = θ l-1 i+1 -ϵ(∇C(θ l-1 i+1 )/λ i + ∇L(θ l-1 i+1 )) for l from 1 to s 1.3. θ i+1 = θ s i+1 2. Update the weight λ i+1 = λ i + τ L(θ i+1 ) and increment i ← i + 1 Output: Learned parameters θ
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 θ ⋆ , T θ ⋆ ) is reached by the optimization algorithm such that T θ ⋆ is an ϵ-OT map between µ and its push-forward (meaning ∥T θ ⋆ -T ⋆ ∥ ∞ ≤ ϵ where T ⋆ is the OT map). Then we have, with the same notations as in Proposition 1,

∀x, y ∈ X \ A, ∥T θ ⋆ (x) -T θ ⋆ (y)∥ ≤ O(ϵ + ∥x -y∥ η )
Proof. We simply write the decomposition:

T θ ⋆ (x) -T θ ⋆ (y) = T θ ⋆ (x) -T ⋆ (x) + T ⋆ (x) -T ⋆ (y) + T ⋆ (y) -T θ ⋆ (y)
and use the triangular inequality: the first and third terms are smaller than ϵ by hypothesis while Hölder continuity applies for the second by Proposition 1.

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 λ to L and minimizing C(θ) + λL(θ). This provides a simpler and quite competitive benchmark (see Appendix F.3). We use an orthogonal initialization [START_REF] Saxe | Exact solutions to the nonlinear dynamics of learning in deep linear neural network[END_REF] 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.

Experiments

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. 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 λ 0 = 0.1, τ = 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). 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 + 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 = ReLU(x m + i w m,i (x m )
) and the encoder is no longer fixed. We found that penalizing 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 λ 0 = 1, τ = 0.1 and s = 5) corrects the overfitting of the bigger ResNeXt network compared to ResNet.

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.

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 batchnormalization for up to 50 blocks.

LAP-ResNets are also compared in this regard to the small step method of [Zhang andWynter, 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.

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

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.

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.

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 [START_REF] Teng | The layerwise training convolutional neural networks using local loss for sensor-based human activity recognition[END_REF], 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 [START_REF] Jaderberg | Decoupled neural interfaces using synthetic gradients[END_REF]. Fourthly, it prohibits training the layers in parallel. These limitations follow from the locking problems [START_REF] Jaderberg | Decoupled neural interfaces using synthetic gradients[END_REF]] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF], Pyeon et al., 2021], outperforming more complicated methods addressing the locking problems e.g. synthetic [START_REF] Jaderberg | Decoupled neural interfaces using synthetic gradients[END_REF], Czarnecki et al., 2017] and delayed [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF], Huo et al., 2018a] gradients. We can also deduce theoretical results about networks of greedilytrained shallow modules [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], 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 [START_REF] Arora | Understanding deep neural networks with rectified linear units[END_REF], 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 D, to solve one after the other, for 1≤k≤K, Problems

(T k , F k ) ∈ arg min T,F x∈D L(F, T • G k-1 (x)) (20)
where

G k = T k • ... • T 1 for 1≤k≤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 • G k-1 (x)) = l(F • T • G k-1 (x), y)
where l is a machine learning loss such as cross-entropy or mean squared error. See Figure 16 Indeed, module-wise training suffers often from a well-documented stagnation problem [START_REF] Marquez | Deep cascade learning[END_REF], 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. [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] 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. To solve this issue, InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] 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 [START_REF] Du | Information bottleneck theory based exploration of cascade learning[END_REF]. To avoid early overfitting, we propose a different, transportbased 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[START_REF] He | Deep residual learning for image recognition[END_REF], which remain competitive to this day [START_REF] Wightman | Resnet strikes back: An improved training procedure in timm[END_REF], and similar models such as ResNeXts [Xie et al., 2017], Wide ResNets [START_REF] Zagoruyko | Wide residual networks[END_REF] and vision transformers that are made up essentially of residual connections [START_REF] Liu | Swin transformer: Hierarchical vision transformer using shifted windows[END_REF], 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 [START_REF] Simonyan | Very deep convolutional networks for large-scale image recognition[END_REF]. 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 [START_REF] Nøkland | Training neural networks with local error signals[END_REF] and InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF], 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.

Related Work

Module-wise training methods. Local layer-wise training was initially considered as an early training algorithm for unsupervised learning [START_REF] Hinton | The "wake-sleep" algorithm for unsupervised neural networks[END_REF], as a pre-training and initialization method [START_REF] Bengio | Greedy layer-wise training of deep networks[END_REF], Marquez et al., 2018], and as a complement to end-to-end training to avoid vanishing gradients in the Inception architecture [START_REF] Szegedy | Inception-v4, inception-resnet and the impact of residual connections on learning[END_REF].

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[START_REF] Nøkland | Training neural networks with local error signals[END_REF]. This has led to it being considered in practical settings with limited resources such as embedded training on mobile devices [START_REF] Teng | The layerwise training convolutional neural networks using local loss for sensor-based human activity recognition[END_REF], Tang et al., 2021] and dealing with very large whole slide images in histopathology [START_REF] Zhang | Gigapixel whole-slide images classification using locally supervised learning[END_REF].

Many papers consider using a different auxiliary loss, instead of or in addition to the classification loss: kernel similarity [Mandar [START_REF] Kulkarni | [END_REF], pairwise losses [Duan et al., 2022], class separability [START_REF] Lengellé | Training mlps layer by layer using an objective function for internal representations[END_REF], informationtheory-inspired losses [START_REF] Xu | Training mlps layerby-layer with the information potential[END_REF], Löwe et al., 2019, Nguyen and Choi, 2019, Ma et al., 2020, Wang et al., 2021] and biologically plausible losses [START_REF] Löwe | Putting an end to end-to-end: Gradient-isolated learning of representations[END_REF][START_REF] Nøkland | Training neural networks with local error signals[END_REF], Gupta, 2020[START_REF] Illing | [END_REF], Mostafa et al., 2018]. Paper [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF] reports the best experimental results when solving the layer-wise problems sequentially. Methods PredSim [START_REF] Nøkland | Training neural networks with local error signals[END_REF], DGL [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF], Sedona [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF] and InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] report the best results when solving the module-wise problems in parallel. [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], Belilovsky et al., 2020] do it simply through architectural considerations regarding the auxiliary networks. Sedona [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF] uses architecture search to chose 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, [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF] 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 [START_REF] Huang | Learning deep resnet blocks sequentially using boosting theory[END_REF]] 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 [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], Belilovsky et al., 2020].

In more recent developments, [START_REF] Patel | Local learning with neuron groups[END_REF] attempts to also split the layers width-wise into groups of neurons and [START_REF] Ren | Scaling forward gradient with local losses[END_REF], Fournier et al., 2023] combine layer-wise training with forward gradient learning [START_REF] Baydin | Gradients without backpropagation[END_REF], another alternative to end-to-end backpropagation. Finally, we mention that module-wise training has been extended to the training of spiking neural networks [START_REF] Guo | Efficient training of spiking neural networks with temporallytruncated local backpropagation through time[END_REF] and of neural search networks [START_REF] Chen | Layer-wisely supervised learning for one-shot neural architecture search[END_REF].

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 [START_REF] Xiong | Loco: Local contrastive representation learning[END_REF], 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 [START_REF] Jaderberg | Decoupled neural interfaces using synthetic gradients[END_REF], Czarnecki et al., 2017], DDG [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF], FR [Huo et al., 2018a] and ADL [START_REF] Zhuang | Accumulated decoupled learning with gradient staleness mitigation for convolutional neural networks[END_REF], 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 [START_REF] Sun | Layer-parallel training of residual networks with auxiliary-variable networks[END_REF] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF]. The high dimension of the predicted gradient which scales with the size of the network renders [START_REF] Jaderberg | Decoupled neural interfaces using synthetic gradients[END_REF], 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 [START_REF] Dun | Resist: Layer-wise decomposition of resnets for distributed training[END_REF] and DPA [START_REF] Ni | Deep model assembling[END_REF]. 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, [START_REF] Gao | Deep generative learning via variational gradient flow[END_REF], 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.

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.

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 = id+r), its kinetic energy is simply the squared norm of its residue r = T -id, which we add to the loss L in the target of the greedy problems [START_REF]Boundary tilting perspective of adversarial samples[END_REF]. Given τ >0 used to weight the regularization, we now solve, for 1≤k≤K, Problems

(T τ k , F τ k ) ∈ arg min T,F x∈D L(F, T • G τ k-1 (x)) + 1 2τ ∥T • G τ k-1 (x)-G τ k-1 (x)∥ 2 (21)
where

G τ k =T τ k •..•T τ 1 for 1≤k≤K and G τ 0 =id. The final network is now F τ K •G τ K .
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 [START_REF] Jastrzebski | Residual connections encourage iterative inference[END_REF], 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 ρ, which can be discrete or continuous, and the distribution-wide loss L that arises from the point-wise loss L. Then Problem ( 21) is equivalent in the discrete case to the following problem

(T τ k , F τ k ) ∈ arg min T,F L(F, T ♯ ρ τ k ) + 1 2τ Ω ∥T (x) -x∥ 2 dρ τ k (x) (22)
with

ρ τ k+1 =(T τ k ) ♯ ρ τ k and ρ τ 1 =ρ, i.e data embedding distributions ρ τ k are pushed forward by maps T τ k , and L(F, T ♯ ρ τ k ) = L(F, T (x)) dρ τ k (x) = L(F, .) dT ♯ ρ τ k .
The difference with Problem [START_REF]Module-wise training[END_REF] 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.

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 Z(µ) = min F L(F, µ), which represents the loss of the best classifier on distribution µ. If we restrict ourselves to linear classifiers, Z(ρ τ k ) represents the linear separability of representation ρ τ k at module k of the data distribution ρ. When auxiliary networks are not necessary, for example in generative tasks where the output and the target have the same shape, Z is simply L. When the auxiliary network F is pre-trained and fixed, Z = L(F, µ). The MMS was introduced in [START_REF] Giorgi | Problemi di evoluzione in spazi metrici e curve di massima pendenza[END_REF], 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 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 (Ω) = (P(Ω), W 2 ), where Ω ⊂ R d is a convex compact set, P(Ω) is the set of probability distributions over Ω and W 2 is the Wasserstein distance over P(Ω) derived from the optimal transport problem with Euclidean cost:

W 2 2 (µ, ν) = min T s.t. T ♯ µ=ν Ω ∥T (x) -x∥ 2 2 dµ(x) (23)
where we further assume that ∂Ω is negligible and that the distributions we are dealing with are absolutely continous.

Definition 3. Given Z : W 2 (Ω) → R, the minimizing movement scheme is a discretized gradient flow that is well-defined in non-Euclidean metric spaces and minimizes (under some conditions) Z starting from ρ τ 1 ∈ P(Ω). It is given by

ρ τ k+1 ∈ arg min ρ∈P(Ω) Z(ρ) + 1 2τ W 2 2 (ρ, ρ τ k ) (24)
Equation ( 24) can be seen as a proximal step in the Wasserstein distance for minimizing Z. Equation ( 24) can also be interpreted as a non-Euclidean version of the implicit Euler scheme for following the gradient flow of Z. Indeed, the Euclidean version of ( 24) for F : R d → R is

x τ k+1 ∈ arg min x∈R d F (x) + 1 2τ ∥x -x τ k ∥ 2 (25)
And the optimality condition of problem ( 25) is

x τ k+1 -x τ k τ = -∇F (x τ k+1 ) (26)
which is exactly the formula of the implicit Euler scheme (see Appendix A.1) with step size τ for the Cauchy problem of continuous-time gradient descent:

x ′ (t) = -∇F (x(t)) (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 τ goes to 0 to a gradient flow x of F (i.e. to a solution of x ′ (t) = -∇F (x(t))). Therefore if both sets of assumptions on F are satisfied, x τ k converges to a minimizer of F as t tens to infinity and τ 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 [START_REF]The Monge problem from µ to ν[END_REF].

In the Wasserstein case, if Z is lower-semi continuous then problems [START_REF]The Monge problem from µ to ν[END_REF] always admit a solution because P(Ω) is compact. If Z is also λ-geodesically convex for λ>0, we have convergence of ρ τ k as k→∞ and τ →0 to a minimizer of Z, potentially under more technical conditions (see Appendix C for details). Even though a machine learning loss will usually not satisfy these conditions, this analyses 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 τ as small as possible and many modules T τ k in practice.

Link with the Minimizing Movement Scheme

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

Proposition 3. The distributions ρ τ k+1 = (T τ k ) ♯ ρ τ k
, where the functions T τ k are found by solving (22) and ρ τ 1 = ρ is the data distribution, coincide with the minimizing movement scheme [START_REF]The Monge problem from µ to ν[END_REF] for Z = min F L(F, .).

Proof. The minimizing movement scheme ( 24) is equivalent to taking

ρ τ k+1 = (T τ k ) ♯ ρ τ k where T τ k ∈ arg min T :Ω→Ω Z(T ♯ ρ τ k ) + 1 2τ W 2 2 (T ♯ ρ τ k , ρ τ k ) (28)
under conditions that guarantee the existence of a transport map between ρ τ k and any other measure. Absolute continuity of ρ τ k suffices for this, and the loss can ensure that ρ τ k+1 is absolutely continuous. Among the functions T τ k that solve problem [START_REF]Decoding the internal representations of a network[END_REF], is the optimal transport map from ρ τ k to ρ τ k+1 . To solve specifically for this optimal transport map, we have to solve the equivalent Problem

T τ k ∈ arg min T Z(T ♯ ρ τ k ) + 1 2τ Ω ∥T (x) -x∥ 2 dρ τ k (x) (29)
Problems ( 28) and ( 29) have the same minimum value, but the minimizer of ( 29) is now an optimal transport map between ρ τ k and ρ τ k+1 . 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 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 [START_REF] Arora | Understanding deep neural networks with rectified linear units[END_REF], Bach, 2017, Janzamin et al., 2016, Ge et al., 2018, Du and Goel, 2018].

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 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 F, that ρ τ k is absolutely continuous, that L is continuous and non-negative, that Ω is convex and compact and that ∂Ω is negligible. We then have Theorem 5. Problem (22) has a minimizer (T τ k , F τ k ) such that T τ k is an optimal transport map. For any minimizer (T τ k , F τ k ), T τ k is an optimal transport map.

Proof. Take a minimizing sequence ( Fi , Ti ), i.e. such that C( Fi , Ti ) → min C, where C ≥ 0 is the target function in ( 22) and denote β i = Ti♯ ρ τ k . Then by compacity Fi → F ⋆ and β i ⇀ β ⋆ in duality with C b (Ω) by Banach-Alaoglu. There exists T ⋆ an optimal transport map between ρ τ k and β ⋆ and we have that

C(F ⋆ , T ⋆ ) ≤ lim i→∞ C( Fi , Ti ) = min C by continuity of L and because Ω ∥T ⋆ (x) -x∥ 2 dρ τ k (x) = W 2 2 (ρ τ k , β ⋆ ) = lim i→∞ W 2 2 (ρ τ k , β i ) ≤ lim i→∞ Ω ∥ Ti (x) -x∥ 2 dρ τ k (x)
as W 2 metrizes weak convergence of measures. We take

(F τ k , T τ k )=(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 is η-Hölder continuous almost everywhere and if the optimization algorithm we use to solve the discretized problem ( 21) returns an approximate solution pair ( F τ k , T τ k ) such that T τ k is an ϵ-optimal transport map, i.e. ∥ T τ k -T τ k ∥ ∞ ≤ ϵ, then we have (using the triangle inequality) the following stability property of the module T τ k :

∥ T τ k (x) -T τ k (y)∥ ≤ 2ϵ + C∥x -y∥ η (30)
for almost every x, y ∈ supp(ρ τ k ) 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 and T τ k allows to get stability bounds for the composition networks

G τ k and Gτ k = T τ k • .. • T τ 1 .
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.

Practical implementation

Multi-block Modules

For simplicity, we presented in [START_REF]Transport cost histogram for clean and attacked test samples . . 22 Circles test set and OOD points after blocks 6 and 9[END_REF] 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 [START_REF]31 Internal representations of ResNet-9 on CIFAR10[END_REF] 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 ≤ m < M , we have seen in Section 3 that ∥r m (x m )∥ 2 is its discrete kinetic energy. If ϕ x m denotes the position of a point x after m ResBlocks, then regularizing the kinetic energy of multi-block modules means solving

(T τ k , F τ k ) ∈ arg min T,F x∈D (L(F, T (G τ k-1 (x)) + 1 2τ M -1 m=0 ∥r m (ϕ x m )∥ 2 ) (31) s.t. T = (id + r M -1 ) • ... • (id + r 0 ), ϕ x 0 = G τ k-1 (x) ϕ x m+1 = ϕ x m + r m (ϕ x m )
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 , F τ k ) ∈ arg min T,F L(F, T ♯ ρ τ k ) + 1 2τ 1 0 ∥v t ∥ 2 L 2 ((ϕ • t ) ♯ ρ τ k ) dt (32) s.t. T = ϕ • 1 , ∂ t ϕ x t = v t (ϕ x t ), ϕ • 0 = id where ρ τ k+1 = (T τ k ) ♯ ρ τ
k and r m is the discretization of vector field v t at time t = m/M . Here, distributions ρ τ k are pushed forward through the maps T τ k which correspond to the flow ϕ at time t = 1 of the kinetically-regularized velocity field v t . We recognize in the second term in the target of [START_REF]Internal representations of LAP-ResNet-9 on CIFAR10[END_REF] the optimal transport problem in its dynamic formulation [START_REF] Benamou | A computational fluid mechanics solution to the monge-kantorovich mass transfer problem[END_REF], and given the equivalence between the Monge OT problem (23) and the dynamic OT problem [START_REF]Decomposition of a neural network[END_REF] in Appendix B, Problem [START_REF]Internal representations of LAP-ResNet-9 on CIFAR10[END_REF] 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 ( Fi , ṽi ) and the corresponding maps Ti as in the proof of Theorem 5 to get the result).

Solving the Module-wise Problems

The module-wise problems can be solved in one of two ways. In sequential modulewise 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 and its auxiliary classifier F τ k , meaning that T τ k forwards its output to T τ k+1 immediately after it computes it. This is called decoupled greedy training in [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF], 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 [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] 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 modulewise training, our method improves the test accuracy of all three module-wise training regimes.

Varying the Regularization Weight

The discussion in Section 4.3.3 suggests taking a fixed τ that is as small as possible. However, instead of using a fixed τ , we might want to vary it along the depth k to further constrain with a smaller τ 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 τ 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 λ k,i :=2τ k,i given to the loss. We denote θ k,i the parameters of both T k and F k at SGD iteration i. We also denote L(θ, x) and C(θ, x) respectively the loss and the transport regularization as functions of parameters θ and data x. We now increase the weight λ 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 θ k,1 , initial weights λ k,1 ≥0, learning rates (η i ) and batches (x i ), we apply for module k and i≥1:

θ k,i+1 = θ k,i -η i ∇ θ (λ k,i L(θ k,i , x i ) + C(θ k,i , x i )) λ k,i+1 = λ k,i + hL(θ k,i+1 , x i+1 ) if i mod s = 0 else λ k,i (33)
The weights λ k,i will vary along modules k even if we use the same initial weights λ k,1 = λ 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 λ 1 = λ k,1 for all modules so that this method requires choosing three hyper-parameters (h, s and λ 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 τ for the first K/2 modules and double it (i.e use 2τ ) for the last K/2 modules. A single hyper-parameter has then to be chosen.

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 [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], 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 [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF]. 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 [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] 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 crossentropy 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 [START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF], 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 [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF], 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF], which only focuses on the architecture of the auxiliary classifier, and whose auxiliary classifier we use.

-PredSim [START_REF] Nøkland | Training neural networks with local error signals[END_REF], which adds a similarity matching loss to the loss of every module, requiring a second auxiliary network.

-DDG [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF] 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 [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF].

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

Parallel Module-wise Training with Few Modules

The first experiment is training in parallel 3 residual architectures and a VGG-19 [START_REF] Simonyan | Very deep convolutional networks for large-scale image recognition[END_REF] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF], PredSim [START_REF] Nøkland | Training neural networks with local error signals[END_REF] and Sedona [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF], from Table 2 in [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF]. 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. 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 [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF] 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. The third experiment is training in parallel a ResNet-110 divided into two, four, eight and sixteen modules on STL10 [START_REF] Coate | An analysis of single layer networks in unsupervised feature learning[END_REF]. We compare in Table 6 our results in this setting to the recent methods InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] and DGL [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF] The fourth experiment is training (from scratch) in parallel a Swin-Tiny Transformer [START_REF] Liu | Swin transformer: Hierarchical vision transformer using shifted windows[END_REF] divided into 4 modules on three datasets. We compare in Table 7 our results in this setting with those of InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] and InfoProL, a variant of InfoPro proposed in [START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF], from Table 3 in [START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF]. 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 [START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF]. 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. 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 τ over a large scale of values.

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 [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] 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 [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF]] 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≤4 modules. When the modules are even in weight and not in depth, keeping the notations of [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF], we call the training methods VanGL*, TRGL* and InfoPro*. As said above, this leads to shallower early modules which hurts performance according to [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF]. 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. [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF] 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 [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF] is architecture-wise essentially identical to VanGL and consumes the same amount of memory. [START_REF] Teng | The layerwise training convolutional neural networks using local loss for sensor-based human activity recognition[END_REF], Tang et al., 2021]. We test our regularization therefore in this section on sequential block-wise training.

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 [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF].

We check in 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 [START_REF] Marquez | Deep cascade learning[END_REF]. 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 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 and26 in Appendix G.2). The 88% accuracy of sequential training on CIFAR10 in Table 11 is the same in table 2 of [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], which is the best method for layerwise sequential training available, with VGG [START_REF] Simonyan | Very deep convolutional networks for large-scale image recognition[END_REF] 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 [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF]).

Avoiding Early Overfitting

We verify that our method avoids the stagnation in test accuracy with increasing depth. In Figure 18

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 ̸ = 2, i.e. regularizing with a norm ∥.∥ p with p ̸ = 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.

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 [START_REF] Karkar | Adversarial sample detection through neural network transport dynamics[END_REF]] at ECML 2023.

Introduction

Neural networks have improved performances on many learning tasks, including image classification. They are however vulnerable to adversarial attacks [START_REF] Szegedy | Intriguing properties of neural networks[END_REF]. 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 word, for example causing autonomous cars to misclassify road signs [START_REF] Eykholt | Robust physicalworld attacks on deep learning visual classification[END_REF], and can be generated without having access to the network [START_REF] Liu | Delving into transferable adversarial examples and black-box attacks[END_REF]. 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 word [START_REF] Amodei | Concrete problems in ai safety[END_REF].

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 [START_REF] Nguyen | Deep neural networks are easily fooled: High confidence predictions for unrecognizable images[END_REF]. 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 [START_REF] He | Deep residual learning for image recognition[END_REF], 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.

[ [START_REF] Wu | Skip connections matter: On the transferability of adversarial examples generated with resnets[END_REF] 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 [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF] and natural scene statistics [START_REF] Kherchouche | Detection of adversarial examples in deep neural networks with natural scene statistics[END_REF] 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 [START_REF] Pang | Towards robust detection of adversarial examples[END_REF].

Related Work

An adversarial attack aims at slightly modifying an input image at test time by no more than ϵ 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 ϵ > 0, an adversarial sample y constructed from a clean sample x is y = x + δ, such that ∥δ∥ ≤ ϵ and f (y) ̸ = f (x). The maximal perturbation size ϵ 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. White-box attacks are attacks where the adversary has access to the network's architecture, weights and gradients. The Fast Gradient Method (FGM) [START_REF] Goodfellow | Explaining and harnessing adversarial examples[END_REF]] takes a perturbation step in the direction of the gradient that maximizes the loss (see Figure 19 above). Projected Gradient Descent (PGD) [START_REF] Madry | Towards deep learning models resistant to adversarial attacks[END_REF] and the Basic Iterative Method (BIM) [START_REF] Kurakin | Adversarial examples in the physical world[END_REF] are iterative versions of FGM. Auto-PGD (APGD) [Croce and [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF]] is a variant of PGD that uses an adaptive step size. Two slower but more effective attacks are DeepFool (DF) [START_REF] Moosavi-Dezfooli | Deepfool: A simple and accurate method to fool deep neural networks[END_REF], which iteratively perturbs an input in the direction of the closest decision boundary, and Carlini-Wagner (CW) [START_REF] Carlini | Towards evaluating the robustness of neural networks[END_REF], which solves an optimization problem to find the perturbation.

Black-box attacks don't have any knowledge about the the network and can only query it. They include the Boundary Attack (BA) [START_REF] Brendel | Decisionbased adversarial attacks: Reliable attacks against black-box machine learning models[END_REF], 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 [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF] 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) [START_REF] Andriushchenko | Square attack: a query-efficient black-box adversarial attack via random search[END_REF].

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 [START_REF] Szegedy | Intriguing properties of neural networks[END_REF][START_REF] Goodfellow | Explaining and harnessing adversarial examples[END_REF]), regularization [START_REF] Gu | Towards deep neural network architectures robust to adversarial examples[END_REF], distillation [START_REF] Papernot | Distillation as a defense to adversarial perturbations against deep neural networks[END_REF] and denoising [START_REF] Meng | Magnet: A two-pronged defense against adversarial examples[END_REF]]. 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 59 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 [START_REF] Metzen | On detecting adversarial perturbations[END_REF]. However, this network can be adversarially attacked as easily as the first network. More recent popular statistical approaches include LID [START_REF] Ma | Characterizing adversarial subspaces using local intrinsic dimensionality[END_REF], Mahalanobis [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF], Ren et al., 2021], MMD [START_REF] Grosse | On the (statistical) detection of adversarial examples[END_REF] and KD [START_REF] Feinman | Detecting adversarial samples from artifacts[END_REF]. 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, [START_REF] Gretton | A kernel two-sample test[END_REF]) between shuffled inputs. KD combines kernel density (KD, [Parzen, 1962, Rosenblatt, 1956]) estimates and Bayesian uncertainty estimates obtained via dropout [START_REF] Srivastava | Dropout: A simple way to prevent neural networks from overfitting[END_REF] 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. [START_REF] Carrara | Adversarial examples detection in features distance spaces[END_REF] 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 [START_REF] Harder | Spectraldefense: Detecting adversarial attacks on cnns in the fourier domain[END_REF] or on natural scene statistics [START_REF] Kherchouche | Detection of adversarial examples in deep neural networks with natural scene statistics[END_REF], Kherchouche et al., 2022] have been proposed. See Appendix E.2 for details and [START_REF] Aldahdooh | Adversarial example detection for dnn models: a review and experimental comparison[END_REF] for a review.

60

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 [START_REF] Weng | Evaluating the robustness of neural networks: An extreme value theory approach[END_REF], 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 [START_REF] Hein | [END_REF]Andriushchenko, 2017, Cisse et al., 2017]. These regularizations [Cisse et al., 2017, Gu and[START_REF] Gu | [END_REF] 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 [START_REF] Gilmer | Adversarial spheres: The relationship between high-dimensional geometry and adversarial examples[END_REF], 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 [START_REF] Tanay | A boundary tilting perspective on the phenomenon of adversarial examples[END_REF], 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). To the best of our knowledge, [START_REF] Pang | Towards robust detection of adversarial examples[END_REF] is the only other method that 61 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 crossentropy 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.

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 [START_REF]Accuracy of each module after module-wise training[END_REF] to improve detectability of adversarial attacks.

Detection

Given a residual network that applies x m+1 = x m + hr m (x m ) to an input x 0 for 0 ≤ m<M , we consider the embeddings x m for 0<m ≤ M , or the residues r m (x m ) for 0 ≤ 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 1 d m ∥r m (x m )∥ 2 , cos r m (x m ), 1 m M -1 m=0 ( 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 1 dm 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 62 does not fully satisfy the dynamic view, as the activation is applied after the skip-connection, i.e. x m+1 = 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 ( [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF]) and LID ( [START_REF] Ma | Characterizing adversarial subspaces using local intrinsic dimensionality[END_REF]) 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.

Regularization

Regularity of neural networks (typically Lipschitz continuity) has been used as a measure of their robustness to adversarial attacks [START_REF] Weng | Evaluating the robustness of neural networks: An extreme value theory approach[END_REF], 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)∥ ≤ 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 [START_REF] Hein | Formal guarantees on the robustness of a classifier against adversarial manipulation[END_REF] and [START_REF] Cisse | Parseval networks: Improving robustness to adversarial examples[END_REF]. 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 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 andWynter, 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 η-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

min θ C(θ) = x∈D M -1 m=0 ∥r m (ϕ x m )∥ p subject to ϕ x m+1 = ϕ x m + hr m (ϕ x m ) and ϕ x 0 = x for all x ∈ D L(θ) = 0 (35) 63 which is the discretization of inf v,F 1 0 ∥v t ∥ p L p ((ϕ • t ) ♯ µ) dt subject to ∂ t ϕ x t = v t (ϕ x t ) ϕ • 0 = id (ϕ • 1 ) ♯ µ ∈ S F,L (36)
which is equivalent to inf

T,F C(T ) = Ω c(x, T (x)) dµ(x) subject to T ♯ µ ∈ S F,L (37)
We solve again [START_REF]Weight distribution of the classifier of a vanilla ResNet[END_REF] through the method of multipliers

θ i+1 = arg min θ C(θ) + λ i L(θ) λ i+1 = λ i + τ L(θ i+1 ) (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 [START_REF]Early overfitting of vanilla module-wise training[END_REF].

We take Ω ⊂ R d convex and compact and the data distribution α ∈ P(Ω) absolutely continuous and such that δΩ is α-negligible. We suppose that there exists an open bounded convex set X ⊂ Ω such that α is bounded away from zero and infinity on X and is zero on its complement X ∁ . 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 α and β := T ♯ α. We suppose that β is absolutely continuous and that there exists an open bounded convex set Y ⊂ Ω such that β is bounded away from zero and infinity on Y and is zero on Y ∁ . In the rest of this section, v solves [START_REF]Weight distribution of the classifier of a LAP-ResNet[END_REF] and we suppose that we find a solution to the discretized problem [START_REF]Weight distribution of the classifier of a vanilla ResNet[END_REF] that is an ε/2-approximation of v, i.e. ∥r m -v tm ∥ ∞ ≤ ε/2 for all 0 ≤ m < M , with t m = m/M .

In Theorem 6, we show that the regularization makes the residual blocks of the network η-Hölder (with an error of ε) 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 R d .

Theorem 6. For a, b ∈ support(α tm ) where α t := (ϕ • t ) ♯ α with ϕ solving (36) along with v, we have

∥r m (a) -r m (b)∥ ≤ε + K∥a -b∥ ζ1 if ∥a -b∥ ≤ 1 ∥r m (a) -r m (b)∥ ≤ε + K∥a -b∥ ζ2 if ∥a -b∥ > 1 for constants K > 0 and 0 < ζ 1 ≤ ζ 2 ≤ 1,

64

Proof. A solution v to [START_REF]Weight distribution of the classifier of a LAP-ResNet[END_REF] exists and is linked to an optimal transport map T that is a solution to [START_REF]Distances between successive point clouds in a ResNet-9[END_REF]

through v t = (T -id) • T -1 t
where T t := (1 -t)id + tT which is invertible (see Appendix B). By Theorem 1 in Appendix B, being an optimal transport map, T is η-Hölder on X. So for all t ∈ [0, 1[ and all a, b ∈ support(α t ), where α t = (ϕ • t ) ♯ α = (T t ) ♯ α with ϕ solving (36) together with v, we have

∥v t (a) -v t (b)∥ ≤ ∥T -1 t (a) -T -1 t (b)∥ + C∥T -1 t (a) -T -1 t (b)∥ η (39)
Since (α t ) 1 t=0 is a geodesic between α and β = α 1 = T ♯ α, then (α s ) t s=0 is a geodesic between α and α t (modulo reparameterization to [0, 1]). And since α s = (T s ) ♯ α, the map T t is an optimal transport map between α and α t . Therefore its inverse T -1 t is an optimal transport map (see Theorem 25 in Appendix B) and is η t -Hölder with 0<η t ≤1 (being a push-forward by T t , the support of α t satisfies the conditions of Theorem 1 in Appendix B). Therefore, for all a, b ∈ support(α t )

∥v t (a) -v t (b)∥ ≤ C t ∥a -b∥ ηt + CC η t ∥a -b∥ ηηt (40)
and for all a, b ∈ support(α tm )

∥r m (a) -r m (b)∥ ≤ ε + C tm ∥a -b∥ ηt m + CC η tm ∥a -b∥ ηηt m (41)
by the hypothesis on r and the triangle inequality. Let K := max m C tm + CC η tm , ζ 1 := η min m η tm and ζ 2 := max m η tm . Then, we have the desired result immediately from [START_REF]Circles test set embeddings with large initialization, BN and LAP 42 Accuracy of each block after block-wise training[END_REF].

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 [START_REF] Huster | Limitations of the lipschitz constant as a defense against adversarial examples[END_REF], or through more complex estimation algorithms [START_REF] Virmaux | Lipschitz regularity of deep neural networks: analysis and efficient estimation[END_REF], Latorre et al., 2020, Bhowmick et al., 2021]. Bound ( 39) allows through T -1 t 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 ∈ X and constants C, L>0,

∥r m (a m ) -r m (b m )∥ ≤ ε + ∥a 0 -b 0 ∥ + C∥a 0 -b 0 ∥ η + + L(∥a m -ϕ a0 tm ∥ + ∥b m -ϕ b0 tm ∥)
The term µ(a 0 ):=∥a m -ϕ a0 tm ∥ (and µ(b 0 ):=∥b m -ϕ b0 tm ∥) is the distance between the point a m after m residual blocks and the point ϕ a0 tm we get by following the 65 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 [START_REF]Circles test set embeddings with large initialization[END_REF]. Bounding the terms µ(a 0 ) and µ(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 [START_REF]Circles test set embeddings with large initialization[END_REF]. 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 ∈R d and constants R, S > 0,

∥r m (a m ) -r m (b m )∥ ≤ ε + LSε + LSRh + ∥a 0 -b 0 ∥ + C∥a 0 -b 0 ∥ η + + LS(dist(a 0 , X) + dist(b 0 , X))
Proof. Since T -1 t (ϕ x t ) = x, we have for any a 0 , b 0 ∈ X by the triangle inequality

∥r m (a m ) -r m (b m )∥ ≤ ∥r m (a m ) -r m (ϕ a0 tm )∥ + ∥r m (ϕ a0 tm ) -v tm (ϕ a0 tm )∥+ + ∥v tm (ϕ a0 tm ) -v tm (ϕ b0 tm )∥ + ∥r m (ϕ b0 tm ) -v tm (ϕ b0 tm )∥+ + ∥r m (b m ) -r m (ϕ b0 tm )∥ ≤ ε + ∥a 0 -b 0 ∥ + C∥a 0 -b 0 ∥ η + + L(∥a m -ϕ a0 tm ∥ + ∥b m -ϕ b0 tm ∥)
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 -ϕ a0 tm ∥ (and likewise ∥b m -ϕ b0 tm ∥) 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 ϕ a0 tm 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, ∥ϕ a0 tm -a m ∥ ≤ ∥ϕ a0 tm -ãm ∥ + ∥ã m -a m ∥ ≤ Sε + SRh where ãm comes from the Euler scheme with access to v (ã m+1 :=ã m +hv tm (ã m ) and ã0 :=a 0 ), R is the consistency constant of the Euler method and S is its zero-stability constant. If a 0 , b 0 / ∈ X, we need to introduce â0 := Proj X (a o ) and b0 := Proj X (b o ) to apply [START_REF]Circles test set embeddings with large initialization[END_REF]. We now get

∥r m (a m ) -r m (b m )∥ ≤ ε + ∥a 0 -b 0 ∥ + C∥a 0 -b 0 ∥ η + L(∥a m -ϕ â0 tm ∥ + ∥b m -ϕ b0 tm ∥)
Bounding the terms ∥a m -ϕ â0 tm ∥ and ∥b m -ϕ b0 tm ∥ now gives

∥ϕ â0 tm -a m ∥ ≤ ∥a m -ãm ∥ + ∥ã m -ϕ â0 tm ∥ ≤ S(∥a 0 -â0 ∥ + ε) + SRh
where ãm now comes from the Euler scheme with access to v that starts at â0 (meaning ãm+1 :=ã m +hv tm (ã m ) and ã0 :=â 0 ). Likewise, we get the same bound for ∥b m -ϕ b0 tm ∥. Since ∥a 0 -â0 ∥ = dist(a 0 , X) and ∥b 0 -b0 ∥ = 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 -ϕ â0 tm ∥ we get the same result. Indeed, if ãm again comes from the Euler scheme with access to v that starts at a 0 (meaning ãm+1 :=ã m +hv tm (ã m ) and ã0 :=a 0 ), we can write, for some constant F > 0

∥ϕ â0 tm -a m ∥ ≤ ∥a m -ãm ∥ + ∥ã m -ϕ a0 tm ∥ + ∥ϕ a0 tm -ϕ â0 tm ∥ ≤ Sε + SRh+ + F ∥a 0 -â0 ∥ since ∥ϕ a0 tm -ϕ â0 tm ∥ ≤ ∥a 0 -â0 ∥ + tm 0 ∥v s (ϕ a0 s ) -v s (ϕ â0 s )∥ ds ≤ ∥a 0 -â0 ∥ + M tm 0 ∥ϕ a0 s -ϕ â0 s ∥ ds ≤ F ∥a 0 -â0 ∥
where we get the last line by Gronwall's lemma.

The terms dist(a 0 , X) and 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.

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 [START_REF] Goodfellow | Explaining and harnessing adversarial examples[END_REF], PGD [START_REF] Madry | Towards deep learning models resistant to adversarial attacks[END_REF] (specifically the Auto-PGD-CE (APGD) variant [Croce and Hein, 2020b]), BIM [START_REF] Kurakin | Adversarial examples in the physical world[END_REF], DF [START_REF] Moosavi-Dezfooli | Deepfool: A simple and accurate method to fool deep neural networks[END_REF], CW [START_REF] Carlini | Towards evaluating the robustness of neural networks[END_REF], and AA [Croce and [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF], and 2 black-box attacks that can only query the network: HSJ [Chen et al., 2020a] and BA [START_REF] Brendel | Decisionbased adversarial attacks: Reliable attacks against black-box machine learning models[END_REF]. 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 ϵ = 0.03. We use the L 2 norm for CW and HSJ and the L ∞ norm for all the other attacks. We use the Adversarial Robustness Toolbox (ART) Python package [START_REF] Nicolae | Adversarial robustness toolbox v1[END_REF] to generate the adversarial samples, except for AA for which we use the authors' original code [Croce and [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF]. We use the default values from the ART package and the code of [Croce and [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF] 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 [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF] and to the detector of [START_REF] Kherchouche | Detection of adversarial examples in deep neural networks with natural scene statistics[END_REF], 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 [START_REF] Pang | Towards robust detection of adversarial examples[END_REF], 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 τ =1, s=1 and λ 0 =1 for all networks. These hyper-parameters are chosen to improve validation accuracy during training and not adversarial detection. 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.

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 ϵ almost perfectly on the LAP-ResNeXt50 trained on CIFAR100, but not on the vanilla ResNeXt50. We visualize this increased difference in behavior between clean and out-ofdistribution (OOD) points with LAP training in 2 dimensions in Figure 22. The experiments is binary classification of the circles dataset from scikit-learn [START_REF] Pedregosa | Scikit-learn: Machine learning in Python[END_REF], where we add OOD points at test time.

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.

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 [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF]], 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 [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF], but did not find a consistent improvement in detection, even only on successfully attacked images. 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 [START_REF] Pedregosa | Scikit-learn: Machine learning in Python[END_REF]) 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,[START_REF]Internal representations of ResNet-9 on CIFAR100[END_REF] 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 and33

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,[START_REF]Weight distribution of the classifier of a LAP-ResNet[END_REF] and 37 in Appendix H.3. On our detector, the detection rate of successful adversarial samples from unseen attacks remains higher than 90% in most cases (Table 39 in Appendix H.4) and the FPR is always lower than 10% (Table 41 in Appendix H.5). The AUROC is in Table 43 in Appendix H.6. Our detector almost always outperforms the other detectors on all these metrics. On the attacks that are not variants of FGM, the NS detector almost never predicts that an input is adversarial.

However, this experiment shows that our approach has some limitations. We see in Tables 35,[START_REF]Weight distribution of the classifier of a LAP-ResNet[END_REF] and 37 in Appendix H.3 that LAP training does not improve detection accuracy as much anymore in this setting, and reduces it in some cases. It improves it for the Mahalanobis detector on all attacks on the ResNet110-CIFAR10 and WideResNet-TinyImageNet experiments (by up to 10 percentage points), but not on most attacks on the ResNeXt50-CIFAR100 experiment. It often reduces detection generalization on our detector. But it always does better than RCE training in Tables 35 and36 in Appendix H.3.

We claim this is because these methods reduce the variance of features extracted during training on the seen attack, thus harming generalization to unseen attacks. This explains why detection of PGD and BIM, which are variants of FGM, still often improves. Fixing this could be an area of future research. That LAP training improves detection of unseen OOD samples in some cases in the out-ofdistribution experiment in Tables 14 and15 in Section 5.4.4 indicates that this problem is not general.

Detection of Out-of-Distribution Samples

Since our analysis applies to all out-of-distribution (OOD) samples, we test detection of OOD samples in a similar setting to the Mahalanobis paper [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF]. We use the same ResNet110 and ResNeXt50 models trained on CIFAR10 and CIFAR100 respectively. Since the detectors need to be trained, we are in the OOD setting where we have a first dataset for training the network (CIFAR10 in Table 14 and CIFAR100 in Table 15) and a second dataset from another distribution that is not the test OOD distribution to train the detector on. This could be another dataset (CIFAR100 in the first column of Table 14), some images found in the wild, or a perturbation of our dataset that we generate using an adversarial attack (CW on CIFAR10 in the second column of Table 14, and AA and CW on CIFAR100 in Table 15). Detectors can then be used by training them to distinguish between these first two datasets, and then testing them on distinguishing between the first dataset and a third unseen dataset (SVHN [START_REF] Netzer | Reading digits in natural images with unsupervised feature learning[END_REF] in both tables). Results are in Tables 14 and15. Our detector performs very well and better than the MH detector in three of the four experiments, and in the fourth case, the MH detector benefits from LAP training by 8 percentage points (first experiment in Table 15).Without any extra data available, using the CW adversarial attack allows to detect OOD samples from an unseen distribution with more than 90% accuracy and an FPR of less than 10% at a fixed TPR of 95%. The choice of the attack is also important, as CW allows for much better detection of unseen samples from SVHN than AA.

The AUROC and the false positive rate at a fixed true positive rate of 95% are in Appendix H.7, in Tables 44 and45 for ResNet100, and in Tables 46 and47 for ResNeXt50. They show the same advantage for our detector and regularization over the Mahalanobis detector and RCE training respectively.

Attacking the Detector

We consider here the case where the attacker also attacks the detector (adaptive attacks). We try two such attacks on the TR and MH detectors on ResNet110 trained on CIFAR10. Both attacks are white-box with respect to the network. The first is black-box with respect to the detector. It only knows if an adversarial sample has been detected or not. The second has some knowledge about the detector. It knows what features it uses and can attack it directly to find adversarial features. We test these attacks by looking at the percentage of detected successful adversarial samples that they turn into undetected successful adversarial samples that fool both the network and the detector.

The first attack proceeds as follows. A strong white-box attack (CW) is used on the network on image x that has label y. If it finds a successful adversarial image x that fools the network into predicting ỹ ̸ = y but is detected by the detector, the attacker will attempt to modify this image x so that the network and the detector are both fooled. For this, the image x is used as the initialization for an attack (HSJ with a budget of 50 iterations and 10000 evaluations) on a black-box Network-Detector system. The attacker considers that the Network-Detector behaves as follows: it outputs the class prediction of the network if the detector does not detect an attack and outputs an additional 'detected' class if the detector detects an attack. The attacker attacks this Network-Detector on image x targeting the ỹ label. This way the network makes a mistake and the 'detected' class is avoided. On the vanilla ResNet110, this attack turns 16.5% of 1700 detected successful adversarial samples x into undetected successful adversarial samples on our detector, compared to 25.7% on the Mahalanobis detector. These percentages are lower on the LAP-ResNet110 as they drop to 6.8% on our detector and 12.9% on the Mahalanobis detector. This shows that LAP training improves the robustness of both adversarial detectors to being attacked themselves, and that the Transport detector is more robust than the MH detector.

The second attack is very similar to the adaptive attack used in [Carlini and Wagner, 2017a] to break the Kernel Density detector of [START_REF] Feinman | Detecting adversarial samples from artifacts[END_REF]. It proceeds as follows. A strong white-box attack (CW) is used on the network on image x that has label y. If it finds a successful adversarial image x that fools the network but is detected by the detector, the detection features z that x generates when run through the network are used as the initialization for a black-box attack (HSJ with a budget of 50 iterations and 10000 evaluations) on the detector. If successful adversarial detection features z * that fool the detector are found, the attacker has to find an adversarial perturbation of x that still fools the network and that generates these features z * (or close features that also fool the detector) when run through the network. We do this as in [Carlini and Wagner, 2017a] by solving the following optimization problem:

min x * -L(N (x * ), y) + c 1 ∥D(x * ) -z * ∥ + c 2 ∥x * -x∥ ( 42
)
where L is the cross-entropy loss, N is the network, and D is the (differentiable) function that returns the detection features of its input. This optimization problem is differentiable and we try differentiable optimization algorithms such as BFGS and NR to solve it. The initial detected successful adversarial image x is used as initialization as in [Carlini and Wagner, 2017a]. This attack turns 14% of detected successful adversarial samples x into undetected successful adversarial samples on our detector on the LAP-ResNet110.

Given that initial detection rates of successful adversarial samples are very often close to 100% (see Tables 38 and39 in Appendix H.4), this shows that adaptive attacks do not (at least not easily) circumvent the detector, as detection rates drop to 85% at worst. The second attack is stronger than the first one, but it can probably still be improved by using a white-box attack that is specific to random forests for attacking the detector such as [START_REF] Kantchelian | Evasion and hardening of tree ensemble classifiers[END_REF] or [START_REF] Zhang | An efficient adversarial attack for tree ensembles[END_REF], or a different loss than cross-entropy such as the one used in the CW attack. The choice of the differentiable term representing the detector (here ∥D(x * ) -z * ∥ in ( 42)) is also important. However, the difficulty of combining the attack on the network with that on the detector remains. It is the non-differentiability of the random forest that forces either this separate treatment of network and detector then the use of a proxy differentiable term for the detector in (42) to combine both, or the use of a black-box method as in the first attack. Also, we did not consider here the ensemble of the class-conditional detector and the general detector, which is the best performing version of the detector (see Section 5.4.2), and should be more robust to adaptive attacks, as the attacker will have to fool two random forest detectors at once and target a particular label, constraining further the optimization problem he solves. This is similar to [START_REF] Eniser | Raid: Randomized adversarial-input detection for neural networks[END_REF] where a random detector is selected at test time from a pool of trained detectors to render adaptive attacks on the detector more difficult.

Time Comparison

With a ResNeXt50 on CIFAR100 and a Tesla V100 GPU, it takes our method (including the time to generate FGM attacks) 66 seconds to extract its features from both the clean and the adversarial samples, while it takes the Mahalanobis method 110 seconds. Mahalanobis also extracts some statistics from the training set prior to adversarial training, which takes an additional 35 seconds. Our feature vector is of size 32, compared to 5 for the Mahalanobis detector. So our random forest takes only 4 more seconds to train than the Mahalanobis one (7 vs 3 seconds). Computation of the features our detector uses (norms and cosines) is in O(M D), where M is the number of residual blocks and D is the largest embedding dimension inside the network.

Discussion and Conclusion

To conclude, we proposed a method for detecting abnormal samples, based on the dynamical viewpoint of neural networks. The method examines the discrete vector field moving the inputs to distinguish clean and adversarial samples. The detector requires minimal processing and computation to extract the features it needs for detection and achieves state-of-the-art detection performances on 3 different architectures and datasets, on seen and unseen attacks, and also on out-of-distribution detection.

The low dimension of the features we extract (2 per residual block) allows to keep the features from all blocks as opposed to other detection methods that only experimentally select which layers to use. This could explain the superior performance of our detector, as we closely monitor the evolution of the input's embeddings inside the network.

We have also used the transport regularization that was previously shown to improve test accuracy to improve the detection accuracy of adversarial detectors. This regularization provably makes the embeddings closer to each other on the support of the input distribution, thus making them more distinguishable from those of abnormal samples. The main difference with other regularizations that make the network more Lipschitz by regularizing its weights directly is that we only endue the network with additional regularity on the support of the data distribution and not on the entire space.

A possible simple extension of our detection method would be to add cosines to other orthogonal vectors as detection features. A trade-off between the accuracy of the detector and the size and computation time of the features would then have to be taken into account.

Conclusion

Summary

Residual connections are among the workhorses that make neural networks successful. We have expanded on interpretations of their functioning, namely, the iterative refinement bias (i.e. the bias to remain close to the identity function) and the discrete dynamical system interpretation. We use optimal transport theory to theoretically analyze their functioning. In optimal transport terms, these two interpretations are unified to show that residual connections induce a desirable bias towards networks with low kinetic energy and transport cost. This analysis allowed us to propose improvements to the accuracy and stability of residual networks and applications to module-wise training that make it competitive with end-to-end training with only a fraction of the memory requirements and to detection of adversarial and out-of-distribution samples.

In a first work [START_REF] Karkar | A principle of least action for the training of neural networks[END_REF] detailed in Section 3, we empirically motivated and formulated the following learning problem: among all the networks that solve the task, find the one that transforms the data with the lowest cost. This amounts to regularizing the transport cost of the network. We proved existence and regularity of models trained this way, studied their behavior in low-dimensional settings and shown their efficiency on standard classification tasks. Training is stabilized in an adaptive fashion without being slowed down.

In a second work [Karkar et al., 2023a] detailed in Section 4, we used this regularization for module-wise training of neural networks without end-to-end backpropagation, including on Transformers, saving up to 60% memory usage. We proved that this makes every module apply a proximal step in the Wasserstein space to maximize the separability of the data.

In a third work [START_REF] Karkar | Adversarial sample detection through neural network transport dynamics[END_REF] detailed in Section 5, we proposed a detector of adversarial and out-of-distribution samples based on the dynamic viewpoint of neural networks. We also used our detector successfully on detection of unseen out-of-distribution samples. The transport regularization provably makes the activations of out-of-distribution samples more distinguishable from those of clean samples, thus improving the accuracy of adversarial detectors.

Future Directions

Besides the direct possible extensions mentioned in the conclusions to Sections 3 to 5, other avenues of research follow from our analysis and point of view.

Adaptive Depth During Training

Introduction

The dynamical point of view of neural networks, initiated by the analogy between residual networks and the Euler scheme, has led to new architectures [START_REF] Lu | Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations[END_REF], Haber et al., 2019]. While these new architectures often imitate numerical solvers of higher order than the Euler scheme through skip connections, adaptive solvers have only yet been explored in a few papers. We propose to use them to adjust the step size and the depth of the network during training, possibly reducing training cost. Given that adaptive methods are made up of multiple imbedded schemes, this also allows to ensemble different pre-classification forward passes at inference, without an additional training cost.

Method

This is currently a work in progress with positive results on the CIFAR100 dataset. We build a neural network that imitates the structure of an adaptive Runge-Kutta scheme (see Appendices A.2 and A.3). The block we use contains 4 convolutional functions denoted f i and does the following computations for an input x with a step size h:

                             p 1 = f 1 (x) y euler = x + hp 1 p 2 = f 2 (x + 1 2 hp 1 ) y midpoint = x + hp 2 p 3 = f 3 (x + 3 4 hp 2 ) y order3 = x + 2 9 hp 1 + 1 3 hp 2 + 4 9 hp 3 p 4 = f 4 (y order3 ) y order2 = x + 7 24 hp 1 + 1 4 hp n,2 + 1 3 hp 3 + 1 8 hp 4 (43)
Depending on the numerical scheme we want to use, the block forwards one of the outputs y to the next block. Only layers f 1 and f 2 will be trained for example if we want to use the midpoint scheme. This allows to compare the efficiency of the different schemes for solving classification tasks, i.e. is it better to train 1 block of order 3 or 3 blocks of order 1 (since the order 3 method contains 3 functions f i )?

If all the functions f i are learned, and since all the values y approximate the same differential equation, it makes sense to average them after the last block before the classification layer as an ensembling method.

The depth is adaptively adjusted during training as follows. First, a tolerance hyperparameter ε > 0 is fixed. Then, on the validation set, every S epochs, the average error err = ∥y order3 -y order2 ∥ 2 for each block is computed. If it is larger than ε then the step size is halved and the block is duplicated, i.e. a block with the same weights is added right after it. If the error is smaller than ε/3, then the step size is doubled and the following block is deleted.

To summarize, the contributions we aim to verify on larger datasets are:

-Using neural blocks of higher order is more efficient than using normal Euler residual blocks, meaning that even though the blocks are larger, fewer are required to reach higher accuracy with less memory usage.

-Since adaptive scheme are made-up of interwoven schemes of different order, different forward passes representing the schemes of different order can be ensembled at test time, through averaging the final embeddings before classification, to improve accuracy.

-Adaptively adjusting the step-sizes and depth of the network allows to reduce, since the first training epochs, the size of a large network, while still obtaining the same performance at the end of training.

-This architecture also benefits from LAP training as in Section 3.

Related Work

[Yang et al., 2020a] introduce a time-step controller that is an LSTM [START_REF] Hochreiter | Long short-term memory[END_REF] which is dependent on the convolution weights of the current and previous layers and that is jointly optimized with the network. In contrast, our time-step is adjusted as in adaptive numerical schemes, therefore requiring no extra training, and also adjusts the depth of the network during training. [START_REF]Multi-level residual networks from dynamical systems view[END_REF] use the analogy between ResNets and numerical schemes for differential equations to accelerate training by starting the training with a shallow network with a large step-size and then doubling the depth and halving the step-size during training. This is not an adaptive method to adjust the depth and step-size, as it is not dependent on the stability of the approximated differential equation. [START_REF] Zhu | Convolutional neural networks combined with runge-kutta methods[END_REF] simply implement an architecture that imitates a Runge-Kutta scheme. However, the coefficients are not fixed to the theoretical values of the scheme and are instead learned inside the convolutions. The relations between them necessary to the theoretical validity of the scheme are therefore not enforced. There is also no modification of the step-size and depth. Likewise, [START_REF] Lu | Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations[END_REF] use an architecture that imitates the linear multi-step numerical method and [START_REF] Haber | Imexnet -a forward stable deep neural network[END_REF] imitate the implicit-explicit method. Modifying the step-size has shown its benefits in [Zhang et al., 2019, Zhang and[START_REF] Zhang | [END_REF], where fixing it to a small value around 0.1 makes training more stable.

Somewhat similar to our proposed method is the adaptive computation time method applied to residual networks [START_REF] Figurnov | Spatially adaptive computation time for residual networks[END_REF], inspired by a similar idea applied to recurrent neural networks [Graves, 2017, Fojo et al., 2018, Jernite et al., 2017]. In these methods, secondary neural networks are trained to choose (minimize) the number of times the state transition is applied for an input step in a recurrent network, and the number of residual blocks applied for an input in a residual network. The goal in [START_REF] Figurnov | Spatially adaptive computation time for residual networks[END_REF] is faster computation time and not the adjustment during training of the size of the network to the dataset at hand, as the number of residual blocks applied is chosen for each input (or even patch of the input), so the entire network still has to be trained.

Similarity of Neural Network Representations

We propose to use the study of the dynamics inside a neural network, especially as used in Section 5 to detect adversarial and out-of-distribution samples, to consider the problem of measuring similarity between the representations learned by neural networks. This area of research [START_REF] Raghu | Svcca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability[END_REF], Kornblith et al., 2019, Laakso and Cottrell, 2000, Li et al., 2015, Morcos et al., 2018, Wang et al., 2018] tries to answer questions such as:

-Do different architectures learn similar representations when trained on the same dataset?

-Does the same architecture tend to learn similar representations when trained on similar datasets ?

-How much does changing the initialization change the learned representations?

-Can we establish correspondences between layers of different networks?

The available methods usually propose similarity measures (e.g. SVCCA [START_REF] Raghu | Svcca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability[END_REF] and CKA [START_REF] Kornblith | Similarity of neural network representations revisited[END_REF]) between activations that have certain desirable invariances such as invariance to orthogonal transformations.

Optimal transport theory was used in [START_REF] Singh | Model fusion via optimal transport[END_REF] to establish correspondences between the neurons of two networks, which is then useful for fusing these networks into a network that keeps as much of the capabilities of each network as possible. However, [START_REF] Singh | Model fusion via optimal transport[END_REF] do this only layer by layer and not across layers, which could be an extension of their method.

Optimal Transport between Different Dimensions

Another avenue of investigation is to use optimal transport between distributions that are supported in spaces of different dimensions [McCann andPass, 2019, Pass, 2010] to regularize pooling and embedding layers that change the dimension of their input. This could extend the work done here to different architectures, and could be particularly useful to regularizing generative network where we would like to start generating in lower resolutions in the earlier layers before upsampling to use less resources and computation time.

Simulating Continuity Partial Differential Equations

In Appendix C.3, the minimizing movement scheme defined in Definition 3 in Section 4.3.2 as

ρ τ k+1 ∈ arg min ρ∈P(Ω) Z(ρ) + 1 2τ W 2 2 (ρ, ρ τ k )
for Z : W 2 (Ω) → R, is shown to converge, under some conditions on Z, as τ goes to 0, to a curve ρt in W 2 (Ω) that solves the continuity equation

∂ t ρ t (x) -∇ • (ρ t (x)∇( δZ δρ (ρ t ))(x)) = 0
where δZ δρ is the first variation of Z defined in Definition 20 in Appendix C.3. For example, if Z(ρ) = g(ρ(x)) dx with g(x) = x log x, we obtain the heat equation

∂ t ρ t -∆ρ = 0. And if Z(ρ) = g(ρ(x)) dx + V dρ, we obtain the Fokker-Planck equation ∂ t ρ t -∆ρ -∇ • (ρ∇V ) = 0.
This suggests that solving the problems in the minimizing movement scheme by approximating optimal transport maps through residual neural networks could be a way to discretize and numerically solve such continuity partial differential equations.

Conclusion

This principle of least action is central in physics. It can be found in classical and relativistic mechanics, thermodynamics and quantum mechanics [START_REF] García-Morales | Thermodynamics based on the principle of least abbreviated action: Entropy production in a network of coupled oscillators[END_REF], Feynman, 1942, Gray, 2009]. It states that the dynamical trajectory of a system is one that makes a certain action locally stationary. Since essentially all the laws of physics can be derived through some version of this principle, this suggests that nature operates in an optimal and efficient way. The principle of least action was therefore extended to derive laws of mathematical economy [Kombarov, 2021] and of evolution and natural selection [START_REF] Kaila | Natural selection for least action[END_REF]. Neural networks that respect conservation laws have been proposed to approximate physical systems where conservation of energy is important [START_REF] Greydanus | Hamiltonian neural networks[END_REF].

Here, we have found that this principle is also central to finding good solutions to machine learning problems outside of physics, and to doing so while minimizing the resources and computations required. In particular, it explains why many techniques such residual connections, batch normalizations and orthogonal initialization have been so successful in deep learning, and opens the door to new considerations when designing novel neural networks and their training methods. A Background on Numerical Methods for Differential Equations

List of Acronyms

We refer to [START_REF] Quarteroni | Numerical Mathematics[END_REF] for this quick background on numerical methods for solving ordinary differential equations.

Given an interval I = [t 0 , t 0 + T ] and a function f :

I × R d → R d , we consider the initial value problem of finding a function x : I → R d that satisfies x ′ (t) = f (t, x(t))
with initial condition x(t 0 ) = x 0 , or equivalently in integral form, a function x that satisfies

x(t) = x 0 + t t0 f (t, x(t)) dt
It is often difficult to find closed-form solutions to initial value problems. We wish then to approximate the values x takes at certain time points. We consider a subdivision t 0 < t 1 < .. < t N = t 0 + T of I. Denote the time-steps (or step-sizes) h n := t n+1 -t n for 0 ≤ n < N and define h max := max h n .

A.1 The Euler Method

In a one-step method, an approximation of x(t n ) is x n given iteratively starting from a known initial condition x 0 by

x n+1 -x n h n = ϕ(t n , x n , t n+1 , x n+1 , h n ) If f (x n+1
) does not appear in the definition of ϕ, the method is called explicit. Otherwise, it is called implicit, as x n+1 depends then implicitly on itself through f and an equation needs to be solved to find it, for example, using a fixedpoint [Shashkin, 1991] or a Newton-Raphson [Ypma, 1995] method. So for ϕ(t n , x n , t n+1 , x n+1 , h n ) = f (t n , x n ), we get the explicit Euler method:

x n+1 = x n + h n f (t n , x n ) And for ϕ(t n , x n , t n+1 , x n+1 , h n ) = f (t n+1 , x n+1 ), we get the implicit Euler method x n+1 = x n + h n f (t n+1 , x n+1 )
We only need the convergence analysis of explicit methods, so in the following we will limit ourselves to methods given by

x n+1 -x n h n = ϕ(t n , x n , h n )
Definition 4. For a one-step method, the consistency errors e n , for 0 ≤ n < N , are

e n := x(t n+1 ) -Φ(t n , x(t n ), h n ) h n = x(t n+1 ) -x(t n ) h n -ϕ(t n , x(t n ), h n )
where x is solution. The local (truncation) errors are h n e n . The method is consistent if max |e n | goes to zero as h max goes to zero. For p ∈ N * , the method has order p if max |e n | ≤ Ch p max for a constant C that depends on f, t 0 and T . Theorem 9. If f and ϕ are continuous, then the one-step method is consistent if and only if ϕ(t, x, 0) = f (t, x) for all (t, x).

Theorem 10. If f is C p and ϕ is C p in h then the one-step method is of order p if and only if ∂ k h ϕ(t, x, 0) = 1 k+1 f [k] (t,
x) for all (t, x) and 0 ≤ k < p where

f [0] = f and f [k] = ∂ t f [k-1] + f ∂ x f [k-1] . Corollary 11. If f is continuous then Euler's method is consistent. If f is C 1 then Euler's method has order 1. Definition 5. A one-step method is zero-stable (or stable) if ∃ S > 0 such that for all (x n ) 0≤n≤N , (x n ) 0≤n≤N and (ε n ) 0≤n<N satisfying x n+1 -x n h n = ϕ(t n , x n , h n ) and xn+1 -xn h n = ϕ(t n , xn , h n ) + ϵ n (44)
for 0 ≤ n < N , we have max n ∥x n -x n ∥ ≤ S(∥x 0 -x 0 ∥ + T max n |ϵ n |), where ϵ n = ε n /h n . The constant S is the stability constant of the method.

Theorem 12. If ϕ is uniformly L-Lipschitz in its second variable, then the one-step method is stable with constant e LT .

Corollary 13. If f is Lipschitz in its second variable, then Euler's method is stable. Definition 6. A numerical method converges if its global error max n ∥x(t n )-x n ∥ goes to zero as h max goes to zero.

Theorem 14. If a method is consistent and stable with stability constant S, then it converges and

max n ∥x(t n ) -x n ∥ ≤ ST max |e n |. If the method is of order p with constant C, then max n ∥x(t n ) -x n ∥ ≤ ST Ch p max Corollary 15. Euler's method converges if f is C 0 and Lipschitz in x. If f is also C 1 then it converges with speed O(h max ).
A.2 Runge-Kutta Methods Definition 7. Given q ∈ N * , a Runge-Kutta method with q stages generates the estimates (x n ) 0≤n≤N through, for 0≤n<N ,

t n,i = t n + c i h n , x n,i = x n + h n i-1 j=1 a i,j p n,j , p n,i = f (t n,i , x n,i ), 0 ≤ i ≤ q 120 t n+1 = t n + h n , x n+1 = x n + h n q i=1 b i p n,i
where c i , b i , a i,j ∈ [0, 1] for all 0 ≤ i ≤ q and 1 ≤ j < i with

c 1 = 0 so that t n,1 = t n , x n,1 = x n and p n,1 = f (t n , x n ) for 0 ≤ n < N i-1 j=1 a i,j = c i for 0 ≤ i ≤ q and q i=1 b i = 1
This is an explicit one-step method with ϕ given by

             ϕ(t, x, h) = q i=1 b i f (t + c i h, x i ) x i = x + h i-1 j=1 a i,j f (t + c j h, x j ) for 1 ≤ i ≤ q (45)
Remark 16. The conditions on (a i,j ) and (b i ) are so that the numerical integration methods used are of order 1, i.e. exact on constant functions, which is equivalent to the coefficients summing up to the length of the integration interval.

Remark 17. For q = 1 we find Euler's method.

Theorem 18. If f is continuous then Runge-Kutta methods are consistent.

Definition 8. The Runge-Kutta method with q = 2 stages, c 2 = a 2,1 = 1 and b 1 = b 2 = 1/2 is called Heun's method and is given by ϕ

(t, x, h) = 1 2 f (t, x) + 1 2 f (t + h, x + hf (t, x)). Proposition 4. If f is continuous then Heun's method is consistent. If f is C 2
then Heun's method has order 2.

Definition 9. The Runge-Kutta method with q = 4 stages,

c 2 = c 3 = a 2,1 = a 3,2 = 1/2, c 4 = a 4,3 = 1, b 1 = b 4 = 1/6, b 2 = b 3 = 1/3
and all other coefficients equal to zero is the classic RK4 method. It is given by

       t n,1 = t n x n,1 = x n p n,1 = f (t n,1 , x n,1 ) t n,2 = t n + hn 2 x n,2 = x n + hn 2 p n,1 p n,2 = f (t n,2 , x n,2 ) t n,3 = t n + hn 2 x n,3 = x n + hn 2 p n,2 p n,3 = f (t n,3 , x n,3 ) t n,4 = t n + h n x n,4 = x n + h n p n,3 p n,4 = f (t n,4 , x n,4 )
And finally t n+1 = t n + h n and x n+1 = x n + hn 6 (p n,1 + 2p n,2 + 2p n,3 + p n,4 ) Proposition 5. If f is continuous then the RK4 method is consistent. If f is C 4 then the RK4 method has order 4.

121

Theorem 19. If f is continuous and L-Lipschitz in its second variable and the time-steps are bounded then a Runge-Kutta method with q stages is stable with stability constant S = e KT where

K = L q i=1 b i i-1 j=0 L j h j max Corollary 20. Runge-Kutta methods converge if f is C 0 and Lipschitz in x. Corollary 21. Heun's method converges if f is C 0 and Lipschitz in x. If f is also C 2 then it converges with speed O(h 2 max ). Corollary 22. The RK4 method converges if f is C 0 and Lipschitz in x. If f is also C 4 then it converges with speed O(h 4 max ).

A.3 Adaptive Runge-Kutta Methods

Adaptive methods [Söderlind, 2006] estimate the local truncation error of a single Runge-Kutta step and decrease the step size if the error estimate is large (which indicates that the differential equation is not approximated well at that point and more conservative steps need to be taken) and increase it if it is low. This is done by using two interwoven Runge-Kutta methods that have common intermediate stages, meaning that their coefficients p n,i in Definition 7 are the same but the weights b i in front of them differ. The local error is estimated by the difference between the predictions of the two methods. This way, estimating the error has little or negligible computational cost compared to a step of the higher-order method which already requires the computation of the lower-order method. For example, Heun's method includes Euler's method.

During the integration, the step size is adapted such that the estimated error stays below a chosen threshold. If the estimated error is too high, a step is repeated with a lower step size. If the estimated error is much smaller than the threshold, the step size is increased to save time. This results in an (almost) optimal step size, which saves computation time.

For example, a step of the adaptive Bogacki-Shampine adaptive Runge-Kutta method starting from x n at time t n is computed as follows. The gradient evaluation are given by

         p n,1 = f (t n , x n ) p n,2 = f (t n + 1 2 h n , x n + 1 2 h n p n,1 ) p n,3 = f (t n + 3 4 h n , x n + 3 4 h n p n,2 ) p n,4 = f (t n + h n , x n + 2 9 h n p n,1 + 1 3 h n p n,2 + 4 9 h n p n,3 ) (46)
And the predictions for x(t n+1 ) are given by

         x euler n+1 = x n + h n p n,1 x midpoint n+1 = x n + h n p n,2 x order3 n+1 = x n + 2 9 h n p n,1 + 1 3 h n p n,2 + 4 9 h n p n,3 x order2 n+1 = x n + 7 24 h n p n,1 + 1 4 h n p n,2 + 1 3 h n p n3 + 1 8 h n p n,4 (47)
The difference between x order3 n+1 and x order2 n+1 is used to adapt the step-size. Different rules exist. They all require the user to fix a tolerance ε > 0. The simplest [START_REF] Quarteroni | Numerical Mathematics[END_REF] is to simply halve the step-size if the error err = ∥x order3 n+1 -x order2 n+1 ∥ is larger than ε, to double it if the error is smaller than ε/K, and to keep it as is if the error is between ε and ε/K. Usually K = 2 p+1 , where p is the order of the method. Other rules [Söderlind, 2006] are

h n+1 = ( ε err ) p h n and err = 1 d d i=1 x (Order3) n+1,i -x (Order2) n+1 sc i 2 where sc i = ε(1 + max(|x (Order3) n+1,i |, |x (Order2) n+1,i

|)

and then

h new = h n min(facmax, max(facmin, α( 1 err ) 1 p )
where usually facmax ∈ [1. [START_REF]Swin Transformer architecture with 4 residual stages[END_REF][START_REF]Swin Transformer architecture with 4 residual stages[END_REF] and α ∈ {0.8, 0.9, 0.25 1/3 , 0.38 1/3 }. Then, if err ≤ 1, the step-size h n is kept as is and h n+1 is set to h new . And if err > 1, then h n is set to h new and the step is redone.

(50). Then under some conditions the Kantorovich cost J is the largest lower semi-continuous functional on Π(µ, ν) that is smaller than the Monge cost K.

For example, if X = Y = Ω ⊂ R d with Ω compact and c is continuous and µ atomless, then problems ( 48) and ( 49) have the same infimum and the set {γ T = (id, T ) # µ | T # µ = ν} is dense (in a weak- * sense) in Π(µ, ν) so any transport plan can be approximated by transport maps. This allows to prove that J is a relaxation of K.

We have general existence results for the Kantorovich problem. We will prove the following one.

Theorem 23. If X and Y are separable compact metric spaces and c is continuous (or lower semi-continuous), then the Kantorovich problem admits a solution.

Proof. The proof uses the direct method of calculus of variations. First, we know by Riesz representation theorem that, for a separable locally compact metric space E, the space of finite signed Borel measures on E is the dual of C 0 (E) (the set of continuous real valued functions on E that vanish at infinity, which is equal to the set of continous bounded functions C b (E) and to the set of continuous functions C(E) when E is compact). The norm of an element of the dual is the total mass of the measure. Π(µ, ν) is then included in the dual of C(X × Y ) and is bounded since all its elements have mass equal to 1. By the Banach-Alaoglu theorem the unit ball of the dual space is weak- * compact in duality with C(X × Y ).

Denote J(γ) := X×Y c dγ and l the value of the infimum in (49). Take a minimizing sequence (γ n ) n of problem (49), i.e. a sequence such that J(γ n ) → l and γ n ∈ Π(µ, ν). By compacity, we can assume up to sub-sequencing that γ n * ⇀ γ. This implies for c continuous that J(γ n ) → J(γ) and therefore that J(γ) = l. It remains to show that γ ∈ Π(µ, ν). First, γ is clearly a probability measure since by taking the constant function equal to 1 everywhere we have

γ(X × Y ) = X×Y 1 dγ = lim X×Y 1 dγ n = 1. Then take φ ∈ C(X). We have X×Y φ(x) dγ(x, y) = lim X×Y φ(x) dγ n (x, y) = φ dµ, because γ n is in Π(µ, ν). Likewise for ψ ∈ C(Y ), X×Y ψ(y) dγ(x, y) = ψ dν. We used that (π x ) # γ = µ is equivalent to X×Y φ(x) dγ(x, y) = φ dµ for all test functions φ and that γ n * ⇀ γ by definition means X×Y f dγ n → X×Y f dγ for all f ∈ C(X × Y ). Finally γ ∈ Π(µ, ν) and so γ is optimal.
Existence for the Kantorovich problem can be extended to the case of complete and separable metric spaces (Polish spaces) X and Y via Prokhorov's theorem. For the Monge problem, we have the following result which we will need later.

Theorem 24. Consider X = Y = Ω ⊂ R d with Ω compact and c(x, y) = h(x-y) with h strictly convex. If c is continuous we have the existence of an optimal transport plan γ. If µ is absolutely continuous with respect to the Lebesgue measure and ∂Ω is µ-negligible, then the optimal transport plan γ is unique and is of the form (id, T ) # µ for a transport map T , meaning that it is constructed from a unique optimal transport map T for the Monge problem.

A final result that we mention is the following, which says that the inverse of the optimal transport map between µ and ν is the optimal transport map from ν to µ,

= id

B.2 Duality

The Kantorovitch problem is a linear optimization problem with convex constraints. Duality theory was hence used to study it. We simply state here the duality result, which will be needed to differentiate Wasserstein distances later.

The first step is to rewrite the constraint γ ∈ Π(µ, ν) as follows

sup φ,ψ X φ dµ + Y ψ dν - X×Y φ(x) + ψ(y) dγ(x, y)
as this quantity is equal to 0 if γ ∈ Π(µ, ν) and to ∞ otherwise. The Kantorovitch problem can then be rewritten

inf γ X×Y c dγ + sup φ,ψ X φ dµ + Y ψ dν - X×Y (φ(x) + ψ(y)) dγ(x, y) (51)
Interchanging the infimum and supremum in the above expression, we get

sup φ,ψ X φ dµ + Y ψ dν + inf γ X×Y (c(x, y) -φ(x) -ψ(y)) dγ(x, y) (52)
The dual of the Kantorovitch problem comes from showing that the maximum of ( 52) is equal to the minimum of the Kantorovitch problem, and from rewriting inf γ X×Y (c(x, y) -φ(x) -ψ(y)) dγ(x, y) as the condition φ(x) + ψ(y) ≤ c(x, y) first, and then as a property known as c-concavity.

Definition 12. Given φ : X → R, its c-transform (c-conjugate) is φ c : Y → R given by φ c (y) = inf x∈X c(x, y) -φ(x)
Similarly, the c-transform of ψ : Y → R is ψ c : X → R given by

ψ c(x) = inf y∈Y c(x, y) -ψ(x)
Definition 13. A function φ : X → R is said to be c-concave if there exists ψ : Y → R such that φ = ψ c. We denote c-conc(X) the set of c-concave functions on X.

Theorem 26. If X and Y are separable compact metric spaces and c is uniformly continuous and bounded, then

min γ∈Π(µ,ν) X×Y c dγ = max φ∈c-conc(X) X φ dµ + Y φ c dν
The functions φ realizing the maximum in the above equation are called the Kantorovitch potentials for the transport from µ to ν.

Theorem 27. Under the conditions of Theorem 24, the optimal transport map between µ and ν is given by T

(x) = x -(∇h) -1 (∇φ(x))
where φ is a Kantorovitch potential for the transport from µ to ν.

B.3 The Wasserstein Space

From now on, we take X = Y = Ω ⊂ R d , with Ω complete (i.e. closed), and c(x, y) = ∥x -y∥ p for p ∈ [1, ∞[ so that an optimal transport plan for the Kantorovich problem exists (by the extension of Theorem 23 to Polish spaces). Note that if p > 1 then c is of the form c(x, y) = h(x -y) with h strictly convex and Theorem 24 stands if Ω is compact and µ absolutely continuous with respect to the Lebesgue measure. Compactness is not necessary if c(x, y) = ∥x -y∥ 2 .

A distance between probability measures can be deduced from the Kantorovich problem. For p ∈ [1, ∞[, we restrict ourselves to the set

P p (Ω) := {µ ∈ P(Ω) | ∥x∥ p dµ(x) < ∞} (53)
where ∥.∥ is the euclidean norm. For µ, ν ∈ P p (Ω) we define W p (µ, ν) to be the p-th root of the minimal value of the Kantorovich problem between µ and ν for the cost c(x, y) = ∥x -y∥ p , i.e.

It can be shown that W p is a distance over P p (Ω) called the Wasserstein distance.

One way to prove this is through the gluing lemma:

Lemma 1. Let (X 1 , µ 1 ), (X 2 , µ 2 ) and (X 3 , µ 3 ) be Polish probability spaces. Take γ 1 ∈ Π(µ 1 , µ 2 ) ⊂ P(X 1 × X 2 ) and γ 2 ∈ Π(µ 2 , µ 3 ) ⊂ P(X 2 × X 3 ). Then there exists γ ∈ P(X 1 × X 2 × X 3 ) such that (π 1,2 ) # γ = γ 1 and (π 2,3 ) # γ = γ 2 , where π i,j is the projection from X 1 × X 2 × X 3 on X i × X j .

Proposition 6. W p is a distance over P p (Ω).

Proof. W p is clearly positive and finite because c is positive and finite, the constraint set is never empty and because of the restriction to P p (Ω). W p is also clearly symmetric: define F (x, y) := (y, x) then F # γ ∈ Π(ν, µ) if and only if γ ∈ Π(µ, ν) and the transport costs are the same because c is symmetric.

Taking γ = (id, id) # µ ∈ Π(µ, µ) gives W p (µ, µ) = c(x, x) dµ(x) = 0. Conversely, W p (µ, ν) = 0 implies the existence of γ ∈ Π(µ, ν) such that c dγ = 0. This means γ has to be concentrated on c -1 ({0}) = {(x, y) | x = y}. This means that γ = (id, id) # µ and that ν = id # µ = µ.

To prove the triangle inequality, consider µ 1 , µ 2 and µ 3 ∈ P p (Ω). There exists γ 1 ∈ Π(µ 1 , µ 2 ) an optimal transport plan between µ 1 and µ 2 , and γ 2 ∈ Π(µ 2 , µ 3 ) an optimal transport plan between µ 2 and µ 3 . Then by Lemma 1, there exists γ ∈ P(Ω 3 ) with (π 1,2 ) # γ = γ 1 and (π 2,3 ) # γ = γ 2 . In particular, this means that (π 1,3

) # γ ∈ Π(µ 1 , µ 3 ) because (π 1,3 ) # γ(A × Ω) = (π 1 ) # γ(A) = (π 1 ) # ((π 1,2 ) # γ)(A) = (π 1 ) # γ 1 (A) = µ 1 (A)
where π 1 denotes the projection on the first coordinate (from Ω 3 or Ω 2 depending on the context) and we used the general formulas

(f • g) # γ = f # (g # γ) and π 1 = π 1 • π 1,2 .
Likewise (π 1,3 ) # γ(Ω × B) = µ 3 (B). Since (π 1,3 ) # γ ∈ Π(µ 1 , µ 3 ) but is not necessarily optimal we can write (denoting d(x, y) := ∥x -y∥)

W p (µ 1 , µ 3 ) ≤ d(x, z) p d(π 1,3 ) # γ(x, z) 1 p = d(x, z) p dγ(x, y, z) 1 p ≤ (d(x, y) + d(y, z)) p dγ(x, y, z) 1 p ≤ d(x, y) p dγ(x, y, z) 1 p + d(y, z) p dγ(x, y, z) 1 p = d(x, y) p dγ 1 (x, y) 1 p + d(y, z) p dγ 2 (y, z) 1 p = W p (µ 1 , µ 2 ) + W p (µ 2 , µ 3 )
where the last equality comes from the optimality of γ 1 and γ 2 and the third inequality comes from the Minkowski inequality ∥f + g∥ p ≤ ∥f ∥ p + ∥g∥ p .

If we take g(t) = d(w(0), w( 1)) for all t ∈ [0, 1], it is immediate that a constantspeed geodesic w is absolutely continuous. Its metric derivative is clearly equal to d(w(0), w(1)) at all times. This implies via Proposition 7 that its length is d(w(0), w(1)) and therefore that it is a geodesic between its endpoints. In fact, we have two other equivalent definitions of constant-speed geodesics.

Proposition 8. Fix p > 1 and consider a curve w connecting x to y in a metric space (E, d), then the following three points are equivalent -w is a constant-speed geodesic -w is absolutely continuous and |w ′ |(t) = d(w(0), w(1)) for almost every t

-w solves min{ 1 0 |v ′ | p (t) dt | v ∈ AC(E), v(0) = x, v(1)
= y} Now we prove that W p (Ω) is a geodesic space when Ω is convex.

Theorem 28. W p (Ω) for Ω convex is a geodesic space, i.e. there exists a geodesic curve between every two points in W p (Ω).

More precisely, for t ∈ [0, 1], define π t : Ω×Ω → Ω through π t (x, y) = (1-t)x+ty. For any µ, ν ∈ W p (Ω), there exists by the previous section an optimal transport plan γ ∈ Π(µ, ν) between µ and ν for the cost c(x, y) = ∥x -y∥ p , i.e. γ is optimal for the problem in (54). The curve defined by µ t := (π t ) # γ for t ∈ [0, 1] is then a constant-speed geodesic between µ and ν.

Proof. Take µ, ν ∈ P p (Ω) and γ ∈ Π(µ, ν) optimal for problem (54). Define the curve (µ t ) t through µ t := (π t ) # γ. It is immediate that µ 0 = µ and µ 1 = ν. We start by proving that for every t < s, W p (µ t , µ s ) ≤ |t -s| W p (µ, ν). First, take γ s t := (π t , π s ) # γ. Then γ s t ∈ Π(µ t , µ s ) because for every measurable set A ⊂ Ω we have

γ s t (Ω × A) = γ((π t , π s ) -1 (Ω × A)) = γ(π -1 s (A)) = ((π s ) # γ)(A) = µ s (A)
and likewise γ s t (B × Ω) = µ t (B). Since γ s t is an admissible transport plan between µ t and µ s but not necessarily an optimal one we can write W p (µ t , µ s ) ≤ ∥x -y∥ p dγ s t (x, y)

1 p = ∥π t (x, y) -π s (x, y)∥ p dγ(x, y) 1 p = |t -s| ∥x -y∥ p dγ(x, y) 1 p = |t -s|W p (µ, ν)
The first equality comes from the definition of γ s t as a push-forward of γ, i.e from the formula F (x, y) dγ s t (x, y) = F (π t (x, y), π s (x, y)) dγ(x, y) for any function F . The second equality comes from the definitions of π t (x, y) and π s (x, y) and the relation ∥(1 -t)x + ty -(1 -s)x -sy∥ = |t -s|∥x -y∥. The third equality comes from the fact that γ is an optimal transport plan between µ and ν.

The inequality W p (µ t , µ s ) ≤ |t -s| W p (µ, ν) for t < s is actually an equality, because if it were a strict inequality we could write

W p (µ, ν) ≤ W p (µ, µ t ) + W p (µ t , µ s ) + W p (µ s , ν) < W p (µ, ν)(t + (s -t) + (1 -s)) = W p (µ, ν)
and so we would have W p (µ, ν) < W p (µ, ν) which is impossible. The equality W p (µ t , µ s ) = |t -s| W p (µ, ν) for all t, s shows that (µ t ) t is a constant-speed geodesic between µ and ν according to Definition 18 and concludes the proof that W p (Ω) for Ω convex is a geodesic space. Convexity of Ω is needed so that π t (Ω × Ω) ⊂ Ω.

If Ω is convex and p > 1, then every constant-speed geodesic in W p (Ω) is of the form ((π t ) # γ) t for an optimal transport plan γ.

B.4 The Dynamic Formulation of Optimal Transport

Chapter 5 of [Santambrogio, 2015] also shows a link between geodesics in W p (Ω) (so optimal transport plans) and solutions of the continuity equation

∂ t µ t + ∇ • (µ t v t ) = 0, µ 0 = µ (61)
where for every t ∈ [0, 1], µ t is a probability measure on Ω and v t : Ω → R d is a vector field. When looking for solutions (µ t , v t ) t to the continuity equation, µ t does not have to be a density as we need weak solutions where we only integrate against µ t .

More precisely, we say that (µ t , v t ) t solves the continuity equation in the weak sense if for any test function ψ ∈ C 1 c ( Ω), the function t → ψ dµ t is absolutely continuous in t and for almost every t ∈ [0, 1] we have

d dt Ω ψ dµ t = Ω ∇ψ • v t dµ t ( 62
)
If µ is Lipschitz continuous in (t, x) and v is Lipschitz continuous in x then being a weak solution is equivalent to being a solution almost everywhere.

If is well known that under some conditions (Ω bounded and v Lipschitz and bounded in x uniformly in t with a flow that remains in Ω), the continuity equation describes the evolution of the density µ t of particles distributed initially according to µ 0 = µ and moving according to the velocity field v t .

More precisely, given a time-dependent vector field v, if points x ∈ Ω are distributed according to µ 0 = µ and move tangentially to v at all times, i.e satisfy

∂ t ϕ t (x) = v t (ϕ t (x)), ϕ 0 (x) = x (63)
where ϕ t (x) is the position of point x at time t, then their distribution across time ((ϕ t ) # µ) t solves the continuity equation ( 61) for the given vector field. This is the unique solution if we restrict ourselves to distributions that are absolutely continuous with respect to the Lebesgue measure. A solution ϕ to ( 63) is called the flow of the differential equation ( 63) determined by v.

The link with optimal transport comes from the equivalence between absolutely continuous curves in W p (Ω) and solutions of the continuity equation.

Theorem 29. Take p > 1 and Ω compact.

-If (µ t ) t is an absolutely continuous curve in W p (Ω), then there exists a vector field

v t ∈ L p (R d ; µ t ) for every t ∈ [0, 1] such that ∂ t µ t +∇•(µ t v t ) = 0
is satisfied in a weak sense.

-Conversely, if (µ t ) t is a family of measure in P p (Ω) and we have for every

t ∈ [0, 1] a vector field v t ∈ L p (R d ; µ t ) with 1 0 ∥v t ∥ L p (µt) dt < ∞ such that ∂ t µ t + ∇ • (µ t v t ) = 0 is satisfied in a weak sense, then (µ t ) t is an absolutely continuous curve in W p (Ω).
-Furthermore, (µ t ) t and (v t ) t solving the continuity equation together must satisfy ∥v t ∥ L p (µt) = |µ ′ |(t) for almost every t.

If Ω is convex and compact and p > 1, we can combine this result with the fact that W p (Ω) is a geodesic space to obtain another formula for the Wasserstein distance.

Theorem 30. On a convex and compact set Ω, for all µ, ν ∈ P p (Ω) with p > 1 we have

W p p (µ, ν) = min 1 0 ∥v t ∥ p L p (ρt) dt | ∂ t ρ t + ∇ • (ρ t v t ) = 0, ρ 0 = µ, ρ 1 = ν
This is known as the Benamou-Brenier formula.

Proof. First, given Definition 17 of geodesic spaces and formula (58) for the length of absolutely continuous curves we can write for any µ, ν ∈ P p (Ω)

W p p (µ, ν) = min 1 0 |ρ ′ |(t) dt | ρ ∈ AC(W p (Ω)), ρ 0 = µ, ρ 1 = ν p
Since we proved the existence of constant-speed geodesics between every two points of W p (Ω), the minimum can be restricted to these geodesics. Given the second equivalent definition we gave of constant-speed geodesics when p > 1 in Proposition 8, we can write

W p p (µ, ν) = min 1 0 |ρ ′ | p (t) dt | ρ ∈ AC(W p (Ω)), ρ 0 = µ, ρ 1 = ν
Finally, the equivalence between absolutely continuous curves in W p (Ω) and solutions of the continuity equation allows us to minimize over theses solutions

W p p (µ, ν) = min 1 0 ∥v t ∥ p L p (ρt) dt | ∂ t ρ t + ∇ • (ρ t v t ) = 0, ρ 0 = µ, ρ 1 = ν
where we also used the relation ∥v t ∥ L p (µt) = |ρ ′ |(t) between (ρ t ) t and (v t ) t that jointly solve the continuity equation.

In this dynamic formulation of the Kantorovich problem we solve for pairs (ρ t , v t ) t satisfying the continuity equation with border conditions ρ 0 = µ and ρ 1 = ν that minimize the energy 1 0 ∥v t ∥ p L p (µt) which is the integral across time of the integral across points of the norms of the tangent vectors to the movement of points distributed initially according to µ, so another way of expressing the total distance traveled by all points.

From an optimal transport plan γ between µ and ν we can construct the geodesic ((π t ) # γ) t and the associated velocity field v t for t ∈]0, 1] without solving the continuity equation. In fact, if an optimal transport map T exists, then v t = (T -id) • T -1 t , where T t = (1 -t)id + tT and is invertible. Since the continuity equation needs to hold only weakly, mass can still be split between different destinations. However, we want to solve directly for (µ t , v t ) t . As we have mentioned above, if we restrict ourselves to time-dependent vector fields v that are Lipschitz and bounded in x uniformly in t and have a flow ϕ that remains in Ω and to distributions that are absolutely continuous with respect to the Lebesgue measure, then the continuity equation has a unique solution (µ t ) for a given v. The minimization is therefore only over vector fields and we have the problem

inf v 1 0 ∥v t ∥ p L p (µt) dt subject to ∂ t µ t + ∇ • (µ t v t ) = 0, µ 0 = µ, µ 1 = ν (64)
Since in our machine learning setting distributions µ and ν are always discrete samples, we prefer to use the constraint (63) on the movement of points:

inf v 1 0 ∥v t ∥ p L p ((ϕt) # µ) dt subject to ∂ t ϕ t (x) = v t (ϕ t (x)) ∀x ∈ Ω, ϕ 0 = id, (ϕ 1 ) # µ = ν (65)
between µ and ν is C 0,η over X. If Y is not convex, there exist two relatively closed sets A, B in X, Y respectively such that T ∈ C 0,η (X \ A, Y \ B), where A and B are of null Lebesgue measure.

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

137

C Background on Gradient Flows

We follow [Santambrogio, 2015, Santambrogio, 2016, Ambrosio et al., 2005, Jordan et al., 1998] for this background on Wasserstein gradient flows.

C.1 Gradient Flows in Euclidean Space

Given F : R d → R and x 0 ∈ R d , a gradient flow is a curve x solution of the Cauchy problem x ′ (t) = -∇F (x(t))

x(0) = x 0 (66)

A solution exists and is unique if ∇F is Lipschitz or F is convex. Given τ > 0 and x τ 0 = x 0 define a sequence (x τ k ) k through the minimizing movement scheme

x τ k+1 ∈ arg min x∈R d F (x) + 1 2τ ∥x -x τ k ∥ 2 (67)
F lower semi-continous and F (x) ≥ C 1 -C 2 ∥x∥ 2 guarantees existence of a solution of (67) for τ small enough. F λ-convex meets these conditions and also provides uniqueness of the solution because of strict convexity of the target.

The optimality condition of problem ( 67) is

x τ k+1 -x τ k τ = -∇F (x τ k+1 ) (68)
which is exactly the formula of the implicit Euler scheme (see Appendix A.1) with step size τ .

We interpret the point x τ k as the value of a curve x at time kτ . We then construct a curve x τ as the piecewise constant interpolation of the points

x τ k x τ (t) = x τ k for t ∈](k -1)τ, kτ ], k ∈ {1, .., K} (69)
We can also construct a curve xτ as the affine interpolation of the points

x τ k xτ (t) = x τ k-1 + (t -(k -1)τ ) x τ k -x τ k-1 τ for t ∈](k -1)τ, kτ ], k ∈ {1, .., K} (70)
If F (x 0 ) < ∞ and inf F > -∞ then x τ and xτ converge uniformly to the same curve x as τ goes to zero (up to extracting a subsequence). If F is C 1 , then the limit curve x is a solution of (66) (i.e. a gradient flow of F ). If F is not differentiable then x is solution of the problem defined using the subdifferential of F , i.e. x satisfies x ′ (t) ∈ -∂F (x(t)) for almost every t.

If λ > 0, then the solution to (66) converges exponentially to the unique minimizer of F (which exists by coercivity). So by taking τ → 0 and k → ∞, we tend towards the minimizer of F .

C.2 Gradient Flows in Metric Spaces

The advantage of the minimizing movement scheme ( 67) is that it can be adapted to metric spaces. If (E, d) is a metric space, x 0 ∈ E and F : E → R ∪ {∞} is lower semi-continuous and bounded from below, we can define a sequence (x τ k ) k with x τ 0 = x 0 and

x τ k+1 ∈ arg min x∈E F (x) + 1 2τ d(x, x τ k ) 2 (71)
If E is compact then a solution to (71) exists. We study the limit as τ goes to 0 of the piecewise constant curve x τ constructed from these points by ( 69). In the euclidean setting, we had uniform convergence to a gradient flow of F .

A curve x : [0, T ] → E is called a minimizing movement for F and x 0 if there exists a sequence of time steps (τ j ) j that goes to 0 such that the sequence of piecewise constant curves (x τj ) j uniformly converges to x as j → ∞.

Again, if F (x 0 ) < ∞ and inf F > -∞ then a minimizing movement exists because we can extract a uniformly converging subsequence of curves from the family of piecewise constant curves (x τ ) τ as τ → 0.

If we add some structure to (E, d) by assuming it is a geodesic space, we can also construct a metric counterpart of the affine interpolation (70) using geodesics

xτ (t) = w x τ k-1 ,x τ k t -(k -1)τ τ for t ∈](k -1)τ, kτ ] and k ∈ {1, .., K} (72)
where w x τ k-1 ,x τ k is any constant-speed geodesic between x and y. The curve xτ is continuous and locally Lipschitz (because geodesics are) and coincides with the piecewise constant curve x τ at times kτ . If F (x 0 ) < ∞ and inf F > -∞ then (x τ ) τ is also compact and converges to the same minimizing movement as x τ when τ → 0.

However, if x is a minimizing movement for F we cannot say that it is a gradient flow of F , i.e. that it satisfies (66), because x ′ = -∇F does not mean anything in a metric space. We characterize gradient flows in a metric space by finding equivalent conditions to x ′ = -∇F in the euclidean setting that have a metric counterpart. The first such condition is obtained through the Cauchy-Schwartz inequality and is called the energy dissipation equality (EDE)

F (x(s)) -F (x(t)) = t s 1 2 ∥x ′ (r)∥ 2 + ∥∇F (x(r))∥ 2 dr ∀ s < t (73)
In the metric setting, we use the metric derivative of a curve |x ′ |(t) and the descending slope of a real-valued function

|∇ -F |(x) := lim sup y→x (F (x) -F (y)) + d(x, y) (74)
Under some conditions on the descending slope of F (in particular lower semicontinuity) a minimizing movement for F , i.e. the limit of an interpolation of the minimizing movement scheme, satisfies the EDE condition (73).

If F is λ-convex, another condition which is equivalent to x ′ (t) ∈ -∂F (x(t)) in the euclidean setting is the evolution variational inequality (EVI)

1 2 d dt ∥x(t) -y∥ 2 ≤ F (y) -F (x(t)) - λ 2 ∥x(t) -y∥ 2 ∀ y ∈ R d (75)
In the metric setting, we replace the euclidean distance by the distance d and obtain a different definition of gradient flows in metric spaces. The EVI condition leads to uniqueness and stability results for gradient flows and implies the EDE condition, but more assumptions are needed to prove that minimizing movements satisfy it.

C.3 Gradient Flows in the Wasserstein Space

In the geodesic metric space W 2 (Ω) with Ω compact and convex, for the functional F : W 2 (Ω) → R ∪ {∞} lower semi-continuous for the weak convergence of measures in duality with C(Ω) and ρ τ 0 = ρ 0 ∈ P(Ω), the minimizing movement scheme (71) becomes

ρ τ k+1 ∈ arg min ρ∈P2(Ω) F (ρ) + 1 2τ W 2 2 (ρ, ρ τ k ) (76)
This problem has a solution because the objective is lower semi-continuous and the minimization is over P(Ω) which is compact by Banach-Alaoglu.

We can construct a piecewise constant interpolation between the measures ρ τ k as in (69), or a geodesic interpolation as in ( 72), where we travel along a geodesic between ρ τ k-1 and ρ τ k in W 2 (Ω) (see Theorem 28 in Appendix B.3 for how to construct geodesics in W 2 (Ω) from optimal transport maps). Again, if F (x 0 ) < ∞ and inf F > -∞ then both interpolations converge uniformly to a limit curve (ρ t ) as τ goes to zero. If F is also λ-convex along geodesics, then we might be able to prove the EVI condition for the minimizing movement ρ.

However, the space W 2 (Ω) allows us to consider another possible property of the minimizing movement that can also provide uniqueness and stability. We first need to define a notion of differentiation for functionals defined on probability spaces.

Definition 19. Given F : P(Ω) → R ∪ {∞}, we say that ρ ∈ P(Ω) is regular for We can now say that the minimizing movement ρ may satisfy the following continuity equation:

F if F ((1 -ε)ρ + ερ) < ∞ for all ε ∈ [0, 1] and all ρ ∈ P(Ω) ∩ L ∞ c (
∂ t ρ t (x) -∇ • (ρ t (x)∇( δF δρ (ρ t ))(x)) = 0 (78)
And we can then use the two following results. First from [Santambrogio, 2016]:

Proposition 9. If the functional F : P(Ω) → R ∪ {∞} is λ-geodesically convex and the curves ρ 0 and ρ 1 satisfy the continuity equation ( 78) then

1 2 d dt W 2 2 (ρ 0 t , ρ 1 t ) ≤ - λ 2 W 2 2 (ρ 0 t , ρ 1 t )
and by Gronwall's inequality

W 2 2 (ρ 0 t , ρ 1 t ) ≤ W 2 2 (ρ 0 0 , ρ 1 0 )e -λt
So we have uniqueness of the solution to (78) for the same initial condition ρ 0 and stability with exponential convergence as t → ∞ if λ > 0.

And from [Santambrogio, 2015]:

Proposition 10. Given a functional F : P(Ω) → R ∪ {∞} and a probability measure ρ ⋆ that is regular for F and that minimizes it, suppose δF δρ (ρ ⋆ ) exists and let l = ess inf δF δρ (ρ ⋆ ) be its essential infimum with respect to the Lebesgue measure.

Then if

δF δρ (ρ ⋆ ) ∈ C(Ω) then ∀x ∈ Ω, δF δρ (ρ ⋆ )(x) ≥ l and ∀x ∈ support(ρ ⋆ ), δF δρ (ρ ⋆ )(x) = l.
And if δF δρ (ρ ⋆ ) is only measurable but ρ ⋆ is absolutely continuous, then for almost every x ∈ Ω, δF δρ (ρ ⋆ )(x) ≥ l and for almost every x ∈ {ρ ⋆ > 0}, δF δρ (ρ ⋆ )(x) = l. Therefore if ρ ⋆ minimizes F while satisfying the conditions of Proposition 10, then the constant curve ρt = ρ ⋆ satisfies the continuity equation ( 78) in the weak sense or almost everywhere as both terms are equal to zero. If the minimizing movement ρ also satisfies the continuity equation ( 78) then by Proposition 9, W 2 (ρ t , ρ ⋆ ) ---→ t→∞ 0 exponentially if F is λ-geodesically convex for λ > 0.

Why does it make sense for ρ to satisfy the continuity equation (78)? By Proposition 10 and Remark 33, the solution ρ τ k+1 to problem (76) satisfies the optimality condition δF δρ

(ρ τ k+1 ) + φ τ = const ( 79
)
where φ is the Kantorovich potential associated with the optimal transport S from ρ τ k+1 to ρ τ k . By Theorem 27 in Appendix B.2 (with h being the Euclidean norm), S(x) = x -∇φ(x), we consider the velocity v of particles moving from ρ τ k to ρ τ k+1 . According to Appendix B.4, it is given by

-v(x) = S(x) -x τ = - ∇φ(x) τ = ∇( δF δρ (ρ τ k+1 ))(x) (80)
Therefore, as τ → 0, we can expect to get in ρ a solution to the continuity equation ( 78), but this has to be proven rigorously for a given F .

D Background on Module-wise Training

The typical setting of (sequential) module-wise training for minimizing a loss L, is, given a dataset D, to solve one after the other, for 1≤k≤K, Problems

(T k , F k ) ∈ arg min T,F x∈D L(F, T (G k-1 (x))
where G k = T k • ... • T 1 for 1≤k≤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 the loss can be computed. The final network trained this way is F K • G K , but we can stop at any depth k and use

F k • G k if it performs better.
In fact, module-wise training suffers often from a well-documented stagnation problem observed in [START_REF] Marquez | Deep cascade learning[END_REF], 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 significantly, or even degrade it.

We present below some recent methods that aim at alleviating this problem and improving the accuracy of module-wise training, either by adding other terms to the loss used to train the modules, or simply through architectural considerations.

D.1 Additional Loss Terms

InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF]. InfoPro start by verifying that vanilla modulewise training outperforms end-to-end training in the first modules, but then destroys too much information between the learned features and both the inputs and the labels and stagnates and gets overtaken by end-to-end training in the later modules. See Figure 27 below. They then propose to train each module to maximize the mutual information between the learned features h and both the input x and its label y. This means minimizing -λ 1 I(h, x) -λ 2 I(h, y)

where I is the mutual information given by I(a, b) = H(b) -H(b|a) where H is the entropy [Shannon, 1948]. The mutual information I(h, y) with the label is approximated using an auxiliary classifier with the cross-entropy loss or a contrastive loss inspired by [Chen et al., 2020b, Khosla et al., 2020, He et al., 2020]. The mutual information I(h, x) with the input is approximated using an auxiliary decoder that tries to reconstruct x from h with a reconstruction loss that is binary cross-entropy.

In [START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF], this method is slightly modified in that the auxiliary decoder now tries to reconstruct the previous feature (i.e the input to the module) and not the initial input x to the network.

PredSim [START_REF] Nøkland | Training neural networks with local error signals[END_REF]. PredSim adds an auxiliary similarity matching loss inspired by similarity losses used in neuroscience [START_REF] Patel | Local learning with neuron groups[END_REF]. The similarity matching loss is the distance between matrices whose entries are meant to measure the similarity between inputs in a mini-batch. The target for such a matrix is then the similarity matrix constructed from the labels of the inputs, which contains 1 at position (i, j) if inputs i and j have the same label, and 0 otherwise. More precisely, given a matrix A made up of stacked vectors a i , we denote S(A) the adjusted cosine similarity matrix or correlation matrix, whose entries are given by S(A) i,j = S(A) j,i = ã⊺ i ãj ∥ã i ∥ 2 ∥ã j ∥ 2 where we mean-center vector a i into vector ãi . Given a mini-batch, we denote Y the one-hot encoded matrix of its labels, and H the matrix of the stacked embeddings h i of the elements of the mini-batch after the current module. The similarity matching loss is then ∥S(Conv(H) -S(Y )∥ 2 F , meaning that the embeddings go through a convolutional auxiliary layer to reduce their dimension. This loss is added to the cross-entropy classification loss for the training of every module.

D.2 Architectural Approaches

DGL [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF]. DGL, and previously [START_REF] Belilovsky | Greedy layerwise learning can scale to imagenet[END_REF], simply consider the best auxiliary architectures for module-wise training. They find that deeper convolutional auxiliary classifiers lead to the best performances. However, these auxiliary classifiers are quite large which reduces the computational savings we desire from module-wise training. We, like InfoPro [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF], limit ourselves to one convolutional layer. DGL then introduce an auxiliary architecture made up of staged resolution reduction phase before the linear classification layer, which they find offers a good trade-off between size and performance.

Sedona [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF]. Sedona design light auxiliary classifiers inspired by the architecture of MobileNetV2 [START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF] that relies on depth-wise and point-wise convolutions and inverted residual blocks. They also introduce an architecture search phase that chooses the best auxiliary classifier architecture from a pool of candidates, and where exactly to insert them in the network. This search phase however leads to a deep early module which reduces the memory savings of module-wise training.

E Background on Adversarial Attacks

E.1 Adversarial Attacks

We present here how some popular adversarial attacks on neural networks that solve an image classification task work. We refer to [START_REF] Aldahdooh | Adversarial example detection for dnn models: a review and experimental comparison[END_REF], Li et al., 2022] for more details.

Below, L is the classification loss (usually cross-entropy), x is a clean image, y is its label and d its dimension as a vector, x is the adversarial image found from x, δ = x -x is the perturbation, and ϵ > 0 is the maximal allowed perturbation size in a given L p norm. For targeted attacks, t ̸ = y is the target label the attack wants the network to predict for x. We denote the pre-softmax outputs of the network Z, while f denotes the softmax outputs.

E.1.1 White-box Attacks

White-box attacks are attacks where the adversary has access to the network's architecture, weights and gradients. We present below three gradient-based attacks (FGM, BIM and PGD), one optimization-based attack (CW), which is one of the strongest and most widely used attacks, and one boundary-based attack (DF).

FGM [START_REF] Goodfellow | Explaining and harnessing adversarial examples[END_REF]. The Fast Gradient Method is a simple early attack. It finds the update direction that maximizes the loss L at each pixel of image x and takes one step in that direction. So

x = x + ϵ sign(∇ x L(x, y)) such that x ∈ [0, 1] d This way ∥x -x∥ ∞ ≤ ϵ.
BIM [START_REF] Kurakin | Adversarial examples in the physical world[END_REF]. The Basic Iterative Method iterates the FGM attack k times. So for 0 ≤ i < k and x0 = x,

xi+1 = xi + α sign(∇ x L(x i , y)) such that xi+1 ∈ [0, 1] d
where α = ϵ/k and x = xk PGD [START_REF] Madry | Towards deep learning models resistant to adversarial attacks[END_REF]. Projected Gradient Descent is similar to BIM but starts from a random perturbation in an L p ball around x. Auto-PGD [Croce and [START_REF] Hein ; Croce | Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks[END_REF]] adds a momentum to the gradient step and adapts the step size α across iterations.

CW [START_REF] Carlini | Towards evaluating the robustness of neural networks[END_REF]. The Carlini-Wagner attack is one of the strongest attacks available. It considers a loss

g(x) = max(max j̸ =t (Z j (x)) -Z t (x), -c)
where c > 0 is a confidence parameter. Minimizing g finds an image with a high softmax output the target t. So the attack then solves the optimization problem

min δ ∥δ∥ p + c g(x + δ) such that x = x + δ ∈ [0, 1] d
The attack also transforms the problem from a box-constrained one to an unconstrained one by rewriting the perturbation as δ = 1 2 (tanh(w) + 1) -x. DF [START_REF] Moosavi-Dezfooli | Deepfool: A simple and accurate method to fool deep neural networks[END_REF]. The DeepFool attack approximates the decision regions of a neural network by polyhedrons, taking inspiration from the case of an affine classifier. The attack then accumulates perturbations that are the shortest projections to the boundary of the polyhedron of the correct label. First the attack computes w j = ∇f j (x) -∇f y (x) and d j = f j (x) -f y (x) for all labels j that are different from the ground truth label y. It then finds the closest boundary thus

k = arg min j̸ =y |d j | ∥w j ∥ 2 The perturbation added to x is then δ = |d k | ∥w k ∥ 2 2 w k
This process is repeated starting from x + δ until the predicted label changes.

Among other white-box attack we mention JSMA (Jacobian Saliency Map Attack) [START_REF] Papernot | The limitations of deep learning in adversarial settings[END_REF], UAP (Universal Adversarial Perturbation) [START_REF] Moosavi-Dezfooli | Universal adversarial perturbations[END_REF], FA (Feature Adversary) [START_REF] Sabour | Adversarial manipulation of deep representations[END_REF], ATN (Adversarial Transformation Networks) [START_REF] Baluja | Adversarial transformation networks: Learning to generate adversarial examples[END_REF] and FAB (Fast Adaptive Boundary) [Croce and Hein, 2020a].

E.1.2 Black-box Attacks

Black-box attacks don't have any knowledge about the the network and can only query it. We present below the substitute model family of black-box attacks, a random search attack and an optimization-based attack.

Substitute model attacks [Papernot et al., 2016a, Papernot et al., 2017, Hu and Tan, 2022]. These attacks simply train a network (called a substitute model) that imitates the predictions of the targeted network. The attacker then uses a white-box attack to find an adversarial sample for the substitute model and feeds this sample to the targeted network.

SA [START_REF] Andriushchenko | Square attack: a query-efficient black-box adversarial attack via random search[END_REF]. The Square Attack selects via random search local square-shaped ϵ-bounded perturbations at random positions to solve the optimization problem

min x Z y (x) -max j̸ =y Z j (x) • Z y (x) such that x ∈ [0, 1] d and ∥x -x∥ p ≤ ϵ
The attack has an L 2 variant and an L ∞ variant.

ZOO [START_REF]Zoo: Zeroth order optimization based black-box attacks to deep neural networks without training substitute models[END_REF]. The Zeroth Order Optimization attack estimates the gradients of the network f using finite differences. Indeed

df (x) dx i ≃ f (x + he i ) -f (x -he i ) 2h and d 2 f (x) d 2 x i ≃ f (x + he i ) -2f (x) + f (x -he i )
h 2 where h > 0 is small and e i is the i-th standard basis vector. In one variant of the attack, these estimations are used to perform stochastic coordinate descent optimization methods [Bertsekas, 2016] to minimize a similar target to that of the Carlini-Wagner attack:

g(x) = max(max j̸ =t (log f i (x)) -log f t (x), -c)
Among other black-box attacks we mention Pixel Attacks [Su et al., 2019, Kotyan andVargas, 2020], ST (spatial transformation) [START_REF] Engstrom | Exploring the landscape of spatial robustness[END_REF], HSJ (Hop-Skip-Jump) [Chen et al., 2020a] and BA (Boundary Attack) [START_REF] Brendel | Decisionbased adversarial attacks: Reliable attacks against black-box machine learning models[END_REF].

E.2 Adversarial Detectors

We describe here how some recent adversarial detectors operate. Generally, an adversarial detector extracts statistics from the activations of an input going through a network, before training a classier on these statistics.

E.2.1 Statistical Approaches

Mahalanobis [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF]. The Mahalanobis detector uses Mahalanobis distances between the activations and a Gaussian fitted to them during network training as detection features. The means of the Gaussian distributions are fitted for each class c thus

μc,m = 1 N c x∈Dc h m (x)
where D c is the set of samples x with class c, N c = #D c is its cardinality, and h m is the embedding function at depth m. The covariance matrix is the same for all classes and is fitted thus

Σm = 1 N (x,y)∈D (h m (x) -μy,m )(h m (x) -μy,m ) ⊺
where D is the dataset of sample-label pairs, and N = #D is the number of samples. The Mahalanobis distances to these Gaussians for an input x are then for every class c

Mc,m (x) = -(h m (x) -μc,m ) ⊺ Σ-1 (h m (x) -μc,m )
The confidence scores for an input x to be classified as adversarial (or out-ofdistribution) or not are then

Ŝm (x) = max c -Mc,m (x)
In practice, the detector (a logistic regression) is trained on Mahalanobis scores Ŝm (x) at five different depths m.

[ [START_REF] Ren | A simple fix to mahalanobis distance for improving near-ood detection[END_REF] propose an improvement over the Mahalanobis distance for outof-distribution detection called RMD (Relative Mahalanobis Distance). They fit another Gaussian to all training samples, regardless of class. So the estimated mean is now μm

= 1 N (x,y)∈D h m (x)
And the estimated covariance is now

Σm = 1 N (x,y)∈D (h m (x) -μm )(h m (x) -μm ) ⊺
Then they compute the Mahalanobis distance to this Gaussian

Mm (x) = -(h m (x) -μm ) ⊺ Σ-1 (h m (x) -μm )
And define the relative Mahalanobis distance as follows

RMD c,m (x) = Mc,m (x) -Mm (x)
The confidence scores used as features for detection are now

Sm (x) = max c -RMD c,m (x)
LID [START_REF] Ma | Characterizing adversarial subspaces using local intrinsic dimensionality[END_REF]. Given a point x sampled from a probability distribution P , and a batch B = {x i } k i=1 of points also sampled i.i.d from P , the MLE estimator of the local intrinsic dimensionality (LID) of P at point x is given by

L(x, B) = - 1 k k i=1 log r i (x) r k (x) -1
where r i (x) is the distance from x to its i-th nearest neighbor in B. For adversarial detection, the features used are L(a, B) where a is an activation at a certain depth from an input (clean or adversarial) going through the network, and B is always a batch of activations of clean samples.

E.2.2 Computer Vision Approaches

NSS [START_REF] Kherchouche | Detection of adversarial examples in deep neural networks with natural scene statistics[END_REF]. The natural scene statistics (NSS) detector uses image quality assessment statistics that detect distortions to detect adversarial samples. It extracts 18 of these statistics from the images themselves and not from their embeddings. We describe here how the first two statistics are extracted.

The mean subtracted contrast normalized (MSCN) coefficients [START_REF] Mittal | Noreference image quality assessment in the spatial domain[END_REF] of an image I are given by

Î(i, j) = I(i, j) -µ(i, j) ϕ(i, j) + c
where c > 0 is a small constant, (i, j) is a pixel position, the local mean µ is

µ(i, j) = 3 k=-3 3 l=-3 w(k, l)I(i + k, j + l)
and the local standard deviation ϕ is

ϕ(i, j) = 3 k=-3 3 l=-3 w(k, l)(I(i + k, j + l) -µ(i, j)) 2
where w is a circularly-symmetric Gaussian weighting function. A generalized Gaussian distribution (GGD) is then fitted to Î and its shape parameter and variance are used as the first two features for detection.

Fourier detectors [START_REF] Harder | Spectraldefense: Detecting adversarial attacks on cnns in the fourier domain[END_REF]. The discrete Fourier transform (DFT) in 2 dimensions of an image I ∈ [0, 1] d×d is given by

F(I)(l, m) = d j,k=0 I(j, k) exp -2πi lj + km d
for l, k ∈ {0, .., d -1}. The magnitude of the Fourier coefficients is given by

|F(I)(l, m)| = Re 2 (F(I)(l, m)) + Im 2 (F(I)(l, m))
And the phase is given by ϕ(F(I)(l, m)) = arctan Re(F(I)(l, m)) Im(F(I)(l, m)) [START_REF] Harder | Spectraldefense: Detecting adversarial attacks on cnns in the fourier domain[END_REF] propose two detectors. This first is called MFS (magnitude Fourier spectrum) and uses the magnitudes of the Fourier coefficients as detection features. The second is called PFS (phase Fourier spectrum) and uses the phases of the Fourier coefficients as detection features. They find it is better to extract these features from the embeddings of the images at certain layers than from the images themselves. But given the high dimensionality of these embeddings, they have to heuristically select a subset of layers to consider.

F Additional Experiments from Section 3

F.2 Transport Visualization

If we pretrain an autoencoder on MNIST, we can use its encoder as the encoder of the ResNet and freeze it during training. This makes it possible to visualize the transport of the data by decoding, via the pretrained decoder, the output of each residual block (see Figure 28 below). We show this below on MNIST. In Figure 29, we see the decodings of a vanilla ResNet trained to achieve 99.4% test accuracy. We add the transport cost with λ 0 =5, τ =1 and s=5. The accuracy barely drops (99.3%) and we see in Figure 30 that the decodings change much less. We do the same thing on CIFAR10, first with a vanilla ResNet-9 in Figure 31. Finally, on CIFAR100, the decodings without regularization are in Figure 33. With the transport regularization, the decodings are in Figure 34.

F.3 Experiments with Fixed λ

In this section, we show some additional experimental results with a model where, instead of using an adaptive optimization algorithm, we simply take the transport cost as a regularizer, thus giving us a minimization objective:

L(θ) + λ C(θ)
This is an easier and more straightforward method, simply considering a relaxed constraint in the optimization problem [START_REF]Accuracy of each module after module-wise training[END_REF]. Aside from the advantage of simpler implementation, it allows for easier fine-tuning of the regularization hyper-parameter which is useful when the datasets and networks are big. The adaptivity is lost but this still leads to better test accuracy than the nonregularized networks. Results on the same tasks as in Section 3.5 are below. In 2 dimensions we can compute the exact W 2 distance between successive point clouds in the network before and after each block, and compare it, in Figure 37 below, to the displacement that the block causes. This confirms that residual networks are indeed already close to being optimal transport maps, and that LAP training encourages this bias.

F.5.2 Comparison to Batch Normalization

We compare our method with batch normalization. We find that the two cooperate well to improve test accuracy on the same 2-dimensional task as in Section 3.2 when the model is too small (1 block, Table 19), too big (100 blocks, Table 20), poorly initialized (N (0, 5) initialization, Table 22) and when the data set is small (50 points, Table 21). LAP-ResNets use λ 0 = 0.1, τ = 0.1 and s = 5.

We are interested in cases where adding the transport cost to the loss helps with the generalization of the network. The advantage of Uzawa's algorithm is that fixing λ 0 = 0.1, τ = 0.1 and s = 5 works well in most cases. A fixed λ requires more tinkering, but when a good value is found, it often performs better than Uzawa's algorithm with those values, and finding better values for Uzawa requires a grid search over three parameters instead of just one. Since the results are similar for both the circles and the moons dataset, we report below the results on the circles dataset only.

A first case where adding C to the loss helps is when the network is small. If the network has only two blocks, then it takes on average around 73 epochs to achieve 100% test accuracy and keep that performance for three consecutive epochs (on average over 100 instances of the problem, an instance being a randomly generated circles dataset and a random initialization of the model's weights). If we add batch normalization to the residual blocks as in Figure 1, then the average number of epochs is around 71. But if we add the transport cost C with fixed weight λ = 0.01, then the average number of epochs to achieve the same performance drops to 54. Adding both batch normalization and the transport cost at the same time does not help, as the average number of epochs goes up to 64.

If the network has only one block (Table 19), then it never reaches 100% test accuracy. The highest test accuracy it reaches after 100 epochs is on average 76.6%. With batch normalization, this figure drops to 75.4%, but with Uzawa's algorithm, it goes up to 82.1%. With both batch normalization and Uzawa, the average best accuracy is even higher at 84.6%. A second case is when the network is too big (Table 20). The network with 100 blocks and no batch normalization has an average test performance of only 89.1%. As expected, adding batch normalization helps greatly as the average test accuracy rises to 99.4%. This performance is improved slightly by adding a transport cost via Uzawa's algorithm and reaches 99.8%. Note that adding the transport cost without batch normalization hinders the test performance in this case. A third instance of the transport cost improving generalization is when we use a big initialization (Table 22). The network has 9 blocks and we initialize the weights and biases of the fully connected layers by drawing from a centered normal distribution with standard deviation equal to 5. Here, batch normalization is essential for the network to recover from this initialization and get close to 100% test accuracy inside 100 epochs. If we further add a transport cost with λ = 0.04, the average highest test accuracy improves from 97.8% to 99.8%. If we add a transport cost via Uzawa's algorithm instead of a fixed weight, we get an average test accuracy of 99.6%. An example of the behavior we want to avoid with a big initialization is in Figure 39 below. With batch normalization (and the same random seeds), we get the behavior in Figure 40. If we also add λ = 0.04, we gain further control as seen in Figure 41. A final case is when the training set is small (Table 21). Here we consider a dataset containing 50 points (40 train and 10 test) and a network with 9 blocks. The test performance of this network is on average 88.2%. We can get to 92.7% average test accuracy with batch normalization but using Uzawa's algorithm and no batch normalization yields the best average performance at 95.7%. Adding batch normalization does not change this.

F.5.3 Additional Transport Visualizations in 2 Dimensions

We provide here further visualizations of transport dynamics in 2 dimensions inside a ResNet-9, as in Section 3.2.2, but now showing point clouds after every block. In Figure 38, we show the normal behavior of a network trained with transport regularization (to be compared with the first row in Figure 12 in Section 3.2.2). The three following figures show how transport regularization and batch normalization help control the behavior of a badly initialized network. Table 24: Test accuracy of parallel TRGL and VanGL with 2 modules, compared to DGL and DDG from [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF] and InfoPro from [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF] Figures 18 and 42 suggest that our regularization helps most when looking at the accuracy of the last block. We confirm this by including in Table 27 the accuracy achieved by the best block for the same experiment as in Table 11 and we notice a more important improvement from the regularization in the accuracy of the last block than in the accuracy of the best block. We also observe that with the regularization the difference between the accuracy of the last block and that of the best block is smaller than without the regularization. We further confirm this through the following experiment. As the network gets deeper (50 blocks trained for 10 epochs each sequentially), we expect training it block-wise to become more difficult, and indeed the improvement from the regularization is slightly larger than usual when looking at the accuracy of the last block for both sequential training methods (Table 28). We also include in Table 29 results from block-wise training of ResNeXt-50-32×4d [Xie et al., 2017], which turns out to be difficult to train block-wise. H Additional Experiments from Section 5

In the tables below, VAN corresponds to detectors trained on a vanilla network, RCE to detectors trained on a network trained with reverse cross entropy and LAP to detectors trained on a transport-regularized network. 'MH' refers to the Mahalanobis detector and 'TR' refers to our detector.

H.1 Implementation Details

For ResNeXt50 ( [Xie et al., 2017]) on CIFAR100 ( [Krizhevsky, 2009]), we train for 300 epochs using SGD with a learning rate of 0.1 (divided by ten at epochs 150, 225 and 250), Kaiming initialization, a batch size of 128 and weight decay of 0.0001. For RCE training, the only changes are that the learning rate is 0.05 and the initialization is orthogonal with a gain of 0.05.

For ResNet110 [START_REF] He | Deep residual learning for image recognition[END_REF]) on CIFAR10 ( [Krizhevsky, 2009]), we train for 300 epochs using SGD with a learning rate of 0.1 (divided by ten at epochs 150, 225 and 250), orthogonal initialization with a gain of 0.05, a batch size of 256, weight decay of 0.0001 and gradient clipping at 5. For RCE training, the only change is that we don't use gradient clipping.

For WideResNet [START_REF] Zagoruyko | Wide residual networks[END_REF]) on TinyImageNet, we train for 300 epochs using SGD with a learning rate of 0.1 (divided by ten at epochs 150, 225 and 250), orthogonal initialization with a gain of 0.1, a batch size of 114 and weight decay of 0.0001.

For the magnitude parameter of the Mahalanobis detector, we try all the values tried in their paper for the magnitude and we report the best results.

H.2 Detection of Seen Attacks

H.4 Detection Rate of Successful Adversarial Samples

As in [START_REF] Lee | A simple unified framework for detecting out-of-distribution samples and adversarial attacks[END_REF], we might be only concerned with detecting adversarial samples that successfully fool the network and that are created from clean samples that are correctly classified. We find that the detection rate of successful adversarial samples is always very high and close to 100% on our detector. The results are in Table 38 for seen attacks and in Table 39 for unseen attacks.

H.5 False Positive Rate

We report here the false positive rate on seen (Table 40) and unseen (Table 41) attacks.

H.6 AUROC

We report in Table 42 the AUROC of seen attacks, and in Table 43 the AUROC of unseen attacks. Note that the AUROC is computed on the class-agnostic random forest detector, not on the ensemble of the class-agnostic and the class-conditional detectors.

The AUROC is the area under the receiver operating characteristic curve (or ROC curve). The ROC curve plots the true positive rate (TPR) against the false positive rate (FPR), at different threshold values from 0 to 1.

For our random forest classifiers, it is the percentage of decision trees in the forest that detect an adversarial sample (or an out-of-distribution sample) that is compared to the threshold. If it is higher than the threshold, than the input is declared adversarial (or out-of-distribution). If it is lower than the threshold, then the input is declared clean.

H.7 Detection of Out-of-Distribution Samples

We report below the AUROC (Table 44 for ResNet110 on CIFAR10 and Table 46 for ResNeXt50 on CIFAR100) and the FPR at a fixed TPR of 95% (Tables 45 for ResNet110 on CIFAR10 and 47 for ResNeXt50 on CIFAR100) for the out-of-distribution detection experiments in Section 5.4.4.

Figure 1 :

 1 Figure 1: Basic residual block architecture proposed in [He et al., 2016a]. BN: batch normalization, ReLU: x → (x) + . Conv: convolution. d: number of channels.

Figure 2 :

 2 Figure 2: Bottleneck residual block (left) and ResNeXt block (right) architectures. From [Xieet al., 2017]. A convolutional layer is shown as (number of input channels, filter size, number of output channels).

Figure 3 :

 3 Figure 3: Bottleneck residual block (left) and inverted residual block (right) architectures.From[START_REF] Sandler | Mobilenetv2: Inverted residuals and linear bottlenecks[END_REF]. The thickness indicates the number of channels.

Figure 4 :

 4 Figure 4: Three consecutive residual blocks.

Figure 5 :

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

Figure 8 :

 8 Figure 8: The Kantorovitch problem from µ to ν. From [Peyre and Cuturi, 2019]. The black dots and curve represent the transport plan γ.

Figure 10 :

 10 Figure 10: Dynamic transport of µ 0 to µ 1 according to velocity field v (ϕ x t is ϕt(x) in the text). From [de Bézenac et al., 2021].

Figure 11 :

 11 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)

Figure 12 :

 12 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 N (0, 5) initialization; third row with a N (0, 5) initialization and the transport cost added to the loss.

Figure 13 :

 13 Figure 13: Decomposition of a neural network

  The base model is a ResNet with 9 residual blocks. Two convolutional layers first encode the image of shape 1 × 28 × 28 into shape 32 × 14 × 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.

Figure 14 :

 14 Figure 14: Test accuracy of ResNets of various depths without batch-normalization on CIFAR10.

  optimization algorithm (19), we solve L(θ) + λ C(θ) with a fixed λ. 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 15 :

 15 Figure 15: Test accuracy during training of ResNeXt50 models on CIFAR100.

  for a representation of this training. The final network trained this way is F K • G K , but we can stop the forward pass at any depth k and keep the network F k • G k for inference if it performs better.

Figure 16 :

 16 Figure 16: Module-wise training.

Figure 17 :

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

  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 Table6in 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 Table11in 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 Figure42in Appendix G with parallel and multi-lap sequential block-wise training.

Figure 18 :

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

Figure 19 :

 19 Figure 19: Example of an adversarial attack from [Goodfellow et al., 2015a].

Figure 20 :

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

Figure 21 :

 21 Figure 21: Histogram of the transport cost C for clean and FGM-attacked test samples with different values of ϵ 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.

Figure 23 :

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

APGDMLE

  Auto Projected Gradient Descent. ATN Adversarial Transformation Networks. AUROC Area Under Receiver Operating Curve. BA Boundary Attack. BFGS Broyden-Fletcher-Goldfarb-Shanno. BIM Basic Iterative Method. Maximum Likelihood Estimator. MLP Multi-Layer Perception. MLS Multi-Lap Sequential. MMD Maximum Mean Discrepancy. MSA Method of Successive Approximations. MSCN Mean Subtracted Contrast Normalized. NODE Neural Ordinary Differential Equation. NR Newton-Raphson. NSS Natural Scene Statistics. ODE Ordinary Differential Equation. OOD Out-Of-Distribution. OT Optimal Transport. PDE Partial Differential Equation. PFS Phase Fourier Spectrum. PGD Projected Gradient Descent. PMP Pontryagin's Maximum Principle. PredSim Prediction Similarity. RCE Reverse Cross Entropy. ReLU Rectified Linear Unit. ResBlock Residual Block. ResNet Residual Network. RK Runge-Kutta. RMD Relative Mahalanobis Distance. ROC Receiver Operating Characteristic. SA Square Attack. Sedona Searching for Decoupled Neural Architectures. SGD Stochastic Gradient Descent. ST Spatial Transformation. SVCCA Singular Vector Canonical Correlation Analysis. SW-MSA Shifted Window Multi-head Self-Attention. TPR True Positive Rate. TR Transport. TRGL Transport-Regularized Greedy Learning. UAP Universal Adversarial Perturbation. VanGL Vanilla Greedy Learning. W-MSA Window Multi-head Self-Attention. ZOO Zeroth Order Optimization.

Theorem 25 .

 25 If µ and ν are absolutely continuous measures supported respectively on compact subsets X and Y of R d with negligible boundaries, and if c(x, y) = h(x -y) with h strictly convex, then there exists a unique couple (T, S) of functions such that the five following points hold -T : X → Y and S : Y → X -T # µ = ν and S # ν = µ -T is optimal for the Monge problem from µ to ν with the cost c -S is optimal for the Monge problem from ν to µ with the cost c -T • S ν-a.s. = id and S • T µ-a.s.

  ∥x -y∥ p dγ(x, y)

  Ω) (absolutely continuous probability measures with densities in L ∞ c (Ω)). Definition 20. Given F : P(Ω) → R ∪ {∞} and ρ ∈ P(Ω) that is regular for F we call the first variation of F at ρ and denote δF δρ (ρ), any measurable function, if it exists, such thatd dε F (ρ + εχ) χ = ρ -ρ with ρ ∈ P(Ω) ∩ L ∞ c (Ω). Remark 33.If there is a unique Kantorovitch potential φ up to additive constants from µ to ν (see Theorem 26 in Appendix B.2), then δW p p (.,ν) δρ (µ) = φ.

Figure 27 :

 27 Figure 27: Linear separability (left), mutual information with input (center), and mutual information with output (right) of features learned with module-wise training with K modules. K=1 is end-to-end training. From [Wang et al., 2021].

F. 1

 1 Implementation DetailsThese are further implementation details for the experiments in Section 3.5. Orthogonal initialization with gain 0.01 is used for ResNet models. Kaiming initialization is used for ResNeXt models. SGD is used for training. The momentum is always set to 0.9 and weight decay to 0.0001. For ResNet models, the learning rate is 0.01 and is divided by 5 at epochs 120, 160 and 200 (when the training goes that far). For ResNeXt models, the learning rate is 0.1 and is divided by 10 at epochs 150, 225 and 250. Batch size is 128 for all experiments. Architectures of ResNet[He et al., 2016a] and ResNeXt[Xie et al., 2017] blocks are standard and exactly as in the references. The ResNets used are single representation ResNets (i.e. containing one residual stage) with 9 blocks. The ResNeXt architecture used is the ResNeXt-50-32×4d from[Xie et al., 2017].

Figure 28 :

 28 Figure 28: Decoding the internal representations of a network.

Figure 29 :

 29 Figure 29: Decodings of the internal representations (outputs of the ResBlocks) after training a ResNet-9 on MNIST (og: original image, ae: encoding, b1: output of block 1...)

Figure 30 :

 30 Figure 30: Decodings of the internal representations (outputs of the ResBlocks) after training a LAP-ResNet-9 on MNIST (og: original image, ae: encoding, b1: output of block 1...)

Figure 31 :

 31 Figure 31: Decodings of the internal representations (outputs of the ResBlocks) after training a ResNet-9 on CIFAR10 (og: original image, ae: encoding, b1: output of block 1...)

Figure 32 :

 32 Figure 32: Decodings of the internal representations (outputs of the ResBlocks) after training a LAP-ResNet-9 on CIFAR10 (og: original image, ae: encoding, b1: output of block 1...)

Figure 33 :

 33 Figure 33: Decodings of the internal representations (outputs of the ResBlocks) after training a ResNet-9 on CIFAR100 (og: original image, ae: encoding, b1: output of block 1...)

Figure 34 :

 34 Figure 34: Decodings of the internal representations (outputs of the ResBlocks) after training a LAP-ResNet-9 on CIFAR100 (og: original image, ae: encoding, b1: output of block 1...)

F. 5 . 1

 51 Comparison to Exact Wasserstein Distances

Figure 37 :

 37 Figure 37: Average residual function norm and W 2 distance between successive point clouds in a ResNet-9 trained on the circles test set. Vanilla training on the left and LAP training on the right.

Figure 38 :

 38 Figure 38: Circles test set after each block after training with transport regularization with weight λ = 0.07.

Figure 40 :

 40 Figure 40: Circles test set after each block after training with N (0, 5) initialization and batch normalization.

Figure 41 :

 41 Figure 41: Circles test set after each block after training with N (0, 5) initialization, batch normalization and transport regularization with fixed weight λ = 0.04.

Figure 42 :

 42 Figure 42: Test accuracy after each block of 10-1 ResNet averaged over 10 runs with 95% confidence intervals. Left: multi-lap sequential vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training on 10% of the CIFAR100 training set. Right: parallel vanilla (VanGL, in blue) and regularized (TRGL, in red) block-wise training on 10% of CIFAR100 training set.

Table 1 :

 1 Average highest test accuracy and 95% confidence interval of ResNet-9 over 50 instances on MNIST with training sets of different sizes (in %)

	Train size	ResNet	LAP-ResNet
	500	90.8 ± 0.4	90.9 ± 0.2
	400	88.4 ± 0.5	88.4 ± 0.4
	300	83.5 ± 0.5	86.2 ± 0.4
	200	74.9 ± 1.0	82.0 ± 0.5
	100	56.4 ± 1.5	70.0 ± 1.0

Table 2

 2

	: Average highest test accuracy and 95% confidence interval of ResNet-9 over 20
	instances on CIFAR10 with training sets of different sizes (in %)
	Train size	ResNet	LAP-ResNet
	50 000	91.49 ± 0.10 91.94 ± 0.10
	30 000	88.61 ± 0.14 89.41 ± 0.10
	20 000	85.73 ± 0.14 86.74 ± 0.13
	10 000	79.25 ± 0.25 80.90 ± 0.16
	5 000	70.32 ± 0.32 72.58 ± 0.21
	4 000	67.80 ± 0.26 70.12 ± 0.30

Table 3 :

 3 Average highest test accuracy and 95% confidence interval of ResNeXt50 over 10 instances on CIFAR100 with training sets of different sizes (in %)

	Train size ResNeXt50 LAP-ResNeXt50
	50 000	72.97 ± 1.18	76.11 ± 0.78
	25 000	62.55 ± 2.37	64.11 ± 1.85
	12 500	45.90 ± 2.74	48.23 ± 1.84

Table 4 :

 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 Table2in[START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF], with memory saved compared to E2E as a percentage of E2E memory in red.

	Dataset	Architecture K	Parallel VanGL	Parallel TRGL (ours)	PredSim	DGL	Sedona	E2E
		VGG-19	4	56.17 ± 0.29 (↓ 27%)	57.28 ± 0.20 (↓ 21%)	44.70	51.40	56.56	58.74
	TinyImageNet	ResNet-50 ResNet-101	4 4	58.43 ± 0.45 (↓ 26%) 63.64 ± 0.30 (↓ 24%)	60.30 ± 0.58 (↓ 20%) 63.71 ± 0.40 (↓ 11%)	47.48 53.92	53.96 53.80	54.40 59.12	58.10 62.01
		ResNet-152	4	63.87 ± 0.16 (↓ 21%)	64.23 ± 0.14 (↓ 10%)	51.76	57.64	64.10	62.32

Table 5 :

 5 Accuracy of Parallel TRGL with 2 modules (average and 95% confidence over 3 runs) on CIFAR100, compared to DDG, FR and E2E from Table2in[Huo et al., 2018a].

	Dataset	Architecture K	Par VanGL	Par TRGL (ours) DDG	FR	E2E
	CIFAR100	ResNet-101 ResNet-152	2 2	77.31 ± 0.27 75.40 ± 0.75	77.87 ± 0.44 76.55 ± 1.90	75.75 73.61	76.90 76.52 76.39 74.80

  from Table2in[START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF]. 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 6 :

 6 Accuracy of Parallel TRGL with K modules (average and 95% interval over 5 runs) on STL10, compared to DGL, InfoPro and E2E from Table2in[START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF].

	Dataset Architecture	K	Par VanGL	Par TRGL (ours)	DGL	InfoPro Softmax InfoPro Contrast	E2E
			2	79.85 ± 0.93	80.04 ± 0.85	75.03 ± 1.18	78.98 ± 0.51	79.01 ± 0.64	77.73 ± 1.61
	STL10	ResNet-110	4 8	77.11 ± 2.31 75.71 ± 0.55	79.72 ± 0.81 77.82 ± 0.73	73.23 ± 0.64 72.67 ± 0.24	78.72 ± 0.27 76.40 ± 0.49	77.27 ± 0.40 74.85 ± 0.52	77.73 ± 1.61 77.73 ± 1.61
			16 73.57 ± 0.95	77.22 ± 1.20	72.27 ± 0.58	73.95 ± 0.71	73.73 ± 0.48	77.73 ± 1.61

Table 7 :

 7

	Architecture	Dataset	K	Parallel VanGL	Parallel TRGL (ours)	InfoPro	InfoProL	E2E
		STL10	4	67.00 ± 1.36 (↓ 55%) 67.92 ± 1.12 (↓ 50%) 64.61 (↓ 38%) 66.89 (↓ 45%) 72.19
	Swin-Tiny	CIFAR10	4	83.94 ± 0.42 (↓ 33%) 86.48 ± 0.54 (↓ 29%) 83.38 (↓ 38%) 86.28 (↓ 45%) 91.37
		CIFAR100	4	69.34 ± 0.91 (↓ 33%)	74.11 ± 0.31 (↓ 29%) 68.36 (↓ 38%) 73.00 (↓ 45%) 75.03

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

[START_REF] Pathak | Local learning on transformers via feature reconstruction[END_REF]

, with memory saved compared to E2E as a percentage of E2E memory consumption in red.

Table 8 :

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

	Seq VanGL Seq TRGL MLS VanGL MLS TRGL
	73.31	73.61	74.34	74.78
	4.5.2 Ablation Study and Sensitivity to Hyperparameter τ

Table 9 :

 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%.We summarize inTable 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

			Equally deep modules			Equally heavy modules
	K	Par VanGL	Par TRGL (ours)	InfoPro	Par VanGL* Par TRGL* (ours)	InfoPro*
	4	27% (77.11)	24% (79.72)	18% (78.72)	41% (77.14)	39% (78.94)	33 % (78.78)
	8	50% (75.71)	48% (77.82)	37% (76.40)
	16 61% (73.57)	58% (77.22)	59% (73.95)

  Time of Parallel TRGLSince 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 modulewise 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 TrainingFull 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 modulewise training, it has been used in practice in very constrained settings such as on-device training of sensors

  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.

Table 10 ,

 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 10 :

 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.

	Train size	Seq VanGL	Seq TRGL	MLS VanGL	MLS TRGL	Par VanGL	Par TRGL	E2E
	50000	68.74 ± 0.45	68.79 ± 0.56	69.48 ± 0.53	69.95 ± 0.50 72.59 ± 0.40	72.63 ± 0.40 75.85 ± 0.70
	25000	60.48 ± 0.15	60.59 ± 0.14	61.33 ± 0.23	61.71 ± 0.32 64.84 ± 0.19	65.01 ± 0.27 65.36 ± 0.31
	12500	51.64 ± 0.33	51.74 ± 0.26	51.30 ± 0.22	51.89 ± 0.30 55.13 ± 0.24	55.40 ± 0.35 52.39 ± 0.97
	5000	36.37 ± 0.33	36.40 ± 0.40	33.68 ± 0.48	34.61 ± 0.59 39.45 ± 0.23	40.36 ± 0.23 36.38 ± 0.31

Table 11 :

 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.

	Train	Seq VanGL	Seq TRGL	E2E
	50000 88.02 ± .18	88.20 ± .24 91.88 ± .18
	25000 83.95 ± .13	84.28 ± .22 88.75 ± .27
	10000 76.00 ± .39	77.18 ± .34 82.61 ± .35
	5000	67.74 ± .49	69.67 ± .44 73.93 ± .67
	1000	45.67 ± .88	51.34 ± .90 50.63 ± .98

  WideResNet on TinyImageNet are in the third row of Table12. 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 Table38in Appendix H.4 and are higher than 95% on our detector. False positive rates are in Table40in Appendix H.5 and are always less than 5% on our detector. The AUROC is in Table42in Appendix H.6. On all these metrics, our detector outperforms the other detectors, and LAP-training greatly improves the performance of the Mahalanobis detector.

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

Table 12 :

 12 Average detection accuracy of seen attacks on Network/LAP-Network.

	Attack

Table 13 :

 13 Average detection accuracy of unseen attacks after training on FGM.

	Attack

Table 14 :

 14 Average OOD detection accuracy and standard deviation over 5 runs using ResNet110 trained on CIFAR10.

		OOD Experiment 1	OOD Experiment 2
	Detector	CIFAR100 (seen) SVHN (unseen) CW-CIFAR10 (seen) SVHN (unseen)
	VAN TR	98.30 ± 0.46	97.46 ± 0.49	97.42 ± 0.57	91.38 ± 0.95
	RCE TR	98.42 ± 0.40	98.20 ± 0.39	91.54 ± 6.06	77.58 ± 6.72
	LAP TR	98.30 ± 0.22	98.50 ± 0.47	97.28 ± 0.62	85.46 ± 2.64
	VAN MH	86.88 ± 1.52	91.28 ± 0.92	81.80 ± 1.96	83.76 ± 1.13
	RCE MH	94.82 ± 0.45	92.16 ± 0.57	76.74 ± 2.75	54.24 ± 3.46
	LAP MH	94.84 ± 0.41	90.46 ± 1.45	89.68 ± 0.65	76.72 ± 1.73
	Table 15: Average OOD detection accuracy and standard deviation over 5 runs using ResNeXt50
	trained on CIFAR100.
		OOD Experiment 1	OOD Experiment 2
	Detector	AA-CIFAR100 (seen) SVHN (unseen) CW-CIFAR100 (seen) SVHN (unseen)
	VAN TR	84.48 ± 0.59	75.32 ± 0.62	95.82 ± 0.67	92.92 ± 1.36
	RCE TR	50.04 ± 0.07	55.44 ± 5.76	76.48 ± 0.75	75.66 ± 0.68
	LAP TR	87.10 ± 0.12	72.98 ± 2.72	95.94 ± 0.57	85.94 ± 2.88
	VAN MH	83.44 ± 0.48	78.82 ± 0.48	94.96 ± 0.81	85.10 ± 1.50
	RCE MH	50.04 ± 0.07	58.74 ± 2.36	76.20 ± 0.72	72.20 ± 1.47
	LAP MH	86.04 ± 0.31	86.84 ± 0.68	94.82 ± 0.34	88.92 ± 1.26

Table 16 :

 16 Average highest test accuracy and 95% confidence interval of ResNeXt50 over 10 instances on CIFAR100 with training sets of different sizes (in %)

	Train size	ResNeXt	LAP-ResNeXt50 ResNeXt50 with λ = 0.01
	50 000	72.97 ± 1.18	76.11 ± 0.78	75.96 ± 1.04
	25 000	62.55 ± 2.37	64.11 ± 1.85	64.10 ± 1.74
	12 500	45.90 ± 2.74	48.23 ± 1.84	47.77 ± 1.85

Table 17 :

 17 Average highest test accuracy and 95% confidence interval of ResNet-9 over 10 instances on CIFAR100 with training sets of different sizes (in %)

	Train size	ResNet	LAP-ResNet	ResNet with λ ∈ {0.05, 0.2}
	50 000	72.32 ± 0.24	72.43 ± 0.18	72.62 ± 0.21
	25 000	64.34 ± 0.23	64.34 ± 0.23	64.76 ± 0.24
	10 000	49.27 ± 0.42	50.57 ± 0.23	50.46 ± 0.27
	5 000	34.74 ± 0.84	37.97 ± 0.30	38.44 ± 0.45
	1 000	15.66 ± 0.43	16.42 ± 0.32	16.03 ± 0.49

Table 18 :

 18 Average highest test accuracy and 95% confidence interval of ResNet-9 over 20 instances on CIFAR10 with training sets of different sizes (in %)

	Train size	ResNet	LAP-ResNet	ResNet with λ = 0.2
	50 000	91.49 ± 0.09	91.94 ± 0.10	91.36 ± 0.08
	30 000	88.61 ± 0.14	89.41 ± 0.09	88.50 ± 0.11
	20 000	85.73 ± 0.14	86.74 ± 0.13	85.82 ± 0.11
	10 000	79.25 ± 0.25	80.90 ± 0.16	80.15 ± 0.13
	5 000	70.32 ± 0.32	72.58 ± 0.12	72.03 ± 0.31
	4 000	67.80 ± 0.26	70.12 ± 0.30	69.64 ± 0.30
	1 000	49.22 ± 0.52	51.14 ± 0.45	50.38 ± 0.45
	500	41.55 ± 0.41	42.92 ± 0.37	42.30 ± 0.43
	100	26.98 ± 0.99	25.34 ± 0.76	27.53 ± 0.94

Table 19 :

 19 Average test accuracy and 95% confidence interval over 100 instances on the circles 2D dataset with 1000 points and 1 block (in %)

		No batch normalization Batch normalization
	ResNet	76.6 ± 3.5	75.4 ± 3.1
	ResNet with λ = 0.005	76.5 ± 3.5	75.6, ± 3.3
	LAP-ResNet	82.1 ± 2.6	84.6 ± 3.1

Table 20 :

 20 Average test accuracy and 95% confidence interval over 100 instances on the circles 2D dataset with 1000 points and 100 blocks (in %)

		No batch normalization Batch normalization
	ResNet	89.1 ± 0.9	99.4 ± 0.4
	ResNet with λ = 0.09	69.7 ± 4.0	99.5 ± 0.6
	LAP-ResNet	75.7 ± 2.9	99.8 ± 0.2

Table 21 :

 21 Average test accuracy and 95% confidence interval over 100 instances on the circles 2D dataset with 50 points and 9 blocks (in %)

		No batch normalization Batch normalization
	ResNet	88.2 ± 2.3	92.9 ± 2.0
	ResNet with λ = 0.04	93.5 ± 2.1	94.4 ± 2
	LAP-ResNet	95.8 ± 1.8	96.0 ± 1.3

Table 22 :

 22 Average test accuracy and 95% confidence interval over 100 instances on the circles 2D dataset with a N (0, 5) initialization (in %)

		No batch normalization Batch normalization
	ResNet	90.2 ± 1.3	98.0 ± 0.8
	ResNet with λ = 0.04	89.7 ± 1.5	99.7 ± 0.2
	LAP-ResNet	79.1 ± 3.9	99.4 ± 0.4

Table 25 :

 25 that also split a ResNet-110 in 2 module-wise-parallel-trained modules, on CIFAR 10. Average last block test accuracy and 95% confidence interval of[START_REF]Boundary tilting perspective of adversarial samples[END_REF]-1 ResNet (32 filters, fixed encoder, same classifier) over 20/50 runs on MNIST with different train sizes and parallel (Par) TRGL and VanGL, compared to E2E. .04 99.08 ± .04 99.30 ± .03 30000 98.90 ± .05 98.93 ± .06 99.22 ± .03 12000 98.52 ± .06 98.59 ± .06 98.96 ± .06 6000 98.05 ± .09 98.16 ± .07 98.62 ± .06 1500 96.34 ± .12 96.91 ± .07 97.19 ± .08 1200 95.80 ± .12 96.58 ± .09 96.88 ± .09 600 91.35 ± .99 95.16 ± .15 95.30 ± .17 300 89.81 ± .73 92.86 ± .24 92.87 ± .28 150 81.84 ± 1.22 87.48 ± .42 87.82 ± .59

	Parallel VanGL (ours) Parallel TRGL (ours)	DGL	DDG InfoPro
	94.01 ± .17	94.05 ± .18	93.50	93.41	93.58
	Train size	Par VanGL	Par TRGL		E2E
	60000	99.07 ±

Table 26 :

 26 Average last block test accuracy and 95% confidence interval of 20-1 ResNet (100 filters, fixed encoder, same classifier) over 10 runs on CIFAR10 with different train sizes and parallel (Par) TRGL and VanGL, compared to E2E.

	Train size Par VanGL	Par TRGL	E2E
	50000	85.98 ± .28	86.02 ± .26 93.11 ± .19
	25000	80.94 ± .25	81.09 ± .32 89.10 ± .29
	10000	72.49 ± .46	73.01 ± .31 80.52 ± .46
	5000	62.31 ± .54	64.06 ± .57 69.44 ± .88
	500	38.61 ± .47	41.44 ± .44 40.40 ± .60

Table 27 :

 27 Average highest 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, compared to E2E. .14 88.34 ± .22 91.88 ± .18 25000 84.15 ± .17 84.46 ± .22 88.75 ± .27 10000 76.62 ± .40 77.47 ± .35 82.61 ± .35 5000 69.60 ± .43 70.22 ± .50 73.93 ± .67 1000 51.59 ± .91 52.06 ± .71 50.63 ± .98

	Train	Seq VanGL	Seq TRGL	E2E
	50000 88.14 ±

Table 28 :

 28 Average last block test accuracy and 95% confidence interval of 50-1 ResNet over 10 runs on CIFAR100 with sequential (Seq) and multi-lap sequential (MLS) TRGL and VanGL, compared to E2E.

	Seq VanGL	Seq TRGL	MLS VanGL	MLS TRGL	E2E
	63.40 ± .46	63.86 ± .56	62.59 ± .64	63.24 ± .50 63.34 ± 2.41

Table 29 :

 29

	Seq VanGL	Seq TRGL	MLS VanGL	MLS TRGL	Par VanGL	Par TRGL
	52.29 ± .53	52.42 ± .65	52.59 ± .63	52.84 ± .65	57.86 ± .49	57.93 ± .51

Average test accuracy and 95% confidence interval of ResNeXt50 over 10 runs on CIFAR100 with sequential (Seq), multi-lap sequential (MLS) and parallel (Par) TRGL and VanGL. End-to-end training in this setting achieves an accuracy of 72.97 ± 1.18.

Table 32 :

 32 Average adversarial detection accuracy of seen attacks and standard deviation over 5 runs using ResNet110 on CIFAR10.

	Attack

Table 33 :

 33 Average adversarial detection accuracy of seen attacks and standard deviation over 5 runs using ResNeXt50 on CIFAR100.

	Attack

Table 34 :

 34 Average adversarial detection accuracy of seen attacks and standard deviation over 5 runs using WideResNet on TinyImageNet.

			Attack
	Detector	FGM	APGD	BIM	AA
	VAN TR	95.36 ± 0.36	95.22 ± 0.50	95.26 ± 0.46	81.38 ± 0.38
	LAP TR	95.14 ± 0.49	95.20 ± 0.71	95.12 ± 0.67	81.24 ± 0.53
	VAN MH	81.06 ± 1.10	79.66 ± 1.00	81.20 ± 1.31	78.40 ± 0.72
	LAP MH	85.26 ± 0.96	85.10 ± 0.60	82.46 ± 1.62	78.40 ± 1.04
	VAN NS	94.90 ± 0.70	94.86 ± 0.93	95.00 ± 0.56	81.32 ± 0.22
	H.3 Detection of Unseen Attacks

Table 35 :

 35 Average adversarial detection accuracy of unseen attacks after training on FGM and standard deviation over 5 runs using ResNet110 on CIFAR10.

	Attack

Table 36 :

 36 Average adversarial detection accuracy of unseen attacks after training on FGM and standard deviation over 5 runs using ResNeXt50 on CIFAR100. ± 0.41 93.88 ± 0.97 51.32 ± 0.41 51.62 ± 0.53 51.02 ± 0.43 52.04 ± 0.74 57.9 ± 7.52

	Attack

Table 37 :

 37 Average adversarial detection accuracy of unseen attacks after training on FGM and standard deviation over 5 runs using WideResNet on TinyImageNet.

	Attack

Table 38 :

 38 Average detection rate of successful adversarial samples from seen attacks over 5 runs on Network/LAP-Network.

	Attack

Table 39 :

 39 Average detection rate of successful adversarial samples from unseen attacks after training on FGM over 5 runs.

	Attack

Table 40 :

 40 Average false positive rate of seen attacks over 5 runs on Network/LAP-Network.

	Attack

Table 41 :

 41 Average false positive rate of unseen attacks after training on FGM over 5 runs.

	Attack

Table 42 :

 42 Average AUROC of seen attacks over 5 runs on Network/LAP-Network.

	Attack

Table 43 :

 43 Average AUROC of unseen attacks over 5 runs on Network/LAP-Network.

	Attack

Table 44 :

 44 Average OOD detection AUROC and standard deviation over 5 runs using ResNet110 trained on CIFAR10.

		OOD Experiment 1	OOD Experiment 2
	Detector	CIFAR100 (seen) SVHN (unseen) CW-CIFAR10 (seen) SVHN (unseen)
	VAN TR	99.64 ± 0.13	98.74 ± 0.47	99.32 ± 0.14	96.29 ± 0.71
	RCE TR	99.61 ± 0.09	99.01 ± 0.25	96.53 ± 0.58	86.34 ± 3.98
	LAP TR	99.73 ± 0.09	99.43 ± 0.28	99.31 ± 0.07	96.12 ± 0.59
	VAN MH	92.74 ± 1.50	97.00 ± 0.80	88.38 ± 2.94	88.04 ± 4.03
	RCE MH	97.97 ± 0.25	96.26 ± 0.40	88.28 ± 2.22	79.53 ± 3.78
	LAP MH	98.06 ± 0.38	96.31 ± 0.75	95.22 ± 0.90	86.54 ± 3.60

Table 45 :

 45 Average OOD detection FPR at 95% TPR and standard deviation over 5 runs using ResNet110 trained on CIFAR10.

		OOD Experiment 1	OOD Experiment 2
	Detector	CIFAR100 (seen) SVHN (unseen) CW-CIFAR10 (seen) SVHN (unseen)
	VAN TR	1.16 ± 0.66	2.68 ± 1.05	2.71 ± 1.01	6.68 ± 0.98
	RCE TR	1.16 ± 0.41	1.94 ± 0.86	19.2 ± 2.35	24.98 ± 5.43
	LAP TR	1.02 ± 0.44	1.54 ± 0.56	2.70 ± 0.64	5.68 ± 1.12
	VAN MH	36.42 ± 4.29	15.66 ± 1.69	49.46 ± 4.80	36.78 ± 6.97
	RCE MH	7.34 ± 0.90	20.30 ± 4.45	41.94 ± 6.44	52.72 ± 7.71
	LAP MH	6.98 ± 2.35	17.56 ± 5.76	23.06 ± 6.29	58.40 ± 14.42

Table 46 :

 46 Average OOD detection AUROC and standard deviation over 5 runs using ResNeXt50 trained on CIFAR100.

		OOD Experiment 1	OOD Experiment 2
	Detector	AA-CIFAR100 (seen) SVHN (unseen) CW-CIFAR100 (seen) SVHN (unseen)
	VAN TR	94.96 ± 0.38	78.77 ± 1.15	99.00 ± 0.13	95.17 ± 0.31
	RCE TR	50.12 ± 0.05	50.30 ± 8.13	87.50 ± 1.65	79.99 ± 3.62
	LAP TR	96.45 ± 0.08	76.28 ± 0.73	99.16 ± 0.28	94.84 ± 1.63
	VAN MH	93.32 ± 0.41	85.02 ± 0.97	98.24 ± 0.32	93.07 ± 0.33
	RCE MH	50.15 ± 0.08	58.04 ± 3.56	86.73 ± 2.03	77.42 ± 3.38
	LAP MH	94.87 ± 0.24	92.76 ± 0.39	97.84 ± 0.20	95.92 ± 0.64

Table 47 :

 47 Average OOD detection FPR at 95% TPR and standard deviation over 5 runs using ResNeXt50 trained on CIFAR100.

		OOD Experiment 1	OOD Experiment 2
	Detector	AA-CIFAR100 (seen) SVHN (unseen) CW-CIFAR100 (seen) SVHN (unseen)
	VAN TR	27.68 ± 1.84	32.66 ± 1.41	5.42 ± 0.82	8.90 ± 1.20
	RCE TR	96.12 ± 0.96	94.72 ± 3.10	44.12 ± 3.31	45.68 ± 2.01
	LAP TR	21.62 ± 0.30	29.47 ± 0.59	4.90 ± 0.92	10.52 ± 3.61
	VAN MH	29.75 ± 1.09	39.06 ± 0.96	8.02 ± 0.81	16.38 ± 0.82
	RCE MH	95.28 ± 0.27	91.42 ± 1.26	45.12 ± 3.76	48.45 ± 4.55
	LAP MH	24.74 ± 0.69	34.54 ± 1.16	8.10 ± 0.46	9.80 ± 1.21

List of Tables

Accuracy of LAP-ResNet-9 on MNIST . . . . . . . . . . . . . . . 31 Accuracy of LAP-ResNet-9 on CIFAR10 . . . . . . . . . . . . . . 31 Accuracy of LAP-ResNeXt50 on CIFAR100 . . . . . . . . . . . . 32 Accuracy of parallel TRGL with 4 modules on TinyImageNet . . 50 Accuracy of parallel TRGL with 2 modules on CIFAR100 . . . . 50 Accuracy of parallel TRGL with K modules on STL10 . . . . . . 51 Accuracy of parallel TRGL with 4 modules using a Swin Transformer 51 Accuracy of sequential and MLS TRGL on TinyImageNet . . . . 51 Memory savings of module-wise training on STL10 . . . . . . . . 52 Accuracy of 10-1 ResNet with sequential, MLS and parallel TRGL 54 Last block accuracy of 10-1 ResNet with sequential TRGL . . . . 54 Detection accuracy of seen attacks . . . . . . . . . . . . . . . . . 72 Detection accuracy of unseen attacks . . . . . . . . . . . . . . . . 73 OOD detection accuracy on ResNet110 . . . . . . . . . . . . . . . 74 OOD detection accuracy on ResNeXt50 . . . . . . . . . . . . . . 74 Accuracy of LAP-ResNeXt50 on CIFAR100 . . . . . . . . . . . . 159 Accuracy of LAP-ResNet-9 on CIFAR100 . . . . . . . . . . . . . 159 Accuracy of LAP-ResNet-9 on CIFAR10 . . . . . . . . . . . . . . 159 Accuracy on the circles dataset with 1 block . . . . . . . . . . . . 162 Accuracy on the circles dataset with 100 blocks . . . . . . . . . . 162 Accuracy on the circles dataset with 50 training points . . . . . . 163 Accuracy on the circles dataset with a large initialization . . . . 163 Accuracy of 2-7 ResNet with parallel TRGL . . . . . . . . . . . . 169 Accuracy of parallel TRGL with 2 modules . . . . . . . . . . . . 170 Last block accuracy of 20-1 ResNet with parallel TRGL . . . . . 170 Last block accuracy of 20-1 ResNet with parallel TRGL . . . . . 170 Accuracy of 10-1 ResNet with sequential TRGL . . . . . . . . . . 171 Last block accuracy of 50-1 ResNet with sequential and MLS TRGL171 Accuracy of ResNeXt50 with sequential, MLS and parallel TRGL 171 Accuracy of parallel TRGL* with K modules on STL10 . . . . . 172 Memory savings of module-wise training with 4 modules . . . . . 172 Detection accuracy of seen attacks on ResNet110 . . . . . . . . . 175 Detection accuracy of seen attacks on ResNeXt50 . . . . . . . . . 176 Detection accuracy of seen attacks on WideResNet . . . . . . . . 176 Detection accuracy of unseen attacks on ResNet110 . . . . . . . . 176 Detection accuracy of unseen attacks on ResNeXt50 . . . . . . . 177 Detection accuracy of unseen attacks on WideResNet . . . . . . . 177 Detection rate of successful samples from seen attacks . . . . . . 177 Detection rate of successful samples from unseen attacks . . . . . 178 False positive rate on seen attacks . . . . . . . . . . . . . . . . . 178 False positive rate on unseen attacks . . . . . . . . . . . . . . . . 178 AUROC of seen attacks . . . . . . . . . . . . . . . . . . . . . . . 179 AUROC of unseen attacks . . . . . . . . . . . . . . . . . . . . . . 179

List of Figures

B.1 Monge and Kantorovich Problems

We follow [Santambrogio, 2015, Villani, 2008, Villani, 2021] in this overview of optimal transport. Denote P(E) the space of probability measures on a space E. Given probability measures µ ∈ P(X) and ν ∈ P(Y ) and a cost function c : X × Y → [0, ∞], the starting point of optimal transport is the following problem Definition 10. The Monge problem is inf T :X→Y X c(x, T (x)) dµ(x)

Here T # µ is the push-forward measure of µ through T defined by the formula T # µ(A) := µ(T -1 (A)) ∀ A measurable. For this definition and the constraint to make sense, T has to be measurable with respect to the σ-algebras of µ and ν.

The Monge problem expresses the problem of moving mass distributed according to µ so that it becomes distributed according to ν, while minimizing the sum (so the integral) of efforts made to transport each point, the effort being encoded in the cost function c, which is usually a distance between points.

An admissible T is called a transport map between µ and ν and the objective of the problem is called the transport cost of T .

Figure 24: The Monge problem from µ to ν. From [START_REF] Kolouri | Optimal mass transport: Signal processing and machine-learning applications[END_REF].

The Monge problem is difficult because of the constraint on T which is not closed under weak convergence and a T satisfying it need not even exist (e.g. when µ is a Dirac measure and ν is not). A relaxation exists Definition 11. The Kantorovitch problem is

where

The constraint set Π(µ, ν) of the Kantorovich problem is never empty as it always contains µ ⊗ ν. This relaxation means that instead of moving the entire mass at x to T (x), γ(A × B) denotes the amount of mass moving from A to B, so that two particles located at x ∈ X can move to different destination points in Y (see Figure 25 below). An admissible γ ∈ Π(µ, ν) is called a transport plan between µ and ν. The Kantorovich problem has a smaller infimum than the Monge problem, since if T # µ = ν then γ T := (id, T ) # µ ∈ Π(µ, ν) and the transport costs are equal

If γ ∈ Π(µ, ν) is concentrated on a graph in X × Y , i.e. on a set of the form {(x, T (x)), x ∈ X}, then all the mass in x goes to T (x) and γ = (id, T ) # µ.

The Kantorovich problem is a relaxation of the Monge problem in the following sense: denote the Kantorovich objective J(γ) := c dγ for γ ∈ Π(µ, ν). The Monge objective can also be written as a functional K on Π(µ, ν) defined by K(γ) := J(γ) if γ = γ T for a transport map T and K(γ) := ∞ otherwise.

The Monge problem is equivalent to minimizing K over Π(µ, ν) because of We denote W p (Ω) the metric space P p (Ω) endowed with the distance W p , and call it the Wasserstein space.

Definition 14. A curve w in a metric space (E, d) is a continuous function w : [0, 1] → E. Its metric derivative at time t is

provided the limit exists.

We denote AC(E) the set of absolutely continuous curves valued in E. From the bound in ( 56) we see that absolutely continuous curves admit a metric derivative at all times.

Definition 16. The length of a curve w is

From Definition 15 of absolutely continuous curves and Definition 16 of the length of curve, we can see that every absolutely continuous curve w has finite length, as Length(w) ≤ 1 0 g(s) ds < ∞ for some g ∈ L 1 ([0, 1]). Actually, we have a formula for the length of an absolutely continuous curve.

It is immediate from the definition of length and the triangle inequality that for any x, y ∈ E we have d(x, y) ≤ Length(w) for every curve w connecting x to y.

which means that an absolutely continuous curve must realize the minimum and that its length must be equal to d(x, y). Such a curve is called a geodesic between x and y.

where we used the fact that ((ϕ t ) # µ) t with ϕ solution to ( 63) is a solution to the continuity equation for a given v. Again, since in all generality we don't have uniqueness of ϕ solution to (63), then the optimization is actually over (v t , ϕ t ) t . Since we are going to discretize the problem and use neural networks to approximate the velocity field, we can limit ourselves to uniformly Lipschitz velocity fields v, so that the solution ϕ exists and is unique and the problem as expressed in ( 65) is well defined.

If we want (63) to hold for all t ∈ [0, 1], then this formulation is not exactly equivalent to the first dynamic formulation (64) and to the Kantorovich problem, as all the mass in a position x will follow the velocity field to a unique point y = ϕ 1 (x). We are then back to the Monge problem, but here we look for a velocity field that points distributed according to µ will follow so that they end up distributed according to ν, while minimizing the total energy of this field.

B.5 Regularity of Optimal Transport Maps

Optimal transport maps have some regularity properties under some further assumptions. In particular, Theorem 6.27 of [START_REF] Ambrosio | Gradient Flows in Metric Spaces and in the Space of Probability Measures[END_REF] gives a classical almost-everywhere regularity result:

Theorem 31. If c(x, y) = ∥x-y∥ p for p > 1, and µ and ν have compact supports with d(supp(µ), supp(ν)) > 0, then the optimal transportation map T between µ and ν is ν -a.e. differentiable and its Jacobian ∇T (x) has non-negative eigenvalues ν-almost-surely.

More recently, results summarized below, which correspond to Theorems 4.23, 4.24 and Remark 4.25 of [Figalli, 2017], state that the optimal transportation map has one degree of regularity more than the initial transported density: Theorem 32. Take c(x, y) = ∥x -y∥ p for p > 1. Suppose there are X, Y , bounded open sets, such that the densities of µ and ν are null in their respective complements and bounded away from zero and infinity over them respectively.

Then, if Y is convex, there exists η > 0 such that the optimal transport map T

F.4 Impact on Weight Distribution of Classifier

LAP training moves the weights of the classification layers away from zero: compare Figure 35 for a vanilla network and Figure 36 for a LAP network. but Table 6 where the learning rate is 0.002. For the Swin Transformer in Table 7, we use the AdamW optimizer with a learning rate of 0.007 and a CosineLR scheduler.

For end-to-end training we use a learning rate of 0.1 that is divided by five at epochs 120, 160 and 200. Momentum is always 0.9. For parallel and end-toend training, we train for 300 epochs. For sequential and multi-lap sequential training, the number of epochs varies per module (see Section 4.5.5).

For experiments in Section 4.5.1, we use a batch size of 256, orthogonal initialization [START_REF] Saxe | Exact solutions to the nonlinear dynamics of learning in deep linear neural network[END_REF] with a gain of 0.1, label smoothing of 0.1 and weight decay of 0.0002. Only for Table 6, the batch size changes to 64.

For experiments in Section 4.5.5, we use a batch size of 128, orthogonal initialization with a gain of 0.05, no label smoothing and weight decay of 0.0001.

For the experiments on TinyImageNet in Table 4, we use τ = 500000 for the first two modules and then double it for the last two modules for TRGL.

For the experiments on CIFAR100 in Table 5, we use λ k,1 = 1, h = 1 and s = 50 for TRGL.

For the experiments on STL10 in Table 6, we use τ = 0.5 and double it at the midpoint, expect for the first line where τ = 50.

G.2 Additional Experiments

Table 23: Test accuracy of 2-7 ResNet on CIFAR100 with parallel TRGL and VanGL, compared to DDG and FR from [Huo et al., 2018a] that also split their networks in 2 module-wiseparallel-trained modules.

G.3 Memory Usage

In Table 30 below, we look at the accuracy of a ResNet110 split into K modules that are equally heavy instead of equally deep, which consumes even less memory (see Table 9 in Section 4.5.3). The module-wise training methods are then called VanGL*, TRGL* and InfoPro* as in [START_REF] Wang | Revisiting locally supervised learning: an alternative to end-to-end training[END_REF]. We notice a slight degradation in the performance of module-wise training compared to Table 6, but the benefit of the regularization increases, as we now gain almost two percentage points of accuracy when using 4 modules, and still outperform end-to-end training. Note that methods DDG [START_REF] Huo | Decoupled parallel backpropagation with convergence guarantee[END_REF] 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 [START_REF] Pyeon | {SE-DONA}: Search for decoupled neural networks toward greedy block-wise learning[END_REF] 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 [START_REF] Belilovsky | Decoupled greedy learning of cnns[END_REF] is architecture-wise essentially identical to VanGL and consumes the same amount of memory.

G.4 Sensitivity to Hyperparameter τ

We show in Figure 43 below that TRGL still performs better than VanGL (in the same setting as in Table 6 in Section 4.5.1, i.e. using a ResNet110 trained on STL10 with 16 modules trained in parallel) for values of τ from 0.03 to 100 and is still roughly equivalent to it for values up to 5000.