Designing Stable Neural Networks using Convex Analysis and ODEs
Motivated by classical work on the numerical integration of ordinary
differential equations we present a ResNet-styled neural network architecture
that encodes non-expansive (1-Lipschitz) operators, as long as the spectral
norms of the weights are appropriately constrained. This is to be contrasted
with the ordinary ResNet architecture which, even if the spectral norms of the
weights are constrained, has a Lipschitz constant that, in the worst case,
grows exponentially with the depth of the network. Further analysis of the
proposed architecture shows that the spectral norms of the weights can be
further constrained to ensure that the network is an averaged operator, making
it a natural candidate for a learned denoiser in Plug-and-Play algorithms.
Using a novel adaptive way of enforcing the spectral norm constraints, we show
that, even with these constraints, it is possible to train performant networks.
The proposed architecture is applied to the problem of adversarially robust
image classification, to image denoising, and finally to the inverse problem of
deblurring.
Ferdia Sherry, Elena Celledoni, Matthias J. Ehrhardt, Davide Murari, Brynjulf Owren, Carola-Bibiane Schönlieb
# Designing Stable Neural Networks using Convex Analysis and ODEs

###### Abstract

Motivated by classical work on the numerical integration of ordinary differential equations we present a ResNet-styled neural network architecture that encodes non-expansive (1-Lipschitz) operators, as long as the spectral norms of the weights are appropriately constrained. This is to be contrasted with the ordinary ResNet architecture which, even if the spectral norms of the weights are constrained, has a Lipschitz constant that, in the worst case, grows exponentially with the depth of the network. Further analysis of the proposed architecture shows that the spectral norms of the weights can be further constrained to ensure that the network is an averaged operator, making it a natural candidate for a learned denoiser in Plug-and-Play algorithms. Using a novel adaptive way of enforcing the spectral norm constraints, we show that, even with these constraints, it is possible to train performant networks. The proposed architecture is applied to the problem of adversarially robust image classification, to image denoising, and finally to the inverse problem of deblurring.

keywords: Deep learning, numerical integration of ODEs, convex analysis, monotone operator theory, inverse problems +
Footnote †: journal:

## 1 Introduction

The desire to impose Lipschitz conditions on neural networks has come to the forefront in a number of tasks in recent years, especially because there have been serious concerns about the stability of neural networks ever since it was
shown that high-performance image classifiers may suffer from adversarial examples [16]. These issues need to be satisfactorily resolved before deep learning methods can be considered suitable for application in safety-critical systems. Another important application of Lipschitz neural networks can be found in generative modelling, in particular in models such as Wasserstein generative adversarial networks (GANs) [2]. In these models, the aim is to minimise the Wasserstein distance between the output of a generator neural network and some target distribution:

\[\min_{\Psi}W_{1}(\Psi\#\mu_{\text{latent}},\mu_{\text{true}}), \tag{1}\]

where \(W_{1}\) is the Wasserstein metric, \(\mu_{\text{latent}}\) is a (simple) distribution of latent variables \(Z\in\mathcal{Z}\), \(\Psi\#\mu_{\text{latent}}\) is its pushforward by the generator neural network \(\Psi:\mathcal{X}\rightarrow\mathcal{Z}\) and \(\mu_{\text{true}}\) is the target distribution of \(X\in\mathcal{X}\). Appealing to the Kantorovich-Rubinstein duality, we know that

\[W_{1}(\mu,\nu)=\sup_{f:\mathcal{X}\rightarrow\mathbf{R},\;1-\text{Lipschitz}} \mathbf{E}_{X\sim\mu}[f(X)]-\mathbf{E}_{Y\sim\nu}[f(Y)],\]

where \(f\) is usually called the critic. With this result, eq. (1) becomes the following saddlepoint problem:

\[\min_{\Psi}\sup_{f:X\rightarrow\mathbf{R}\;1-\text{Lipschitz}}\mathbf{E}_{Z \sim\mu_{\text{latent}}}[f(\Psi(Z))]-\mathbf{E}_{X\sim\mu_{\text{true}}}[f(X )].\]

To solve this problem, we are required to flexibly parametrise 1-Lipschitz critic functions \(f:\mathcal{X}\rightarrow\mathbf{R}\).

Another application in which neural network instabilities may cause considerable problems is in the application of deep learning to inverse problems. A prototypical approach to doing this is the so-called Plug-and-Play approach [38; 8; 9; 34], in which some parts of an iterative optimisation algorithm (usually the denoising steps) for a variational reconstruction problem are replaced by a different component, which may be learned separately. An example of this is shown in algorithm 1 (here \(E_{y}\) is a data discrepancy functional measuring the mismatch between an image and corresponding measurements). In this setting, the denoiser \(\Phi\) may be a denoiser trained on natural images.

```
inputs: measurements \(y\), initial estimate \(x^{0}\), denoiser \(\Phi\) \(x\gets x^{0}\) for\(i\gets 1,\dots,\texttt{it}\)do \(x\leftarrow\Phi(x-\tau\nabla E_{y}(x))\) endfor return\(x\)
```

**Algorithm 1** Plug-and-Play proximal gradient method

With such algorithms, we may run into divergent behaviour if we do not restrict \(\Phi\) appropriately [33], but there is recent work showing that the iterative
method will converge as long as certain Lipschitz conditions are imposed on the denoiser \(\Phi\)[30; 20].

Lipschitz continuity is a standard way to quantify the stability of a function. Let us recall its definition and some associated properties: a function \(f:\mathcal{X}\to\mathcal{Y}\) between metric spaces \((\mathcal{X},d_{\mathcal{X}})\) and \((\mathcal{Y},d_{\mathcal{Y}})\) is said to be \(L\)-Lipschitz for some \(L\geqslant 0\) if \(d_{\mathcal{Y}}(f(x_{1}),f(x_{2}))\leqslant Ld_{\mathcal{X}}(x_{1},x_{2})\) for all \(x_{1},x_{2}\in\mathcal{X}\). This notion of stability plays well with the compositional nature of neural networks: if \(f_{1}:\mathcal{X}\to\mathcal{Y}\) and \(f_{2}:\mathcal{Y}\to\mathcal{Z}\) are \(L_{1}\)-Lipschitz and \(L_{2}\)-Lipschitz respectively, their composition \(f_{2}\circ f_{1}\) is \((L_{1}\cdot L_{2})\)-Lipschitz. If \(\mathcal{X}\) and \(\mathcal{Y}\) are in fact normed spaces, we can furthermore see (by definition) that any bounded linear operator \(A:\mathcal{X}\to\mathcal{Y}\) is \(\|A\|\)-Lipschitz, where the norm is the operator norm. In particular, an ordinary feedforward neural network \(\Psi(x)=A^{K}\sigma(b^{K-1}+A^{K-1}\sigma(\ldots+A^{2}\sigma(b^{1}+A^{1}x)))\) with a \(1\)-Lipschitz activation function \(\sigma\) and learnable linear operators \(A^{1},\ldots,A^{K}\) and biases \(b^{1},\ldots,b^{K}\) is \(L\)-Lipschitz, where \(L=\prod_{i=1}^{K}\|A^{i}\|\). This naturally gives rise to the idea of spectral normalisation: if an ordinary feedforward neural network with a given Lipschitz constant \(L\) is required for a specific application, this can be achieved by appropriately normalising the linear operators, as applied to GANs in [22]. It is worth remarking here that we are referring to any \(L\) that satisfies the defining inequality for Lipschitz continuity of \(f\) as a Lipschitz constant of \(f\); often the term Lipschitz constant is used instead to refer only to the infimum of such \(L\), which defines a seminorm on vector spaces of Lipschitz functions. We will refer to this infimum as the optimal Lipschitz constant of \(f\), and note that the statements about the composition of Lipschitz functions, when framed in terms of optimal Lipschitz constants, only give upper bounds in general.

In this work we are focused on the case where \(\mathcal{X}\) and \(\mathcal{Y}\) are equal to each other, as is the case in many image-to-image tasks. Residual networks (ResNets) [19] have proven to be an extremely successful neural network meta-architecture in this setting: a ResNet parametrises a neural network by \(\Psi=(\operatorname{id}+\Psi^{K})\circ\ldots\circ(\operatorname{id}+\Psi^{1})\), where each \(\Psi^{i}\) is a small neural network. Without further constraints, the Lipschitz continuity of such a network may be badly behaved as the depth increases: even if we control each \(\Psi^{i}\) to be \(\varepsilon\)-Lipschitz for some small \(\varepsilon>0\), in the worst case we can not guarantee anything better than that \(\operatorname{id}+\Psi^{i}\) is \((1+\varepsilon)\)-Lipschitz, and that the composition \(\Psi\) is \(L\)-Lipschitz with \(L=(1+\varepsilon)^{K}\), which grows exponentially as \(K\to\infty\). Nevertheless, we show that it is possible to design ResNets that are provably non-expansive (\(1\)-Lipschitz) by discretising non-expansive continuous flows in a sufficiently careful manner (note that it is not guaranteed in general that a discretisation of a continuous flow preserves its structural properties, such as non-expansiveness).

### Related topics

_Lipschitz neural networks_

As mentioned above, within the deep learning community there have been a number of drivers for research into neural networks with controlled Lipschitz constants, such as the desire to increase robustness to adversarial examples [37],
the necessity to model the critic in a Wasserstein GAN as a 1-Lipschitz function [2], and their use in parametrising learnable diffeomorphisms [6]. Spectral normalisation [22] has become a standard approach to constraining the Lipschitz constant of an ordinary feedforward neural network. This approach ensures that the optimal (smallest) Lipschitz constant of a neural network is upper bounded. It is known to be computationally hard to estimate the true optimal Lipschitz constant [39] of a neural network, which has prompted further research into refining Lipschitz neural network architectures.

#### Methods based on continuous dynamical systems

Applied mathematicians and physicists have long studied continuous dynamical systems in the form of ODEs and PDEs, giving rise to an extensive body of research on the structural properties of such systems. More recently, insights from these topics have been used to design neural network architectures which share similar structural properties [29; 10; 7]. The adjoint method for computing gradients has gained widespread use in the deep learning community, after it was shown in the Neural ODEs paper [11] that it is possible to parametrise the vector field defining an ODE by a neural network and differentiate through the flow to learn the vector field. This work has spawned a plethora of works that use learnable continuous dynamical systems.

#### Convex analysis and monotone operator theory

There is a recent line of work investigating the connections between existing deep learning practice and the topics of convex analysis and monotone operator theory. In particular, many of the standard activation functions that are used in neural networks are averaged (in the sense that we define in section 2.1), and further analysis enables one to use this insight to design neural networks that are averaged [12; 26; 18; 20].

### Our contributions

We describe and analyse a family of ResNet-styled neural network architectures that are guaranteed to be non-expansive. The effect of these neural networks on input vectors can be thought of as sequentially composing parts of (discretisations of) gradient flows along learnable convex potentials. We show that it is only necessary to control the operator norms of the learnable linear operators contained in these networks to ensure their non-expansiveness. This task is easily achieved in practice using power iteration.

The most basic such network takes the simple form described in algorithm 2. For this network, we use convex analysis techniques to show that more fine-grained control of the learnable linear operators ensures that each layer of the network is averaged, and as a result that the overall network is averaged.

We demonstrate the use of the proposed architectures by studying their natural application to an image denoising task, focusing on the influence of various tunable aspects in the architectures for this problem, and comparing our approach to a standard approach to the denoising task.
## 2 Methods

Suppose that \(f:\mathbf{R}\times\mathbf{R}^{n}\to\mathbf{R}^{n}\) is a time-dependent vector field and consider the ODE given by the flow along this vector field:

\[\dot{z}(t)=f(t,z(t)). \tag{2}\]

Assuming existence and uniqueness of the solutions to the ODE, we can define the flow map \(\Psi:[0,\infty)\ \times\mathbf{R}^{n}\to\mathbf{R}^{n}\) by \(\Psi(t,x)=z(t)\), where \(z\) solves eq. (2) with the initial condition \(z(0)=x\). Since the vector fields that we will consider are (globally) Lipschitz continuous, global existence and uniqueness is not an issue by the Picard-Lindelof theorem [36]. It is natural to ask when this flow map is non-expansive, in the sense that \(\|\Psi(t,x)-\Psi(t,y)\|\leqslant\|x-y\|\) for all \(t,x,y\) (with \(\|x\|:=\sqrt{\langle x,x\rangle}\)). Letting \(t\to 0\), we see that it is necessary that

\[\langle f(t,x)-f(t,y),x-y\rangle\leqslant 0,\]

and conversely, if this condition holds the flow map is non-expansive since

\[\frac{\mathrm{d}}{\mathrm{d}t}\|\Psi(t,x)-\Psi(t,y)\|^{2}=2\langle f(t,x(t))- f(t,y(t)),x(t)-y(t)\rangle\leqslant 0. \tag{3}\]

In practice, most ODEs of interest are not explicitly solvable and it is necessary to resort to numerical methods to approximate the flow map. A very well-studied class of such numerical integrators is the class of Runge-Kutta methods, which can be defined as follows:

**Definition 2.1** (Runge-Kutta method).: If \(m\in\mathbf{N}\) is a positive integer, an \(m\)-stage Runge-Kutta (RK) method is characterised by a matrix \(\mathcal{A}\in\mathbf{R}^{m\times m}\) and two vectors \(b,c\in\mathbf{R}^{m}\). For a step size \(h>0\), the RK method approximates the step from \(y=\Psi(t,x)\) to \(\Psi(t+h,x)\) as follows:

\[\Phi_{h}(t,y,f)=y+h\sum_{i=1}^{m}b_{i}f(t+c_{i}h,Y_{i}),\]

where \(Y=(Y_{1},\ldots,Y_{m})\) (the set of so-called stages of the method) solves the non-linear system of equations

\[Y_{i}=y+h\sum_{j=1}^{m}\mathcal{A}_{ij}f(t+c_{j}h,Y_{j})\quad\text{for}\quad i =1,\ldots,m.\]

If \(\mathcal{A}\) is strictly lower triangular, these equations are solvable in a single pass and the method is called explicit. Otherwise, the method is called an implicit method. To ensure that the method is at least of order one, we require that \(\sum_{j=1}^{m}b_{j}=1\). Furthermore, an RK method will usually satisfy \(c_{i}=\sum_{j=1}^{m}\mathcal{A}_{i,j}\). If \(c_{1},\ldots,c_{m}\) are distinct, the method is called non-confluent, whereas it is confluent otherwise.
Since we aim to design neural networks that encode non-expansive operators, it is of particular interest to know whether a given numerical integrator preserves the non-expansiveness of a continuous flow for which eq.3 holds. This property of a numerical integrator is called BN-stability and has been studied in detail for RK methods in [4]; for these methods, BN-stability is equivalent to algebraic stability, which is defined by a simple algebraic condition on the coefficients on the method. For methods that are non-confluent, these conditions are also equivalent to the condition of AN-stability (which is a priori a simpler condition). A comprehensive overview of stability properties for RK methods is given in [17; 14]. It is well known (see for instance [24]) however that no explicit RK method can satisfy such an unconditional stability condition. Nevertheless, it was shown in [13] that a conditional stability result can be established for certain explicit RK methods as long as eq.3 is replaced by an alternative that has the effect of controlling the stiffness of the ODE. To state this result, we require the definition of the circle contractivity property of an RK method.

**Definition 2.2** (Circle contractivity).: Suppose that \(\mathcal{A}\in\mathbf{R}^{m\times m}\) and \(b,c\in\mathbf{R}^{m}\) are the matrix and vectors characterising an RK method as in definition2.1. We say that this RK method satisfies the \(r\)-circle contractivity condition for a given \(r\in\mathbf{R}\cup\{\infty\}\) if \(|K(\zeta)|\leqslant 1\) for all \(\zeta\in D(r)^{m}\). Here, the function \(K:\mathbf{C}^{m}\to\mathbf{C}\) is defined 1 as:

Footnote 1: To understand the definition of \(K\), it can be thought of as the action of the RK method on a scalar nonautonomous linear ODE.

\[K(\zeta)=1+b^{\top}\operatorname{diag}(\zeta)(\operatorname{id}-\mathcal{A} \operatorname{diag}(\zeta))^{-1}\mathbf{1},\]

and \(D(r)\) is a generalised disk:

\[D(r)=\begin{cases}\{z\in\mathbf{C}||z+r|\leqslant r\}&\text{when}\quad r \geqslant 0,\\ \{z\in\mathbf{C}|\operatorname{Re}(z)\leqslant 0\}&\text{when}\quad r=\infty,\\ \{z\in\mathbf{C}||z+r|\geqslant-r\}&\text{when}\quad r<0.\end{cases}\]

**Example 2.1**.: _The most basic nontrivial example of an RK method is the forward (explicit) Euler method, given by \(\Phi_{h}(t,y,f)=y+hf(t,y)\). In the notation of definition2.1, we have \(m=1\), \(\mathcal{A}=0\) and \(b=1\), so \(K(z)=1+z\). We conclude that the forward Euler method is \(1\)-circle contractive._

**Remark 1**.: It is straightforward to compute the optimal (in the sense that it gives the largest generalised disk) \(r\) for which a given RK method is \(r\)-circle contractive if we know \(\mathcal{A}\) and \(b\): if we define the symmetric matrix \(Q=\operatorname{diag}(b)\mathcal{A}+\mathcal{A}^{\top}\operatorname{diag}(b )-bb^{\top}\), theorem3.1 from [13] tells us that \(r=-1/\rho\), where \(\rho\) is the largest number such that \(w^{\top}Qw\geqslant\rho w^{\top}\operatorname{diag}(b)w\) for all \(w\in\mathbf{R}^{m}\). Hence, if we can solve the generalised eigenvalue problem \(Qv=\lambda\operatorname{diag}(b)v\), we know that the minimal eigenvalue gives the desired \(\rho\).

With this definition, it is now possible to state the conditional stability result that extends to certain explicit methods:
**Theorem 2.3** (Theorem 4.1 from [13]).: _Suppose that \(\Phi_{h}\) is an RK method satisfying the \(r\)-circle contractivity condition, and that \(f\) satisfies the monotonicity condition_

\[\langle f(t,y)-f(t,z),y-z\rangle\leqslant-\nu\|f(t,y)-f(t,z)\|^{2}. \tag{4}\]

_Then, if \(r\neq\infty\) and \(h/r\leqslant 2\nu\), or if \(r=\infty\) and \(\nu\geqslant 0\),_

\[\|\Phi_{h}(t,y,f)-\Phi_{h}(t,z,f)\|\leqslant\|y-z\|.\]

The idea of using this result to design non-expansive neural networks was recently discussed in [5], though in this work no indication was given of how the vector fields should be parametrised. The monotonicity condition given by eq. (4) is reminiscent of the property of co-coercivity, known mainly from the theory of convex optimisation for its use in the Baillon-Haddad theorem:

**Theorem 2.4** (Corollary 18.16 from [3]).: _Suppose that \(\phi:\mathcal{X}\to\mathbf{R}\) is a Frechet-differentiable convex function on a Hilbert space \(\mathcal{X}\). Then \(\phi\) is \(L\)-smooth for some \(L\geqslant 0\) (equivalently, \(\nabla\phi\) is \(L\)-Lipschitz), meaning that_

\[\phi(y)\leqslant\phi(x)+\langle\nabla\phi(x),y-x\rangle+\frac{L}{2}\|y-x\|^{2},\]

_if and only if \(\nabla\phi\) is \(1/L\)-co-coercive, meaning that_

\[\langle\nabla\phi(y)-\nabla\phi(x),y-x\rangle\geqslant\frac{1}{L}\|\nabla \phi(y)-\nabla\phi(x)\|^{2}.\]

Indeed, if \(f(t,x)=-\nabla\phi(x)\) for a \(1/\nu\)-smooth convex potential \(\phi:\mathbf{R}^{n}\to\mathbf{R}\) (so that we have a gradient flow of a smooth convex potential), then eq. (4) is satisfied. This connection has recently been used to demonstrate in [31] that there is an explicit RK method for which the circle contractivity disk degenerates to a point, by constructing a smooth convex potential for which the non-expansiveness of the flow map is not preserved.

For the purpose of using this observation and theorem 2.3 to design non-expansive neural networks, note the following result:

**Lemma 2.5**.: _Suppose that \(\sigma:\mathbf{R}\to\mathbf{R}\) is a non-decreasing \(L\)-Lipschitz activation function, \(A\in\mathbf{R}^{n\times k}\) is a matrix and \(b\in\mathbf{R}^{n}\) is a bias vector. The vector field \(f_{A,b}(t,x)=-A^{\top}\sigma(Ax+b)\) (where \(\sigma\) is applied separately to each component) satisfies eq. (4) with \(\nu=1/(\|A\|^{2}L)\)._

Proof.: Since \(\sigma\) is non-decreasing and \(L\)-Lipschitz, the function \(\psi:\mathbf{R}\to\mathbf{R}\) given by \(\psi(t)=\int_{0}^{t}\sigma(s)\,\mathrm{d}s\) is convex and \(L\)-smooth. Hence, \(\phi:\mathbf{R}^{n}\to\mathbf{R}\) given by

\[\phi(x)=\sum_{i=1}^{n}\psi(x_{i})\]

is convex and \(L\)-smooth. The functional \(x\mapsto\phi(Ax+b)\) is convex and by the chain rule it has gradient equal to \(-f_{A,b}\) and it is \(\|A\|^{2}L\)-smooth. By the comments preceding this lemma, the vector field \(f_{A,b}\) satisfies eq. (4) with \(\nu=\|A\|^{2}L\)
By the previous observations, we can propose a natural non-expansive neural network architecture as follows: given an \(r>0\) such that we have an \(r\)-circle contractive RK method \(\Phi_{h}\) and an \(L\)-Lipschitz increasing activation function \(\sigma\), consider linear operators \(A^{1},\ldots,A^{\texttt{n\_blocks}}\), biases \(b^{1},\ldots,b^{\texttt{n\_blocks}}\) and stepsizes \(h^{1},\ldots,h^{\texttt{n\_blocks}}\) and define the operator \(\Xi\) by

\[\Xi=\Xi^{\texttt{n\_blocks}}\circ\ldots\circ\Xi^{1},\]

where \(\Xi^{i}(x)=\Phi_{h^{i}}(0,x,f_{A^{i},b^{i}})\) is one numerical integration step along the vector field \(f_{A^{i},b^{i}}\) as defined in lemma 2.5. Lemma 2.5 and theorem 2.3 ensure that \(\Xi\) is non-expansive as long as

\[h^{i}\|A^{i}\|^{2}L\leqslant 2r.\]

There are various ways in which this bound can be maintained during training, and the power method can be used to compute the required operator norm: as an example, it is possible to alternate gradient update steps of an optimiser with steps that scale the operators down to satisfy the bounds that are violated after the gradient update. An alternative method, which is more faithful to the dynamical systems interpretation, similarly keeps track of the operator norms of the weights but splits the interval into multiple smaller time steps to guarantee the bound when the gradient updates cause a violation of the bound. This approach, the adaptive approach, is the one that we will take in section 3. For any explicit RK method, the corresponding neural network \(\Xi\) is a residual network. For the forward Euler method, the network takes the particularly simple form shown in algorithm 2.

```
input: vector \(x\) parameters: step sizes \(h^{i}>0\), linear operators \(A^{1},\ldots,A^{\texttt{n\_blocks}}\) satisfying \(h\|A^{i}\|^{2}L\leqslant 1\) for \(i=1,\ldots,\texttt{n\_blocks}\), and biases \(b^{1},\ldots,b^{\texttt{n\_blocks}}\) \(z^{0}\gets x\) for\(i\gets 1,\ldots,\texttt{n\_blocks}\)do \(z^{i}\gets z^{i-1}-h^{i}(A^{i})^{\top}\sigma(A^{i}z^{i-1}+b^{i})\) endfor return\(\Xi(x)=z^{\texttt{n\_blocks}}\)
```

**Algorithm 2** Forward Euler method for non-expansive ODE networks

As mentioned before, we are focused in this chapter on explicit RK methods since they do not require the solution of a (potentially difficult) non-linear equation at each step. It may be interesting to note, however, what can happen when an implicit numerical method is used, such as the backward Euler method. In that case, each update step in algorithm 2 needs to be replaced by solving the equation

\[z^{i}=z^{i-1}-h^{i}f_{A^{i},b}(z^{i})=z^{i-1}-h^{i}(A^{i})^{\top}\sigma(A^{i}z ^{i}+b^{i}).\]

Recalling from the proof of lemma 2.5 that \(-f_{A,b}\) is the gradient of a convex functional \(\phi(A\cdot+b)\), this shows that the update step is given by

\[z^{i}=(\operatorname{id}+h^{i}\nabla\phi(A^{i}\cdot+b^{i}))^{-1}(z^{i-1})=: \operatorname{prox}_{h^{i}\phi(A^{i}\cdot+b^{i})}(z^{i-1}),\]
which is the defining equation of the proximal operator [23], a mathematical object that has been studied in great detail in the field of convex analysis. Whether considering it from the ODE viewpoint (the backward Euler method is BN-stable [4]) or from the convex analysis and monotone operator theory viewpoint (proximal operators are non-expansive as the resolvents of monotone operators [3, Chapter 23]), proximal operators \(\operatorname{prox}_{h\phi(A\cdot+b)}\) are well-defined and non-expansive regardless of the step size \(h>0\) and the smoothness of \(\phi(A\cdot+b)\). This unconditional stability comes at a cost, though: for general \(A\), computing the proximal operator of \(\operatorname{prox}_{h\phi(A\cdot+b)}\) is not easy (and becomes more difficult as the condition number of \(A\) increases). This issue can be overcome by restricting \(A\) to certain special sets of operators (for instance satisfying certain orthogonality properties), in which case the proximal operator may be explicitly computable. This approach is similar to the one taken in [18; 20] though note that it may be difficult to enforce these constraints on convolution-type linear operators. On the other hand, the operator norm constraints that we are required to enforce with explicit numerical integration methods can be easily controlled using power iteration [15]; all we need is the ability to apply the operator and its adjoint to test vectors.

### A more detailed look at the architecture for the forward Euler method

When the numerical integrator used is the forward Euler method, as described in algorithm 2, straightforward computations can be used to establish the same results guaranteed by the machinery of theorem 2.3, and some more nuanced results. Indeed, it is possible to choose the stepsizes in such a way that the resulting neural network is not just non-expansive, but in fact is also averaged:

**Definition 2.6** (Definition 4.23 from [3]).: Suppose that \(A:\mathcal{X}\to\mathcal{X}\) is an operator mapping a Hilbert space \(\mathcal{X}\) into itself and that \(\alpha\in(0,1)\). We call \(A\) an \(\alpha\)-averaged operator if there is a non-expansive \(T:\mathcal{X}\to\mathcal{X}\) such that \(A=(1-\alpha)\operatorname{id}+\alpha T\). We may also leave \(\alpha\) unspecified, in which case we just call \(A\) an averaged operator if there is an \(\alpha\in(0,1)\) such that \(A\) is \(\alpha\)-averaged.

Note that the triangle inequality shows that an averaged operator is non-expansive. In addition, averaged operators allow for convergent fixed point iterations, whereas ordinarily non-expansive operators enjoy no such guarantees. This is of crucial importance in certain applications, such as Plug-and-Play algorithms, where modelling denoisers using non-expansive operators is not enough to prevent divergence, but using averaged operators can ensure convergence [35; 20]. For our analysis here, let us note the following fact:

**Lemma 2.7**.: _Suppose that \(\Xi:\mathbf{R}^{n}\to\mathbf{R}^{n}\) is \(C^{1}\) with symmetric Jacobian \(D\Xi(x)\) everywhere and that \(\alpha\in(0,1)\). Then \(\Xi\) is \(\alpha\)-averaged if and only if_

\[\operatorname{spectrum}(D\Xi(x))\subset[1-2\alpha,1]\]

_for all \(x\in\mathbf{R}^{n}\). Note that the condition that the Jacobian is everywhere symmetric is equivalent to asking that \(\Xi=\nabla f\) for some underlying functional \(f:\mathbf{R}^{n}\to\mathbf{R}\)._
Proof.: Showing that \(\Xi\) is \(\alpha\)-averaged is equivalent to showing that \(\Theta=(\Xi-\operatorname{id})/\alpha+\operatorname{id}\) is non-expansive. Furthermore \(\operatorname{spectrum}(\Theta)=\operatorname{spectrum}(\Xi)/\alpha+1-1/\alpha\), so if \(\operatorname{spectrum}(D\Xi(x))\subset[1-2\alpha,1]\), then \(\operatorname{spectrum}(D\Theta(x))\subset[-1,1]\). By the fundamental theorem of calculus, we also have that

\[\Theta(x)-\Theta(y)=\Big{[}\int\limits_{0}^{1}D\Theta(y+t(x-y))\operatorname{ d}\!t\Big{]}(x-y),\]

so

\[\|\Theta(x)-\Theta(y)\|\leqslant\Big{[}\sup_{z\in\mathbf{R}^{n}}\|D\Theta(z) \|\Big{]}\|x-y\|\leqslant\|x-y\|,\]

as desired. Conversely, suppose that we do not have that \(\operatorname{spectrum}(D(\Xi))\subset[1-2\alpha,1]\) everywhere. In particular, there is some \(x\in\mathbf{R}^{n}\), such that there is an eigenvector \(v\in\mathbf{R}^{n}\) (assume that \(\|v\|=1\)) of \(D\Theta(x)\) with eigenvalue outside of \([-1,1]\): \(D\Theta(x)v=\lambda v\) with \(|\lambda|>1\). Hence, we have

\[\Theta(x+hv)-\Theta(x)=hD\Theta(x)v+R(h)=\lambda hv+R(h)\]

with \(R(h)\) a remainder term that satisfies \(R(h)=o(h)\) as \(h\to 0\). By the triangle inequality, this implies that

\[\|\Theta(x+hv)-\Theta(x)\|\geqslant|\lambda|\|hv\|-|R(h)|.\]

For small enough \(h\), this inequality tells us that \(\Theta\) is not non-expansive, and by the preceding argument, \(\Xi\) is not \(\alpha\)-averaged.

Recall that a single layer of the proposed architecture is given by \(\Xi(x)=x-hA^{\top}\sigma(Ax+b)\), with the same setting in mind as described in lemma 2.5. There we saw that \(\Xi\) is the gradient of the functional \(x\mapsto\|x\|^{2}/2-h\phi(Ax+b)\), where \(\phi\) is convex and \(L\)-smooth, so that \(\operatorname{spectrum}(D^{2}\phi(x))\subset[0,L]\) for each \(x\in\mathbf{R}^{n}\). Hence, since we have \(D\Xi(x)=x-hA^{\top}D^{2}\phi(Ax+b)A\), we find that

\[\operatorname{spectrum}(D\Xi(x))\subset[1-h\|A\|^{2}L,1].\]

Combining this with lemma 2.7 immediately gives the following result if the activation function \(\sigma\) is \(C^{1}\). This is not required, however, for the result to be valid; any \(L\)-Lipschitz \(\sigma\), such as \(\sigma(x)=\mathtt{ReLU}(x)=\max\{0,x\}\), will work equally well. The argument for general \(L\)-Lipschitz \(\sigma\) is given below:

**Theorem 2.8**.: _Let \(\sigma,A,b\) be as in lemma 2.5 and let \(\alpha\in(0,1)\). A single layer of the proposed architecture, \(\Xi(x)=x-hA^{\top}\sigma(Ax+b)\), is \(\alpha\)-averaged if_

\[h\|A\|^{2}\leqslant 2\alpha/L. \tag{5}\]

Proof.: The argument given above provides some intuition regarding averaged operators, but requires the activation function \(\sigma\) to be \(C^{1}\). Here, we will show that this is not necessary. Indeed, note that an operator \(\Xi\) is \(\alpha\)-averaged if and
only if \((\Xi-\mathrm{id})/\alpha+\mathrm{id}\) is non-expansive. Furthermore, we note that \(\Xi-\mathrm{id}=hf_{A,b}\), where \(-f_{A,b}\) is the gradient of an \(\|A\|^{2}L\)-smooth convex functional, as defined in the proof of lemma 2.5. By theorem 2.4 we have that

\[\frac{1}{\|A\|^{2}L}\|f_{A,b}(x)-f_{A,b}(y)\|^{2}\leqslant\langle-f_{A,b}(x)+f _{A,b}(y),x-y\rangle\leqslant\|A\|^{2}L\|x-y\|^{2},\]

so, if we write \(\Lambda=(\Xi-\mathrm{id})(x)-(\Xi-\mathrm{id})(y)\) to reduce clutter, we have

\[-h\|A\|^{2}L\|x-y\|^{2}\leqslant\langle\Lambda,x-y\rangle\leqslant-\frac{1}{h \|A\|^{2}L}\|\Lambda\|^{2}.\]

In particular, \(\langle\Lambda,x-y\rangle+\|\Lambda\|^{2}/(h\|A\|^{2}L)\leqslant 0\). Upon expanding the squared norm, we find that

\[\|((\Xi-\mathrm{id})/\alpha+\mathrm{id})(x)- ((\Xi-\mathrm{id})/\alpha+\mathrm{id})(y)\|^{2}\] \[=\frac{\|\Lambda\|^{2}}{\alpha^{2}}+2\frac{\langle\Lambda,x-y \rangle}{\alpha}+\|x-y\|^{2}\] \[\leqslant\|x-y\|^{2}+\frac{2}{\alpha}\Big{(}\langle\Lambda,x-y \rangle+\frac{\|\Lambda\|^{2}}{2\alpha}\Big{)}.\]

By the above comments, we see that \((\Xi-\mathrm{id})/\alpha+\mathrm{id}\) is non-expansive when \(2\alpha\geqslant h\|A\|^{2}L\), which can be rewritten into \(h\|A\|^{2}\leqslant 2\alpha/L\). 

Finally, the following result guarantees that the overall network will be averaged as long as each layer is averaged, with a corresponding \(\alpha\) that can be controlled:

**Theorem 2.9** (Proposition 4.32 from [3]).: _Suppose that \(\Xi^{1},\ldots,\Xi^{m}\) are operators \(\Xi^{i}:\mathcal{X}\to\mathcal{X}\) on a Hilbert space \(\mathcal{X}\) and that each \(\Xi^{i}\) is \(\alpha_{i}\)-averaged for some \(\alpha_{i}\in(0,1)\). Then \(\Xi^{m}\circ\ldots\circ\Xi^{1}\) is \(\alpha\)-averaged, where_

\[\alpha=\frac{m}{m-1+\min_{i=1,\ldots,m}(1/\alpha_{i})}.\]

In particular, if we are targeting a certain \(\alpha\in(0,1)\) for which our neural network (\(m\) layers deep) should be \(\alpha\)-averaged, we should ask that each layer is \(\alpha_{i}\)-averaged with \(\alpha_{i}\) at most

\[\alpha_{i}\leqslant\frac{\alpha}{m(1-\alpha)+\alpha}.\]

By theorem 2.8, we see that this implies that we must use a step size \(h=\mathcal{O}(1/m)\) that decreases to \(0\) as the depth \(m\) of the network increases. Alternatively, it is possible to get an averaged operator by appealing to definition 2.6: \((1-\alpha)\,\mathrm{id}+\alpha\Xi\) will be \(\alpha\)-averaged as long as \(\Xi\) is non-expansive, which we have seen can be guaranteed with a step size independent of the depth of the network.
## 3 Experiments

We will study the application of the proposed architectures to tasks where a certain kind of robustness is required: in section 3.2 we study the robustness of image classifiers to adversarial attacks and in section 3.3 we consider the robustness of learned image denoisers. Robustness in the latter setting is of particular importance in downstream tasks, such as when the learned denoisers are used to solve inverse problems, as we will see in section 3.4. First, though, we will describe the methods that we use to train the networks.

### Training methods

We train each network in a supervised manner, by solving a regularised empirical risk minimisation problem. In all experiments, we perform 40,000 iterations of stochastic gradient descent (SGD) with momentum (with momentum parameter \(\beta=0.9\)), with a piecewise linear learning rate schedule that ramps up from a minimum learning rate to a maximum learning rate for the first half of the iterations and then down to 0 for the second half of the iterations. The minimum and maximum learning rates to use are found with the method described in [32]. Every experiment is run with weight decay, with weighting hyperparameters \(\lambda\in\{10^{-5},5\cdot 10^{-5},10^{-4},5\cdot 10^{-4}\}\). The reported results correspond to the setting of the weight decay hyperparameter that performed best in terms of the loss function evaluated on a held-out validation set.

As mentioned before, to guarantee non-expansiveness or averagedness of the network we need to ensure that a bound such as eq.5 in theorem2.8 holds, so we use the power method [22] to compute spectral norms of each of the learnable linear operators: if \(A:\mathbf{R}^{n}\to\mathbf{R}^{m}\) is a linear operator and we have an initial estimate of the first left singular vector \(v^{0}\in\mathbf{R}^{m}\), we iterate

\[u^{k}\leftarrow\frac{A^{\top}v^{k-1}}{\|A^{\top}v^{k-1}\|},\qquad v^{k} \leftarrow\frac{Au^{k}}{\|Au^{k}\|}.\]

Assuming that \(v^{0}\) is not orthogonal to the first left singular vector (this is guaranteed to hold with probability 1 if \(v^{0}\) is randomly selected from a probability distribution that has a density w.r.t. Lebesgue measure), \(u^{k}\) and \(v^{k}\) converge to the first singular vectors of \(A\) as \(k\to\infty\) and \((u^{k})^{\top}Av^{k}\to\|A\|\). In addition, it is reasonable to assume that the weights only undergo incremental updates during each training iteration, so that we can warm start the power method using the estimates of the singular vectors from the previous training iteration, followed by just a single iteration of the power method.

The most straightforward way to satisfy the bounds that we require is by alternating gradient update steps with spectral normalisation steps \(A\mapsto A/\|A\|\). In a similar spirit, it is in principle also possible to differentiate through the spectral normalisation step [22], which results in different training dynamics and may avoid the tendency of the alternating steps to counteract each other's effect.

Instead of these methods, we find that another approach (which we will refer to as the adaptive approach in what follows) similarly ensures that the
bounds of theorem 2.8 hold, with the benefits of empirically resulting in high expressive power, a simple implementation (compared to differentiating through the spectral normalisation) and a close connection to the underlying dynamical system. In the adaptive approach, we still use the power method to keep track of the operator norms, but when the bound in eq. (5) is violated, we split the integration interval into sufficiently many subintervals (each of equal size) to ensure that the bound is satisfied on each subinterval. Let us show an example of what this means for the approach shown in algorithm 2: after updating the weights of each block using SGD and updating the spectral norm estimate of the weights using the power method, the forward propagation shown in algorithm 2 is modified. In a given block with weights \(A^{i}\) and overall step size \(h^{i}\), rather than take a single step \(z^{i}\gets z^{i-1}-h^{i}(A^{i})^{\top}\sigma(Az^{i-1}+b)\), we compute a number of steps \(N=\lceil h^{i}\|A^{i}\|^{2}L/(2r)\rceil\) to subdivide the integration into and take \(N\) steps with stepsize \(h^{i}/N\). In fact, since the ceiling defining \(N\) is not generally attained (this can be checked after training to be certain), each block will naturally satisfy the constraint in theorem 2.8, making the network averaged by theorem 2.9.

As a result of the adaptive approach, the depths of our networks are not fixed during training; although the number of learnable parameters is fixed, the network will become deeper during training if the weights grow. We do not take any additional measures to ensure that the weights stay bounded, but we find that this is not necessary in practice. Furthermore, since the loss functions that we consider are bounded from below and the added weight decay penalty is coercive, we are guaranteed that any minimising sequence of the regularised loss function is bounded. Hence, we do not expect the depths of our networks to blow up during training.

To initialise the weights of the non-expansive networks that we study, we first use the default initialisation method in PyTorch [25] to initialise the convolutional filters and biases, and apply 1,000 iterations of the power method to compute the operator norms of the convolution operations. The filters are then normalised to have operator norm 1 and the singular vectors output by the power method are saved for future iterations.

In all of the networks that we study in the experiments, we choose the activation function \(\sigma\) to be the LeakyReLU, defined as LeakyReLU\((x)=\max\{x,0.01x\}\). Evidently, LeakyReLU is 1-Lipschitz and corresponds to the gradient of a strongly convex functional. For each ODE block in the non-expansive networks, we fix the overall timestep to be the circle contractivity radius \(r\) of the integrator used. Combining the weight initialisation mentioned above and the Lipschitz constant of LeakyReLU, this ensures that the required constraints are satisfied at initialisation. As previously described, the constraints are maintained during training using the adaptive approach.

All experiments have been implemented using PyTorch [25] and were run on a single computational node with a NVIDIA A100-SXM4 GPU with 80 GB of memory. The majority of training runs were finished in fewer than 4 hours (the exceptions being the image denoising experiment with higher order integrators in section 3.3, which took up to 12 hours to complete). The code that we have written to implement the methods and experiments is publicly available
at [https://github.com/fsherry/non-expansive-odes](https://github.com/fsherry/non-expansive-odes).

### Adversarial robustness

Let us study the robustness of the proposed architecture as opposed to a comparable (but unconstrained) classifier based on a ResNet architecture. We will use the CIFAR-10 dataset [21], which consists of 60,000 colour images of size \(32\times 32\), with the standard split into 50,000 training images and 10,000 testing images. Additionally, we will hold out 2,500 training images to use as a validation set. Each image comes with one of 10 possible labels, which we will simply refer to as numbers in \(\{1,\ldots,10\}\). We scale the images so that each channel only takes values in \([0,1]\), but do not perform any additional normalisation or whitening. The neural networks that we consider implement functions between \(\mathbf{R}^{32\times 32\times 3}\) and \(\mathbf{R}^{10}\), the outputs of which are interpreted as the scores assigned to each label. Finally, given such a function \(\Phi:\mathbf{R}^{32\times 32\times 3}\rightarrow\mathbf{R}^{10}\) its classification of an input image \(x\) is given by \(\operatorname*{argmax}_{i\in\{1,\ldots,10\}}\Phi(x)_{i}\).

Oftentimes, the scores output by a classifier are interpreted probabilistically: applying the softmax function, we transform the outputs of the network into probabilities, which can be fit to the true labels by minimising a cross-entropy loss function. In our setting, however, we will not take this approach, instead opting for a loss function that more directly encourages the classifier to have a large margin. Recall that the margin of a classifier at an input \(x\) with true label \(y\) is given by \(\Phi(x)_{y}-\operatorname*{argmax}_{i\in\{1,\ldots,10\}\setminus\{y\}}\Phi(x )_{i}\). With this definition, images that are incorrectly classified correspond to a negative margin, while those that are correctly classified correspond to a positive margin. If we can lower bound the margin at a correctly classfied image, while upper bounding the Lipschitz constant of \(\Phi\) we can certify that \(\Phi\) equally classifies all images in a ball around this image [37]. In particular, maximising the margin of a classifier makes sense when we constrain its Lipschitz constant. In this context, it is natural to use a multi-class hinge loss function:

\[\min_{\Psi}\frac{1}{N_{\text{train}}}\sum_{i=1}^{N_{\text{train}}}\sum_{k=1}^{ 10}\max\{0,\mu-\Phi(x_{i})_{y_{i}}-\Phi(x_{i})_{k}\}, \tag{6}\]

with \(\mu\) a hyperparameter that we fix to be equal to 0.1 in our experiments. Furthermore, we will use data augmentation, effectively replacing the inner term in eq. (6) by an expectation over the augmentations, and \(x_{i}\) by \(\Delta(x_{i})\) with \(\Delta\) the augmentation. We compose an augmentation that randomly erases pixels, an augmentation that randomly crops out parts of the images and an augmentation that applies a random horizontal flip. For the training of all networks we use minibatches of size 128 in each SGD update.

We will consider two comparable network architectures: an architecture we will refer to as ODENet which uses our proposed nonexpansive components, and an architecture we will call ResNet which is similar but unconstrained. In particular, by appropriately setting the weights of the ResNet, we can recover
an ODENet. Both classifiers take the form

\[\Phi=A_{\text{lin}}\circ\text{pool}_{\text{global}}\circ\overset{3}{\underset{i=1 }{\cap}}\Big{[}\Xi^{i}\circ\text{pool}_{2\times 2}\circ A_{\text{lift},i}\Big{]}.\]

Here \(A_{\text{lin}}\) is a final linear layer, \(A_{\text{lift},i}\) are \(1\times 1\) convolutions increasing the number of channels, \(\text{pool}_{2\times 2}\) are average pooling operations with \(2\times 2\) kernels and a stride of 2 and \(\text{pool}_{\text{global}}\) is a global average pooling operation (which collapses the spatial extent of the remaining image into a single pixel). In the case of the ResNet, each \(\Xi^{i}\) is a composition of 5 simple ResNet blocks of the form \(x\mapsto B\sigma(Ax+b)\), whereas in the case of the ODENet, each \(\Xi^{i}\) takes the form shown in algorithm 2 with 5 blocks. With the particular configuration used in our experiments, the ODENet has 254,810 trainable parameters, while the ResNet has 497,290 trainable parameters. Recall from section 3.1 that we use an adaptive training approach for ODENet, meaning that its depth is generally allowed to grow during training, but its number of trainable parameters remains constant throughout training. This is done to enforce the constraints that we require to ensure nonexpansiveness. On the other hand, since we do not enforce constraints on the ResNet, its depth is constant during training.

To implement the adversarial attacks, we use the Foolbox package [27], applying a projected gradient descent attack (which in fact alternates gradient _ascent_ steps on the loss function with projections onto a ball around the clean image). More specifically, we consider an attack that uses an \(\ell^{2}\)-norm constraint and perform 10 iterations of the projected gradient descent method.

Figure 1: A comparison of the adversarial robustness of the proposed ODENet and a similar ResNet trained to classify CIFAR-10 images. The adversarial attack takes the form of a projected gradient descent attack with 10 iterations, with the accuracy curves computed on the test set.
The result of applying such an attack to the networks described above is shown in fig. 1: we compute attacks of various sizes \(\varepsilon\) (corresponding to the size of the \(\ell^{2}\)-ball within which we search for an attack) for each image in the test set and compute how well both networks continue to perform on the attacked images. Clearly, the unconstrained ResNet has higher clean accuracy (accuracy at \(\varepsilon=0\)), namely \(90.2\%\), compared to the ODENet, which has a clean accuracy of \(84.2\%\). On the other hand, the ODENet is considerably more robust than the ResNet as the perturbation size grows: considering the area under the curve as a measure of overall robustness, (which is bounded between \(0\) and \(1\)), we find that the ResNet has an area under the curve of \(0.537\), while the ODENet has an area under the curve of \(0.641\).

### Nonexpansive neural networks for denoising

Let us now proceed by studying the performance of the proposed networks on an image denoising task, comparing to a standard variational denoising approach and an unconstrained deep learning approach. We will use the BSDS500 dataset [1] as training data and test data for these experiments. This dataset consists of \(500\) natural colour images of size \(321\times 481\), split into \(N_{\text{train}}=200\) training images, \(N_{\text{val}}=100\) validation images and \(N_{\text{test}}=200\) test images. In our experiments, we adhere to the same splitting of the dataset. We scale the images so that each channel only contains values in \([0,1]\) and simulate noisy images \(y\) corresponding to each ground truth image \(x\) by adding Gaussian white noise \(\varepsilon\) with a standard deviation of \(0.15\), corresponding to a high noise level (see figs. 3 and 4). For all of the learned approaches, we train the networks \(\Gamma\) involved to solve the empirical risk minimisation problem

\[\min_{\Gamma}\frac{1}{N_{\text{train}}}\sum_{i=1}^{N}\|\Gamma(y_{i})-x_{i}\|^{ 2}, \tag{7}\]

and when we test the performance of the various denoisers on the test set, we essentially use the same loss: we report the peak signal-to-noise ratio (PSNR), defined between a reference \(x^{*}\) and an estimate \(\hat{x}\) as

\[\text{PSNR}(\hat{x},x^{*})=10\log_{10}\Big{(}\frac{\max_{i,j,k}|x^{*}_{i,j,k}| ^{2}}{\frac{1}{3\cdot 321\cdot 481}\sum_{i,j,k}|x^{*}_{i,j,k}-\hat{x}_{i,j,k}| ^{2}}\Big{)}.\]

We use the training approach described in section 3.1 and consistently use mini-batches of size \(5\).

The networks that we propose for this task are of the form

\[\Gamma=A_{\text{project}}\circ\Xi\circ A_{\text{lift}}, \tag{8}\]

where \(A_{\text{lift}}\) is a lifting operation taking the \(3\) input channels to \(64\) channels by appending \(61\) channels filled with zeros, \(A_{\text{project}}\) is a projection operator taking \(64\) channels to the \(3\) output channels by simply dropping the last \(61\) channels, and \(\Xi\) is a network as in algorithm 2 with each \(A^{i}\) a convolution taking \(64\)
channels to 64 channels and each \(b^{i}\in\mathbf{R}^{64}\) (i.e. the biases are spatially constant, as usual). All convolution operators have kernel size \(3\times 3\). This choice of lifting and projection operators ensures that \(\alpha\)-averagedness of \(\Xi\) implies \(\alpha\)-averagedness of \(\Gamma\): we have \(A_{\text{project}}\circ A_{\text{lift}}=\operatorname{id}\) and \(\|A_{\text{lift}}\|=\|A_{\text{project}}\|=1\) so if \(\Xi=(1-\alpha)\operatorname{id}+\alpha T\) is \(\alpha\)-averaged with \(T\) non-expansive, we also have

\[\Gamma =A_{\text{project}}\circ((1-\alpha)\operatorname{id}+\alpha T) \circ A_{\text{lift}}\] \[=(1-\alpha)A_{\text{project}}\circ A_{\text{lift}}+\alpha A_{ \text{project}}\circ T\circ A_{\text{lift}}\] \[=(1-\alpha)\operatorname{id}+\alpha A_{\text{project}}\circ T \circ A_{\text{lift}}.\]

Since \(A_{\text{project}}\) and \(A_{\text{lift}}\) are non-expansive, this shows that \(\Gamma\) is \(\alpha\)-averaged. We will instantiate this architecture using the forward Euler integrator as in algorithm 2 (with n_blocks = 10) and will simply refer to this network by the name Euler in what follows.

Although the concrete architecture obtained when using the forward Euler integrator (as described in algorithm 2) is appealing in its simplicity, the framework laid out in section 2 also allows us to use certain higher order integrators. For instance, consider Heun's method, which is given by

\[\Phi_{h}^{\text{Heun}}(t,y,f)=y+\frac{h}{2}\Big{(}f(t,y)+f(t+h,y+hf(t,y))\Big{)}. \tag{9}\]

This is a 2-stage, second-order, RK method, with

\[\mathcal{A}=\begin{pmatrix}0&0\\ 1&0\end{pmatrix},\quad b=\begin{pmatrix}1/2\\ 1/2\end{pmatrix},\quad\operatorname{diag}(b)\mathcal{A}+\mathcal{A}^{\top} \operatorname{diag}(b)-bb^{\top}=\frac{1}{4}\begin{pmatrix}1&-1\\ -1&1\end{pmatrix},\]

and using remark 1, we conclude that Heun's method is 1-circle contractive, like the forward Euler method is. As a result, \(x\mapsto\Phi_{h}^{\text{Heun}}(0,x,f_{A,b})\) is non-expansive as long as \(h\|A\|^{2}\leq 2\) and algorithm 2 can be adapted to use Heun's method, the only change being that the steps \(z^{i}\gets z^{i-1}-h(A^{i})^{\top}\sigma(A^{i}z^{i-1}+b^{i})\) are replaced by steps of the form \(z^{i}\leftarrow\Phi_{h}^{\text{Heun}}(0,z^{i-1},f_{A^{i},b^{i}})\). We will also instantiate the architecture in eq. (8) with this integrator (again using n_blocks = 10) and will refer to this network by the name Heun. As this integrator takes more evaluations of the right hand side of the ODE, it is more costly to compute the output of Heun than of Euler.

Similarly, we can consider integrators with yet higher orders, such as the fourth-order RK4 integrator \(\Phi_{h}^{\text{RK4}}\), given by definition 2.1 with

\[\mathcal{A} =\begin{pmatrix}0&0&0&0\\ 1/2&0&0&0\\ 0&1/2&0&0\\ 0&0&1&0\end{pmatrix}, b =\begin{pmatrix}1/6\\ 1/3\\ 1/3\\ 1/6\end{pmatrix},\] \[\operatorname{diag}(b)\mathcal{A}+\mathcal{A}^{\top} \operatorname{diag}(b)-bb^{\top} =\frac{1}{36}\begin{pmatrix}-1&4&-2&-1\\ 4&-4&2&-2\\ -2&2&-4&4\\ -1&-2&4&-1\end{pmatrix}.\]
Again using remark 1, we conclude that the RK4 method is 1-circle contractive and we can replace the forward Euler method in algorithm 2 by the RK4 method to obtain a non-expansive neural network. We design a denoiser (with n_blocks = 10) as in eq. (8) with the RK4 integrator and refer to it by the name RK4. RK4 is again more costly to to use since the RK4 integrator takes 4 evaluations of the right hand side of the ODE for each step.

All of our proposed networks have the same number, 369,280, of trainable parameters.

As a benchmark denoising algorithm using the variational approach, we can consider total variation (TV) denoising [28], which gives the denoised image as

\[\hat{u}=\operatorname*{argmin}_{u}\frac{1}{2}\|u-y\|^{2}+\alpha\|\nabla u\|_{1}, \tag{10}\]

where we have tuned \(\alpha\) for optimal reconstruction performance on the training set. This approach solves a convex optimisation problem with a hand-crafted regularisation functional that favours reconstruction of piecewise constant images. Solving this problem accurately is more expensive than the deep learning approaches that we will consider and generally its hand-crafted prior information is not matched as well to reality as the prior information encoded in the learned approaches. On the other hand, this approach can be considered trustworthy for use in downstream tasks: eq. (10) defines the reconstruction map as a proximal operator of a convex functional, so that the reconstruction map is 1/2-averaged.

The DnCNN, introduced in [40], has become a standard benchmark for denoising tasks. A natural comparison to make is between our networks \(\Gamma\) with n_blocks = 10, and the DnCNN \(\Gamma_{\text{DnCNN}}=\operatorname{id}-A_{\text{project}}\circ\Xi_{\text{ DnCNN}}\circ A_{\text{lift}}\) where \(\Xi_{\text{DnCNN}}\) is a 18-layer convolutional neural network without skip connections, the details of which are described in [40]. \(A_{\text{lift}}\) and \(A_{\text{project}}\) are taken as convolutions that lift the channels to 64 channels from the 3 input channels and project the 64 channels back to the 3 output channels respectively. Note that this is a reasonable version of the DnCNN to compare with our networks since each block of our architecture contains a convolution and its transpose, whereas the DnCNN uses one convolution per layer. This DnCNN is trained to solve eq. (7), in the same way as our architectures except that no operator norm constraints are enforced on the convolutions. The DnCNN that we consider has 670,531 trainable parameters.

The results of comparing all of these approaches are shown in fig. 2: all of the learned approaches perform basically on par with each other, slightly outperforming the variational approach. The gap between the learned approaches and the variational approach is not as large as one may expect, but this is a result of the high noise level considered. Figure 3 and fig. 4 show zoomed in examples of test images that are relatively more favourable to the learned approaches and to TV and vice versa. It is remarkable that Euler, Heun and RK4 perform extremely similarly to each other: although the higher order integrators are more costly to evaluate, and attain higher order of approximation of the underlying trajectories, this is not borne out in better denoising performance. This can be
explained by the fact that (at least in this task) we are only interested in the endpoints of the underlying trajectories.

Another point that is worth expanding on is the fact that our proposed networks perform on par with the unconstrained DnCNN. This is perhaps somewhat unexpected: recall from the discussion in section 2 that our approach bears some similarity to the Parseval proximal networks described in [18; 20]. Their approach effectively uses an implicit Euler integrator and restricts the weights to be orthogonal to make the computation of the implicit step computationally feasible, and they observed that it was necessary to scale up their proposed networks by a multiplicative factor and apply residual learning (where the network models the noise instead of the image directly) to obtain networks that perform similarly to DnCNN. These modifications break the desirable non-expansiveness properties of the network, necessitating tricks such as using a so-called "oracle denoiser" to obtain the desired properties. On the other hand, our experiments in this section show that these tricks are not necessary: it is possible to enforce stability constraints while having high denoising performance.

### Plug-and-Play applications of the learned denoisers

Recall from section 1.2 that one of the downstream applications where provably stable neural networks are of great interest is in the Plug-and-Play (PnP) approach to solving ill-posed inverse problems using learned denoisers. With the results from the previous section and the theoretical results of section 2.1, we are now in a position to pursue this further. Indeed, recalling theorem 2.8 and theorem 2.9, we get that the denoiser trained using the Euler integrator in the previous section is an averaged operator, since the adaptive training approach ensures that the condition in theorem 2.8 is satisfied (for an unspecified

Figure 2: Comparison of the denoising performance of the various considered approaches on the test set. The learned approaches all perform similarly and outperform the variational approach that uses TV regularisation.
\(\alpha>0\)) at all times. In particular, it has convergent fixed-point iterations (3, Theorem 5.14). Let us see what this means in practice: fig. 5 shows what happens if we repeatedly apply the learned DnCNN and Euler denoisers to an input image. We clip the values of all images to lie in \([0,1]\) for display purposes, but it is evident from the figure that repeatedly applying DnCNN results in divergence, whereas repeatedly apply Euler converges. This convergence property is a desirable property for a denoiser, and is for instance satisfied by any proximal operator of a convex regularisation functional.

Let us consider the inverse problem of deblurring: we assume that we are given measurements \(y=Kx+\varepsilon\), where \(Kx=k*x\) is a convolution operation representing a motion blur, \(x\) is the ground truth image that we aim to recover as well as possible and \(\varepsilon\) is Gaussian white noise corrupting the measurements. The ill-posedness of this problem is manifested in the instability of the inverse of the convolution; as a consequence of this, a naive inversion of the measurements will blow up the noise in the measurements. The PnP approach uses the previously learned denoiser to regularise this inversion. In particular, we will use the simple PnP proximal gradient method of algorithm 1, with \(E_{y}(x)=\|y-Kx\|^{2}/2\) and the denoiser \(\Phi\) equal to the previously learned Euler denoiser. We choose the step size \(\tau=1/\|K\|^{2}\), which together with the fact that the denoiser is averaged gives a convergence guarantee (35, Proposition 2). The asserted convergence is practically observed in fig. 6: the result \(\hat{x}\) obtained as the limit of the PnP iterations trades off the prior information about natural images encoded in the

Figure 3: Comparison of some of the denoisers considered, on an image that is favourable to the learned approaches. Note that the learned approaches recover more fine details, whereas TV has a tendency to flatten them out. The numbers in the top right corner of each image are the PSNRs (in dB) relative to the ground truth \(x\).
Euler denoiser with the data consistency.

## 4 Conclusions and discussion

We have exhibited a family of ResNet architectures for which it is straightforward to enforce non-expansiveness. The proposed architecture is given by compositions of numerical integration steps along gradient flows in convex potentials. For the main example using the forward Euler method, we have used

Figure 4: Comparison of some of the denoisers considered, on an image that is favourable to TV. The true image to be recovered is relatively well-approximated by a piecewise constant image. Note that the finer details in the image (for example the ridge between the top and bottom part of the beak) are still better recovered by the learned approaches. The numbers in the top right corner of each image are the PSNRs (in dB) relative to the ground truth \(x\).

Figure 5: Repeated application of the unconstrained and constrained denoisers to a given input image gives drastically different results: for the unconstrained DnCNN, this sequence diverges, whereas for the averaged Euler denoiser, this sequence converges.
tools from convex analysis to show that the architecture can be used to encode averaged operators. We have demonstrated the use of the proposed architectures on adversarially robust image classification and on image denoising, with the idea of applying the learned denoisers to ill-posed inverse problems.

With a novel adaptive training approach, we have shown that it is possible to obtain performant neural networks, even when enforcing desirable stability constraints. Although the basic architecture uses the first order forward Euler method as the numerical integrator, it is possible to use higher order methods. In the tasks considered in this work, using higher order integrators did not come with any benefits: the performance of the learned networks was not improved, while the computational cost was significantly increased. It remains a question of interest whether there are tasks where the higher order integrators provide real benefits over the forward Euler integrator.

Future work could also go towards studying the application of these architectures in typical deep learning applications such as GANs, specifically in Wasserstein GANs which require the use of a 1-Lipschitz critic function. We have seen in practice that the proposed architectures are quite expressive, but an interesting direction for future work would also be to study ways in which more general learnable non-expansive flows can be used to motivate the design of provably stable neural network architectures, and provide approximation guarantees for them. Finally, the framework shown in this paper essentially depends

Figure 6: Using the learned Euler denoiser to solve an ill-posed inverse problem (deblurring) in a PnP fashion, with convergence guarantee. The numbers in the top right corner of each image are the PSNRs (in dB) relative to the ground truth \(x\).
on the use of the Euclidean norm, but depending on the application one may want to design neural networks based on the dynamical systems connection that are 1-Lipschitz with respect to different norms.

## Acknowledgments

EC & BO have received support from the European Union Horizon 2020 research and innovation programme under the Marie Sklodowska-Curie grant agreement No. 860124. MJE acknowledges support from EPSRC (EP/S026045/ 1, EP/T026693/1, EP/V026259/1) and the Leverhulme Trust (ECF-2019-478). DM was partially supported by a grant from the Simons Foundation. CBS acknowledges support from the Philip Leverhulme Prize, the Royal Society Wolfson Fellowship, the EPSRC advanced career fellowship EP/V029428/1, EPSRC grants EP/S026045/1 and EP/T003553/1, EP/N014588/1, EP/T017961/1, the Wellcome Innovator Awards 215733/Z/19/Z and 221633/Z/20/Z, the European Union Horizon 2020 research and innovation programme under the Marie Sklodowska-Curie grant agreement No. 777826 NoMADS, the Cantab Capital Institute for the Mathematics of Information and the Alan Turing Institute. FS acknowledges support from the EPSRC advanced career fellowship EP/V029428/1.
