Robust Learning with Implicit Residual Networks

Abstract

In this effort, we propose a new deep architecture utilizing residual blocks inspired by implicit discretization schemes. As opposed to the standard feed-forward networks, the outputs of the proposed implicit residual blocks are defined as the fixed points of the appropriately chosen nonlinear transformations. We show that this choice leads to the improved stability of both forward and backward propagations, has a favorable impact on the generalization power, and allows for control the robustness of the network with only a few hyperparameters. In addition, the proposed reformulation of ResNet does not introduce new parameters and can potentially lead to a reduction in the number of required layers due to improved forward stability. Finally, we derive the memory-efficient training algorithm, propose a stochastic regularization technique, and provide numerical results in support of our findings.

1. Introduction and Related Works

A large volume of empirical results has been collected in recent years illustrating the striking success of deep neural networks (DNNs) in approximating complicated maps by a mere composition of relatively simple functions [1]. Universal approximation property of DNNs with a relatively small number of parameters has also been shown for a large class of functions [2,3]. The training of deep networks nevertheless remains a notoriously difficult task due to the issues of exploding and vanishing gradients, which become more apparent and noticeable with increasing depth [4]. These issues accelerated efforts of the research community in an attempt to explain this behavior and gain new insights into the design of better architectures and faster algorithms. A promising approach in this direction was obtained by casting evolution of the hidden states $y_t\in Y_t$ of a DNN as a dynamical system [5], i.e.,

$$\begin{array}{cc}\hfill y_{t+1}& ={\mathsf{\Phi }}_t({\gamma }_t,y_t),t=0,...,T-1,\hfill \end{array}$$

where for each layer t, ${\mathsf{\Phi }}_t:{\mathsf{\Gamma }}_t\times Y_t\to Y_{t+1}$ is a nonlinear transformation parameterized by the weights ${\gamma }_t\in {\mathsf{\Gamma }}_t$, and $Y_t$, ${\mathsf{\Gamma }}_t$ are the appropriately chosen spaces. In the case of a very deep network, when $T\to \infty$, it is convenient to consider the continuous time limit of the above expression such that

$$\begin{array}{c}\hfill y\left(t\right)=\mathsf{\Phi }\left(\gamma \left(t\right),x\right),t>0,\end{array}$$

where the parametric evolution function $\mathsf{\Phi }:\mathsf{\Gamma }\times Y\to Y$ defines a continuous flow through the input data $y\left(0\right)=x\in Y$. Parameter estimation for such continuous evolution can be viewed as an optimal controlling problem [6], given by

$$\begin{array}{cc}\hfill & \underset{\gamma \left(t\right)}{min}{\mathbb{E}}_{\mu }\left[L\left(y\left(T\right),f\left(x\right)\right)+{\int }_0^{T}R\left(\gamma \left(t\right),y\left(t\right)\right)dt\right],\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill & \mathrm{subject}\mathrm{to}y\left(t\right)=\mathsf{\Phi }\left(\gamma \left(t\right),x\right),\hfill \end{array}$$

(2)

where $L\left(y\left(T\right),f\left(x\right)\right)$ is a terminal loss function, $R\left(\gamma \left(t\right),y\left(t\right)\right)$ is a regularizer, and $\mu$ is a probability distribution of the input–target data pairs $(x,f(x\left)\right)$. More general models additionally consider continuity in the “patial” dimension as well by using differential [7] or integral formulations [8]. A continuous time formulation based on ordinary differential equations (ODEs) was proposed in [9] with the state Equation (2) of the form

$$\begin{array}{c}\hfill \dot{y}\left(t\right)=\mathsf{\Phi }\left(\gamma \left(t\right),y\left(t\right)\right).\end{array}$$

(3)

In the work [9], the authors relied on the black-box ODE solvers and used adjoint sensitivity analysis to derive equations for the backpropagation of errors through the continuous system.

The authors of [10] concentrated on the well-posedness of the learning problem for ODE-constrained control and emphasized the importance of stability in the design of deep architectures. For instance, the solution of a homogeneous linear ODE with constant coefficients

$$\begin{array}{c}\hfill \dot{y}\left(t\right)=Ay\left(t\right)\end{array}$$

is given by

$$\begin{array}{c}\hfill y\left(t\right)=Q{e}^{\mathsf{\Lambda }t}{Q}^{T}x,\end{array}$$

where $A=Q\mathsf{\Lambda }{Q}^T$ is the eigendecomposition of a matrix A, and $\mathsf{\Lambda }$ is the diagonal matrix with the corresponding eigenvalues. The similar equation holds for the backpropagation of gradients. To guarantee the efficient propagation of information through the network, one must ensure that the elements of $e^{\mathsf{\Lambda }t}$ have magnitudes close to one. This condition, of course, is satisfied when all eigenvalues of the matrix A are imaginary with real parts close to zero. In order to preserve this property, the authors of [10] proposed several time continuous architectures of the form

$$\begin{array}{c}\hfill \left\{\begin{array}{c}\dot{y}\left(t\right)={\mathsf{\Phi }}_1\left({\gamma }_1\left(t\right),y\left(t\right),z\left(t\right)\right),\hfill \\ \dot{z}\left(t\right)={\mathsf{\Phi }}_2\left({\gamma }_2\left(t\right),y\left(t\right),z\left(t\right)\right).\hfill \end{array}\right.\end{array}$$

(4)

When ${\mathsf{\Phi }}_1(y,z)={\nabla }_{z}H(y,z)$, ${\mathsf{\Phi }}_2(y,z)=-{\nabla }_{y}H(y,z)$, the equations above provide an example of a conservative Hamiltonian system with the total energy H.

In the discrete setting of the ordinary feed forward networks, the necessary conditions for the optimal solution of (1) and (2) recover the well-known equations for the forward propagation (state Equation (2)), backward gradient propagation (co-state equation), and the optimality condition, to compute the weights (gradient descent algorithm); see, e.g, [11]. The continuous setting offers additional flexibility in the construction of discrete networks with the desired properties and efficient learning algorithms. Classical feed forward networks (Figure 1, left) is just the particular and the simplest example of such discretization, which is prone to all the issues of deep learning. In order to facilitate the training process, a skip-connection is often added to the network (Figure 1, middle) yielding

$$\begin{array}{c}\hfill y=x+h·\mathsf{\Phi }(\gamma ,x),\end{array}$$

(5)

where h is a positive hyperparameter. Equation (5) can be viewed as a forward Euler scheme to solve the ODE in (3) numerically on the time grid with step size h. While it was shown that such residual layers help to mitigate the problem of vanishing gradients and speed-up the training process [12], the scheme has very restrictive stability properties [13]. This can result in the uncontrolled accumulation of errors at the inference stage reducing the generalization ability of the trained network. Moreover, the Euler scheme is not capable of preserving geometric structure of conservative flows and is thus a bad choice for the long time integration of such ODEs [14].

Memory efficient explicit reversible architectures can be obtained by considering time discretization of the partitioned system of ODEs in (4). The reversibility property allows for recovering the internal states of the system by propagating through the network in both directions and thus does not require one to cache these values for the evaluation of the gradients. First, such architecture (RevNet) was proposed in [15], and, without using a connection to discrete solutions of ODEs, it has the form

$$\begin{array}{c}\hfill \left\{\begin{array}{c}{y}_1=x_1+{\mathsf{\Phi }}_1(\gamma ,x_2),\hfill \\ y_2=x_2+{\mathsf{\Phi }}_2(\gamma ,y_1),\hfill \end{array}\right.⟷\left\{\begin{array}{c}{x}_2=y_2-{\mathsf{\Phi }}_2(\gamma ,y_1),\hfill \\ x_1=y_1-{\mathsf{\Phi }}_1(\gamma ,x_2).\hfill \end{array}\right.\end{array}$$

It was later recognized as the Verlet method applied to the particular form of the system in (4), see [10,16]. The leapfrog and midpoint networks are two other examples of reversible architectures proposed in [16].

Other residual architectures can be also found in the literature including Resnet in Resnet (RiR) [17], Dense Convolutional Network (DenseNet) [18] and linearly implicit network (IMEXNet) [19]. For some problems, all of these networks show a substantial improvement over the classical ResNet but still have an explicit structure, which has limited robustness to the perturbations of the input data and parameters of the network. Instead, in this effort, we propose new fully implicit residual architecture, which, unlike the above mentioned examples, is unconditionally stable and robust.

As opposed to the standard feed-forward networks, the outputs of the proposed implicit residual blocks are defined as the fixed points of the appropriately chosen nonlinear transformations as follows:

$$\begin{array}{c}\hfill y=x+\mathsf{\Phi }(\gamma ,x,y).\end{array}$$

(6)

The right part of Figure 1 provides a graphical illustration of the proposed layer. One can immediately recognize the feedback loop which is typical for recurrent neural networks (RNN). The standard approach to train RNNs is by backpropagation through time, the algorithm which requires substantial memory resources for deep unrolled recurrent networks. The authors of [20,21] utilized related to our idea of implicit layers to cope with this issue by directly learning the equilibrium points of “infinite” depth recurrent models with a constant memory complexity. The main difference of our approach is that we design the feedback loop with a specific goal of driving the output of the implicit layer to the stable fixed point. We will discuss the choice of the nonlinear transformation $Phi$ in (6) and the design of the training algorithm in the next section. It is also worth noting that the idea of learning fixed points is not quite new and is rooting way back to the Hopfield networks and content-addressable (“associative”) memory [22,23].

After the first version of this manuscript has appeared, another work proposed the similar idea of implicit Euler skip connections to enhance the adversarial robustness of residual networks [24]. The authors of [24] modified the original residual block by complementing it with a fixed number of steps of the gradient descent algorithm and applied adversarial training to the modified architecture. While both efforts are inspired by implicit numerical schemes for integrating ODEs, our approach is more general as we ensure the convergence of the proposed implicit layer to the stable fixed point. As discussed above and unlike the work in [24], in addition to the enhanced stability properties, this approach does not increase the memory complexity of the original residual networks. We also propose an initialization and regularization strategies which allow for training the network efficiently and admit simple interpretation.

The preliminary results for the work proposed in the current manuscript have been presented at the Second Symposium on Machine Learning and Dynamical Systems in the Fields Institute [25]. Here, we provide a completely revised version of this work including an in-depth description of the method, a new regularization approach, and much extended numerical results.

2. Description of the Method

We first motivate the necessity for our new method by letting the continuous model of a network be given by the ordinary differential equations in (4) that is:

$$\begin{array}{c}\hfill \left\{\begin{array}{c}\dot{y}\left(t\right)={\mathsf{\Phi }}_1\left({\gamma }_1\left(t\right),y\left(t\right),z\left(t\right)\right),\hfill \\ \dot{z}\left(t\right)={\mathsf{\Phi }}_2\left({\gamma }_2\left(t\right),y\left(t\right),z\left(t\right)\right).\hfill \end{array}\right.\end{array}$$

An s-stage Runge–Kutta (RK) method for the approximate solution of the above equations is given by

$$\begin{array}{cc}\hfill k_i& ={\mathsf{\Phi }}_1\left({\gamma }_1(t_0+c_{i}h),y_0+h\sum _{j=1}^s{a}_{ij}{k}_j,z_0+h\sum _{j=1}^s{\widehat{a}}_{ij}{l}_j\right)\hfill \\ \hfill l_i& ={\mathsf{\Phi }}_2\left({\gamma }_2(t_0+{\widehat{c}}_{i}h),y_0+h\sum _{j=1}^s{a}_{ij}{k}_j,z_0+h\sum _{j=1}^s{\widehat{a}}_{ij}{l}_j\right)i,j=1,..,s\hfill \\ \hfill y_1& =y_0+h\sum _{i=1}^s{b}_i{k}_i,z_1=z_0+h\sum _{i=1}^s{\widehat{b}}_i{l}_i.\hfill \end{array}$$

(7)

The order conditions for the coefficients $a_{ij}$, $b_i$, $c_i$, ${\widehat{a}}_{ij}$, ${\widehat{b}}_i$, and ${\widehat{c}}_i$, which guarantee that convergence of the numerical solution is well known and can be found in any topical text, see, e.g., [13]. Note that, when $a_{ij}\ne 0$ or ${\widehat{a}}_{ij}\ne 0$ for at least some $j\ge i$, the scheme is implicit and a system of nonlinear equations has to be solved at each iteration which obviously increases the complexity of the solver. Nevertheless, the following example illustrates the benefits of using implicit approximations.

2.1. Linear Stability Analysis

Consider the following linear differential system:

$$\begin{array}{cc}\hfill \dot{y}\left(t\right)& =-{\omega }^{2}z\left(t\right),\dot{z}\left(t\right)=y\left(t\right)\hfill \end{array}$$

(8)

and four simple discretization schemes [13]:

Forward Euler:	$y_1=y_0-h{\omega }^2{z}_0,z_1=z_0+h{y}_0,$
Backward Euler:	$y_1=y_0-h{\omega }^2{z}_1,z_1=z_0+h{y}_1,$
Trapezoidal:	${\displaystyle y_1=y_0-\frac{h{\omega }^2}{2}\left(z_0+z_1\right),z_1=z_0+\frac{h}{2}\left(y_0+y_1\right),}$
Verlet:	${\displaystyle y_{0.5}=y_0-\frac{h{\omega }^2}{2}{z}_0,z_1=z_0+h{y}_{0.5},y_1=y_{0.5}-\frac{h{\omega }^2}{2}{z}_1.}$

Due to linearity of the system in (8), the numerical solution after n steps can be written as

$$\begin{array}{c}\hfill \left(\begin{array}{c}{y}_n\\ z_n\end{array}\right)=R^n\left(h\omega \right)\left(\begin{array}{c}{y}_0\\ z_0\end{array}\right).\end{array}$$

(9)

The long time behavior of the discrete dynamics is hence determined by the spectral radius of the matrix $R\left(h\omega \right)$ (called the stability matrix), which needs to be less than or equal to one for the sake of stability. For example, we have ${\lambda }_{1,2}=1\pm ih\omega$ for the forward Euler scheme, and the method is unconditionally unstable. The Backward Euler scheme gives ${\lambda }_{1,2}={(1\pm ih\omega )}^{-1}$ and the method is unconditionally stable. The corresponding eigenvalues of the trapezoidal scheme have a magnitude equal to one for all $\omega$ and h. Finally, the characteristic polynomial for the matrix of the Verlet scheme is given by ${\lambda }^2-(2-h^2{\omega }^2)\lambda +1$, i.e., the method is only conditionally stable when $\left|h\omega \right|\le 2$.

Figure 2 illustrates this behavior for the particular case of $\omega =50$. Notice that the flows of the forward and backward Euler schemes are strictly expanding and contracting; if one had to fit the exact flow of the system in (8) using either of these methods, this would result in an inherently ill-posed problem. On the contrary, the implicit trapezoidal and explicit Verlet schemes seem to reproduce the original flow very well, but the latter is conditional on the size of the step h. Another nice property of the trapezoidal and Verlet schemes is their symmetry with respect to the exchanging $y_n\leftrightarrow y_{n-1}$ and $z_n\leftrightarrow z_{n-1}$. Such methods play a central role in the geometric integration of reversible differential flows and are handy in the construction of the memory efficient reversible network architectures. Conditions for the reversibility of general Runge–Kutta schemes can be found in [14].

The discussion above highlights the importance of the appropriate choice for both the structure of a dynamical system (DS) and the corresponding time integrator. The stability of a general Runge–Kutta method is often studied in application to the simpler scalar test equation

$$\begin{array}{c}\hfill \dot{y}\left(t\right)=\lambda y,\end{array}$$

(10)

and by analogy with (9) its stability function is given by

$$\begin{array}{c}\hfill \mathfrak{R}\left(z\right)=\frac{det\left(I-zA+z𝟙b^T\right)}{det\left(I-zA\right)},z=h\lambda ,\end{array}$$

where $𝟙={(1,...,1)}^T$ and $A={\left(a_{ij}\right)}_{i,j=1}^s$, $b={\left(b_j\right)}_{j=1}^s$ are the parameters of the RK method in (7) with the remaining parameters set to zero for the scalar test equation. From this expression, one can see that the stability function of any explicit Runge–Kutta method is a polynomial and hence is bounded only for z in a finite region of the complex plane. The set of all such z with $R\left(z\right)<1$ is called the stability region of the method.

2.2. Implicit ResNet

Motivated by the discussion above, we propose an implicit residual layer in (6) with a nonlinear map $\mathsf{\Phi }(\gamma ,x,y)$ given by

$$\begin{array}{cc}\hfill & \mathsf{\Phi }(\gamma ,x,y):=(1-\theta )F(\gamma ,x)+\theta F(\gamma ,y),\theta \in [0,1],\hfill \end{array}$$

(11)

$$\begin{array}{cc}\hfill \mathrm{or}& \mathsf{\Phi }(\gamma ,x,y):=F(\gamma ,(1-\theta )x+\theta y),\hfill \end{array}$$

(12)

where x, y, $\gamma$ are the input, output, and parameters of the layer, and $F(\gamma ,x)$ is a vector field to be estimated.

Table 1. Derivatives of the nonlinear maps in (11) and (12).

$\mathsf{\Phi }(\gamma ,x,y)$	$(1-\theta )F(\gamma ,x)+\theta F(\gamma ,y)$	$F(\gamma ,z),z=(1-\theta )x+\theta y$
$\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial x}$	$(1-\theta )\frac{\partial F(\gamma ,x)}{\partial x}$	$(1-\theta )\frac{\partial F(\gamma ,z)}{\partial z}$
$\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}$	$\theta \frac{\partial F(\gamma ,y)}{\partial y}$	$\theta \frac{\partial F(\gamma ,z)}{\partial z}$
$\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial \gamma }$	$(1-\theta )\frac{\partial F(\gamma ,x)}{\partial \gamma }+\theta \frac{\partial F(\gamma ,y)}{\partial \gamma }$	$\frac{\partial F(\gamma ,z)}{\partial \gamma }$

shows the derivatives of the nonlinear maps in (11) and (12) with respect to their arguments.

The stability function of the layers in (11) and (12) is given by

$$\begin{array}{c}\hfill \mathfrak{R}\left(z\right)=\frac{1+(1-\theta )z}{1-\theta z}.\end{array}$$

(13)

The corresponding stability regions are illustrated in Figure 3 indicating the improved stability of implicit layers for increasing $\theta$.

Additionally, instead of a single layer in (6), one might consider a block of implicit layers on a given “time interval” $t\in [0,T]$

$$\begin{array}{cc}\hfill & y_t=y_{t-1}+\mathsf{\Phi }({\gamma }_t,y_{t-1},y_t),t=1,\cdots ,T,\hfill \\ \hfill & y_0=x.\hfill \end{array}$$

(14)

Note that, by viewing (14) as a numerical ODE integrator and to get the above expressions, we have assumed that the ODE coefficients $\gamma \left(t\right)$ are piecewise constant cádlág functions, see Figure 4.

We will use (12) as the definition of implicit layers in the rest of this work.

2.2.1. Forward Propagation

Assume that the fixed point in (6) exists and can be computed. To solve the corresponding nonlinear equation, consider the equivalent minimization problem

$$\begin{array}{cc}\hfill & \underset{y}{min}{\parallel r\left(y\right)\parallel }^2,r\left(y\right):=y-x-\mathsf{\Phi }(\gamma ,x,y).\hfill \end{array}$$

(15)

One way to construct the required solution is by applying the descent algorithm

$$\begin{array}{c}\hfill y^{k+1}\leftarrow y^k-{\lambda }^k{s}^k,k=0,1,2,\dots \end{array}$$

where $s^k$ is the descent direction and ${\lambda }^k$ is the corresponding step size. Several common choices are summarized in

Table 2. Common descent algorithms [26].

Method	${\mathit{s}}^{\mathit{k}}$	${\mathit{\lambda }}^{\mathit{k}}$	Parameters
gradient descent	${\mathsf{\Lambda }}^k·{\nabla }_y{\parallel r\left(y^k\right)\parallel }^2+m^k/{\lambda }^k$	Wolfe conditions	scaling matrix ${\mathsf{\Lambda }}^k$ momentum vector $m^k$
quasi-Newton	${\left(H^k\right)}^{-1}·{\nabla }_y{\parallel r\left(y^k\right)\parallel }^2$	Wolfe conditions	approximate Hessian $H^k$
conjugate gradient	${\nabla }_y{\parallel r\left(y^k\right)\parallel }^2-{\beta }^k{s}^{k-1}$	${arg\; min}_{\lambda }$ $\parallel r(y^k+\lambda s^k){\parallel }^2$	conjugate direction parameters ${\beta }^k$

.

The required gradient of the residual norm can be computed as

$$\begin{array}{c}\hfill {\nabla }_y{\parallel r\left(y\right)\parallel }^2=2·{\frac{\partial r\left(y\right)}{\partial y}}^T·r\left(y\right),\end{array}$$

(16)

where the Jacobian of the residual vector is given by (c.f. Table 1)

$$\begin{array}{c}\hfill \frac{\partial r\left(y\right)}{\partial y}=I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}.\end{array}$$

The expression in (16) can be efficiently computed with the automatic differentiation capabilities of any standard deep learning framework making it possible to interface existing iterative solvers. However, optimized implementations of the gradient descent algorithms are readily provided by any such framework. Hence, the forward propagation for implicit layers can be easily implemented using built-in tools native to each framework and without interfacing external solvers. Listing 1 shows the PyTorch pseudocode of the implicit layer and the left part of Figure 5 illustrates its corresponding computational graph.

Listing 1. PyTorch pseudocode of the implicit residual block.

import torch

 

class fpmap (torch.Function):

@staticmethod

def forward (ctx, γ, x, ŷ):

y = x + Φ(γ, x, ŷ)

ctx.save_for_backward (y, ŷ)

return y

 

@staticmethod

def backward (ctx, ∇_yL):

y, ŷ = ctx.saved_tensors

return (I − ∂y/∂ŷ)^−T∇_yL

 

def ImplicitResidualBlock (x):

 

nsolve = lambda γ, x: argmin_z ||z − fpmap (γ, x, z)||²

 

ŷ = nsolve (γ.detach(), x.detach())

 

return fpmap.apply (γ, x, ŷ)

2.2.2. Backpropagation

We now show that, even though the nonlinearity in (6) adds to the complexity of the forward propagation, the direct backpropagation through the nonlinear solver is not required. Firstly, using the chain rule, we can easily find the Jacobian matrices of the implicit residual layer as follows:

$$\begin{array}{cc}\hfill \frac{\partial y}{\partial x}& =I+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial x}+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\frac{\partial y}{\partial x}={{\displaystyle \left(I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\right)}}^{-1}{\displaystyle \left(I+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial x}\right)},\hfill \end{array}$$

and

$$\begin{array}{cc}\hfill \frac{\partial y}{\partial \gamma }& =\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial \gamma }+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\frac{\partial y}{\partial \gamma }={\left(I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\right)}^{-1}\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial \gamma }.\hfill \end{array}$$

The backpropagation formulas then follow immediately:

$$\begin{array}{cc}\hfill & {\nabla }_{x}L={\left(I+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial x}\right)}^T\overline{{\nabla }_{y}L}\mathrm{and}{\nabla }_{\gamma }L={\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial \gamma }}^T\overline{{\nabla }_{y}L},\hfill \end{array}$$

where $\overline{{\nabla }_{y}L}$ is the solution to the linear system

$$\begin{array}{cc}\hfill & {\left(I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\right)}^T\overline{{\nabla }_{y}L}={\nabla }_{y}L.\hfill \end{array}$$

(17)

Note that the custom backpropagation for the $\mathrm{fpmap}$ function in Listing 1 is responsible for the linear solve in (17). The gradients of the loss with respect to the parameters and the input are then computed automatically by the deep learning framework. Finally, DL frameworks allow for the cheap computation of the vector-Jacobian products in (17) and hence for the efficient implementation of iterative linear solvers. For example, we utilize restarted GMRES as a linear solver in our implementation.

2.3. fpResNet

Sophisticated solvers are not required for the nonlinear and linear systems in (15) and (17) when $\mathsf{\Phi }(\gamma ,x,y)$ is a contractive mapping, i.e., when

$$\begin{array}{c}\hfill Lip\left(\mathsf{\Phi }\right):=\underset{y}{sup}{∥\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}∥}_2=\theta \underset{y}{sup}{∥\frac{\partial F(\gamma ,y)}{\partial y}∥}_2<1,\end{array}$$

(18)

where ${\parallel ·\parallel }_2$ is the spectral norm of the matrix equal to its largest singular value. The red color in Figure 3 is used to highlight the part of the complex plane $\{z:|z|<{\theta }^{-1}\}$, where the above condition is satisfied for the Dahlquist test equation in (10). In this case, the Banach fixed-point theorem ensures the convergence of the recurrence relation

$$\begin{array}{c}\hfill y^{k+1}=x+\mathsf{\Phi }(\gamma ,x,y^k),k=0,1,...\end{array}$$

during the forward propagation. The same condition guarantees the validity of the Neumann series expansion for the matrix inverse required for the backpropagation; one gets

$$\begin{array}{c}\hfill \overline{{\nabla }_{y}L}={\left(I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\right)}^{-T}{\nabla }_{y}L=\sum _{i=0}^{\infty }{\left({\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}}^T\right)}^i{\nabla }_{y}L=\sum _{i=0}^{\infty }{\nabla }_y^{i}L,\end{array}$$

with

$${\nabla }_y^{i}L={\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}}^T{\nabla }_y^{i-1}L,{\nabla }_y^{0}L={\nabla }_{y}L.$$

Similarly to (17), each ${\nabla }_y^{i}L$ has a form of the vector-Jacobian product which can be efficiently evaluated with any deep learning framework. In practice, however, such simple iterations converge linearly with the rate proportional to the Lipschitz constant $Lip\left(\mathsf{\Phi }\right)$ which can be rather inefficient when $Lip\left(\mathsf{\Phi }\right)\approx 1$.

2.4. Regularization

One way to ensure the stability of the neural network when viewed as a nonlinear DS is by imposing a hard constraint on its parameters to make it globally dissipative or conservative. This approach has been utilized, for instance, in [7,10] and applied to several explicit, and hence only conditionally stable, residual network architectures. Another approach that we employ here is by imposing the structure by regularization. In this section, we review some common regularization techniques and propose a new one that suits the presented implicit architecture.

2.4.1. Lipschitz Continuous Architectures

Enforcing Lipschitz continuity of neural networks has been recognized as an important component in many applications. For instance, explicit bounds on the Lipschitz constants of the loss function have been utilized to improve the robustness and establish the generalization error of large margin classifiers and GAN discriminators in [27,28,29,30], respectively. Lipschitz continuity has also been considered implicitly in [31,32] to improve the adversarial robustness of DNNs and in [33] for the design of contractive auto-encoders. In all these works, bounding the Lipschitz constants was implemented by penalizing the norm of the Jacobian matrix of the network on the training dataset.

Spectral normalization of weight matrices can be used to provide the uniform bound on the Lipschitz constant of the network. It has been applied to improve the stability of deep networks in [34,35,36] and to construct invertible normalizing flows of probability distributions in [37]. To justify the method, consider the common choice of $F(\gamma ,x)$ in (11) as a composition of affine maps and contractive nonlinear activations, i.e.,

$$\begin{array}{c}\hfill F(\gamma ,x)={\varphi }_n\circ {\varphi }_{n-1}\circ \dots \circ {\varphi }_1\circ x\mathrm{with}{\varphi }_i\circ x=\sigma ({\gamma }_i\circ x+b_i)\mathrm{s}.\mathrm{t}.\parallel \sigma \parallel \le 1.\end{array}$$

The Lipschitz constant of $F(\gamma ,x)$ as a function of x can be bounded from above by

$$\begin{array}{c}\hfill Lip\left(F\right):=\underset{x}{sup}{∥\frac{\partial F(\gamma ,x)}{\partial x}∥}_2\le \prod _{i=1}^n{\parallel {\gamma }_i\parallel }_2=\prod _{i=1}^n\sqrt{\rho \left({\gamma }_i^T{\gamma }_i\right)},\end{array}$$

where ${\parallel ·\parallel }_2$ is the spectral norm and $\rho \left(A\right)$ denotes the spectral radius of a linear map A. Hence, to ensure that $Lip\left(F\right)\le \alpha$, it is enough to take

$$\begin{array}{c}\hfill {\tilde{\gamma }}_i=\frac{{\gamma }_i}{max\left({\displaystyle 1,\frac{\parallel {\gamma }_i{\parallel }_2}{{\alpha }_i}}\right)}\mathrm{s}.\mathrm{t}.\prod _{i=1}^n{\alpha }_i\le \alpha .\end{array}$$

The exact calculation of the operator norm is expensive and one usually appeals to approximate techniques such as the power iteration method [38]. According to this method, the dominant singular vector v and the singular value $\mu$ of a linear operator A are estimated iteratively as

$$\begin{array}{cc}\hfill v^{k+1}& =\frac{A^{T}A{v}^k}{\parallel A^{T}A{v}^k\parallel },{\mu }^k=\sqrt{{\left(v^k\right)}^T{A}^{T}A{v}^k}=\parallel A{v}^k\parallel .\hfill \end{array}$$

In practice, it is enough to take a fixed number (usually 1) of iterations at each weight evaluation during the training stage since the parameters are not expected to change much close to the convergence of the training loop. By observing that

$$\begin{array}{c}\hfill A^{T}A{v}^k=\frac{1}{2}\frac{\partial \parallel A{v}^k{\parallel }^2}{\partial v^k}=\frac{1}{2}\frac{\partial {\left({\mu }^k\right)}^2}{\partial v^k},\end{array}$$

Algorithm 1 provides a simple implementation of this approach.

Algorithm 1 Power iteration method

Input: linear map A, max iterations $k_{max}$ Initialize $v^0$ $v^0\leftarrow v^0/\parallel v^0\parallel$ for$k=1,...,k_{max}$do $\mu \leftarrow \parallel A{v}^k\parallel$ ${\displaystyle v^{k+1}\leftarrow 0.5·\partial {\mu }^2/\partial v^k}$ ${\displaystyle v^{k+1}\leftarrow v^{k+1}/\parallel v^{k+1}\parallel }$ end for Output:$\mu$

More generally, spectral normalization allows for controlling the spread of the Jacobian matrix ${\displaystyle \frac{\partial F(\gamma ,x)}{\partial x}}$, i.e., the largest distance between its eigenvalues

$$s{\displaystyle \left(\frac{\partial F(\gamma ,x)}{\partial x}\right)}:=\underset{i,j}{max}\left\{|{\lambda }_i-{\lambda }_j|\right\}.$$

By definition of the spectral radius, one has

$$\begin{array}{c}\hfill \rho {\displaystyle \left(\frac{\partial F(\tilde{\gamma },x)}{\partial x}\right)}\le {{\displaystyle ∥\frac{\partial F(\tilde{\gamma },x)}{\partial x}∥}}_2\le 1\to Re\left({\lambda }_i\right)\in [-1,1]\forall i.\end{array}$$

Hence, by denoting $F^{-1,1}(\gamma ,x):=F(\tilde{\gamma },x)$, we obtain

$$\begin{array}{c}\hfill F^{\alpha ,\beta }(\gamma ,x)=\frac{\alpha +\beta }{2}x+\frac{\beta -\alpha }{2}{F}^{-1,1}(\gamma ,x)\to Re\left({\lambda }_i\right)\in [\alpha ,\beta ]\forall i,\end{array}$$

so that all eigenvalues of the Jacobian ${\displaystyle \frac{\partial F^{\alpha ,\beta }(\gamma ,x)}{\partial x}}$ are located in the disc with radius $(\beta -\alpha )/2$ centered at $(\alpha +\beta )/2$. In practice, we use the more flexible form of this definition

$$\begin{array}{c}\hfill F^{\alpha ,\beta }(\gamma ,x)=\frac{\alpha +\beta }{2}x+\frac{\beta -\alpha }{2}S\left(\vartheta \right)\odot F^{-1,1}(\gamma ,x),\end{array}$$

(19)

where ⊙ is the Hadamard product, $\vartheta$ are additional learnable parameters, $S\left(\vartheta \right)$ has the same dimensionality as $F^{-1,1}(\gamma ,x)$, and each $S_i\left({\vartheta }_i\right)\in (0,1)$ is the sigmoid function.

2.4.2. Trajectory Regularization

The Lipschitz constant of the proposed residual block in (6) is defined as

$$\begin{array}{c}\hfill Lip\left(y\right):=\underset{x}{sup}{{\displaystyle ∥{{\displaystyle \left(I-\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial y}\right)}}^{-1}{\displaystyle \left(I+\frac{\partial \mathsf{\Phi }(\gamma ,x,y)}{\partial x}\right)}∥}}_2,\end{array}$$

and, for the linear scalar test equation in (10), it is given by the stability function (13) of the method. The hatched circle in Figure 3 contains the spectrum of the 1-Lipschitz function $F(\tilde{\gamma },x)$ and shows that the Lipschitz constant of a residual block can fall outside of its stability region. Moreover, the pole of (13) is located at $z={\theta }^{-1}$ and the stability of implicit layers might actually degrade with increasing $\theta \in [0,1]$ if no additional precautions are taken to isolate the spectrum of the layers from its vicinity. This can be achieved, for instance, by setting $F(\gamma ,x):=F^{\alpha ,{\theta }^{-1}}(\gamma ,x)$ for some $\alpha <{\theta }^{-1}$.

Previous works have focused on improving the efficiency of residual and neural ODE architectures by regularizing their vector fields in a way that leads to a simpler dynamical behavior. For example, the authors of [39] considered the following regularizer

$$\begin{array}{c}\hfill R\left(\gamma \right):={\alpha }_K{\int }_0^T{{\displaystyle ∥F\left(\gamma \right(t),y(t\left)\right)∥}}^{2}dt+{\alpha }_J{\int }_0^T{{\displaystyle ∥\frac{\partial F\left(\gamma \right(t),y(t\left)\right)}{\partial y}∥}}_F^{2}dt.\end{array}$$

The first term encourages the trajectories with origin in the training dataset to follow straight lines, while the second term reduces overfitting by restricting the vector field to be nearly constant in the vicinity of each trajectory. A more general approach has been taken in [40] by directly penalizing the K-th order total derivative of the vector field along the solution trajectories. It has been shown that, by matching K to the order of numerical integrator, it is possible to significantly reduce the cost of solving the learned dynamics without sacrificing the resulting accuracy.

Instead, we consider “discrete-time” regularization of the form

$$\begin{array}{c}\hfill R\left(\gamma \right):=\frac{1}{T}{\displaystyle \left(\sum _{t=0}^T{\phantom{\sum }}^{\prime }{\displaystyle \left[\frac{{\alpha }_{div}}{d}{{\displaystyle \left(\frac{t}{T}\right)}}^p\nabla ·F({\gamma }_t,y_t)+\frac{{\alpha }_{jac}}{d^2}{{\displaystyle ∥\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}∥}}_F^2\right]}+\frac{{\alpha }_{TV}}{T}\sum _{t=1}^T{{\displaystyle ∥{\gamma }_t-{\gamma }_{t-1}∥}}^2\right)},\end{array}$$

(20)

where $\sum {\phantom{\sum }}^{\prime }$ is the trapezoidal quadrature rule, d is the dimension of the vector field $F:\mathsf{\Gamma }\times {\mathbb{R}}^d\to {\mathbb{R}}^d$, $p\ge 0$ is some fixed number, and $\nabla ·F({\gamma }_t,y_t)$ is the divergence of $F({\gamma }_t,y_t)$.

The total variation (TV) like term in (20) is responsible for the temporal regularity of the vector field $F\left(\gamma \right(t),y(t\left)\right)$. Similarly to the approach of [39,40], the second term penalizes the Jacobian matrix of the layer producing vector fields which tend to be constant in the vicinity of the learned trajectories; this results in simpler dynamics and faster convergence. Finally, the first term in (20) promotes the negative divergence of $F({\gamma }_t,y_t)$ along the trajectories. By definition of the divergence of the vector field, one has

$$\begin{array}{c}\hfill \nabla ·F({\gamma }_t,y_t)=\sum _{i=1}^d{{\displaystyle \left(\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)}}_{ii}=Tr{\displaystyle \left(\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)}=\sum _{i=1}^d{\lambda }_i{\displaystyle \left(\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)},\end{array}$$

i.e., it is equal to the sum of eigenvalues of the Jacobian matrix. By minimizing this term, we attempt to push the spectrum of the Jacobian matrix to the negative part of the complex plane so that we can take advantage of the enhanced stability of implicit layers in this region. In addition, note that, by definition, the squared Frobenius norm of a matrix is equal to the sum of its squared singular values

$$\begin{array}{c}\hfill {{\displaystyle ∥\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}∥}}_F^2=\sum _{i=1}^d\sum _{j=1}^d{{\displaystyle \left(\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)}}_{ij}^2=Tr{\displaystyle \left(\frac{\partial F^T({\gamma }_t,y_t)}{\partial y_t}\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)}=\sum _{i=1}^d{\sigma }_i^2{\displaystyle \left(\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}\right)}.\end{array}$$

Hence, the first two terms in (20) are competing with each other, and the coefficient ${(t/T)}^p$ is used to balance these two components by increasing the level of dissipation along the trajectories. The divergence term, however, does not impact the off-diagonal part of the Jacobian matricies, and, for large enough ${\alpha }_{div}$, this will promote their diagonal dominance.

According to the Gershgorin circle theorem, every eigenvalue of a matrix A lies within at least one of the Gershgorin discs $D(a_{ii},r_i)$ with $r_i={\sum }_{i\ne j}\left|a_{ij}\right|$. This means that the optimal solution of the optimization problem with the proposed regularizer in (20) will result in the dynamics that are strongly dissipative in the directions irrelevant for the accurate representation of the training data effectively reducing the dimension of the state space. The behavior of the dynamical system in the remaining directions along the so-obtained low-dimensional state manifold can potentially be arbitrary and, with a proper balance between the loss and regularization, should not decrease the expressive power of the network.

For the efficient evaluation of the divergence and Jacobian regularizers in (20), we utilize the unbiased stochastic Hutchinson trace estimator [41] to obtain

$$\begin{array}{cc}\hfill \nabla ·F({\gamma }_t,y_t)& ={\mathbb{E}}_{z\sim \mathcal{N}(0,1)}{\displaystyle \left[z^T\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}z\right]},\hfill \\ \hfill {{\displaystyle ∥\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}∥}}_F^2& ={\mathbb{E}}_{z\sim \mathcal{N}(0,1)}{\displaystyle \left[{{\displaystyle ∥z^T\frac{\partial F({\gamma }_t,y_t)}{\partial y_t}∥}}^2\right]}.\hfill \end{array}$$

This approach has also been applied in [39,42]. Algorithm 2 provides a more general variant of Hutchinson algorithm which allows for estimating the diagonal of a matrix [43]. This can be convenient when one needs to control the magnitude of the diagonal elements rather than just their sum.

Algorithm 2 Stochastic diagonal estimator

Input: linear map A, max iterations $k_{max}$ Initialize $D^0=0$ for$k=1,...,k_{max}$do $v^k\leftarrow \mathcal{N}(0,1)$ $t^k\leftarrow t^{k-1}+\left(A\left(v^k\right)\odot v^k\right)$ $q^k\leftarrow t^{k-1}+\left(v^k\odot v^k\right)$ $D^k\leftarrow t^k\oslash q^k$ end for Output: Approximate diagonal $D^{k_{max}}$ of A

3. Results

The source code used to generate all examples below can be found at https://github.com/vreshniak/ImplicitResNet.

3.1. Example 1. (Regression)

For the first example, we consider the simple problem that can be easily visualized. The goal is to approximate the one-dimensional sine function in Figure 6 given $N=20$ data points evenly distributed on the interval $x\in [-5,5]$. Following the approach of [44], we augment the original one-dimensional data with an additional dimension initialized with zero. The resulting two-dimensional vector field $F(\gamma ,y):\mathsf{\Gamma }\times {\mathbb{R}}^2\to {\mathbb{R}}^2$ is approximated by a multilayer perceptron using three hidden layers of width 10 and GeLU activation function, i.e.,

$$\begin{array}{c}\hfill F(\gamma ,y)={\gamma }_{out}\sigma {\displaystyle \left({\gamma }_3\sigma {\displaystyle \left({\gamma }_2\sigma {\displaystyle \left({\gamma }_1\sigma ({\gamma }_{in}y+b_0)+b_1\right)}+b_2\right)}+b_3\right)},\end{array}$$

where ${\gamma }_{in}\in {\mathbb{R}}^{10\times 2}$, ${\gamma }_{out}\in {\mathbb{R}}^{2\times 10}$, ${\gamma }_i\in {\mathbb{R}}^{10\times 10}$ and $b_i\in {\mathbb{R}}^{10\times 1}$.

We used $T=5$ residual layers with shared parameters initialized with a Xavier uniform initializer and trained the network for 3000 epochs using Adam optimizer and the regularized loss given by

$$\begin{array}{cc}\hfill L\left(\gamma \right):=\frac{1}{N}& \sum _{i=1}^N{\displaystyle \left({\parallel g\left(y_T^i\right)-f\left(x^i\right)\parallel }^2+\frac{1}{T}\sum _{t=0}^T{\phantom{\sum }}^{\prime }{\displaystyle \left[\frac{{\alpha }_{div}}{d}{{\displaystyle \left(\frac{t}{T}\right)}}^2\nabla ·F(\gamma ,y_t^i)+\frac{0.1}{d^2}{{\displaystyle ∥\frac{\partial F(\gamma ,y_t^i)}{\partial y_t^i}∥}}_F^2\right]}\right)},\hfill \end{array}$$

(21)

where $d=2$ is the dimension of the hidden state space and $g\left(y_T\right)$ gives the last component of $y_T$. The initial learning rate was set to ${10}^{-3}$ and reduced dynamically using ReduceLROnPlateau PyTorch scheduler with patience and cooldown parameters set to 50 epochs.

Figure 7 shows the learned vector fields and eigenvalues of the Jacobian $\frac{\partial F(\gamma ,y)}{\partial y}$ along the learned trajectories for several values of $\theta$ and ${\alpha }_{div}$. Additionally, the top row of Figure 8 depicts the evolution of the loss components in (21), and the number of nonlinear iterations of the trained network for the selected parameter values. One can see that, with the Jacobian regularization alone (${\alpha }_{div}=0$, ${\alpha }_{jac}=0.1$), implicit methods demonstrate similar dynamical behavior for all considered values of $\theta$: (1) the learned vector fields take full advantage of all two available dimensions and tend to be expansive, note the increasing divergence in Figure 8, (2) the learned trajectories mostly follow straight lines, and (3) aside from the fully explicit scheme ($\theta =0$), the costs of all the methods are nearly identical, see the bottom row of Figure 8. By increasing ${\alpha }_{div}$, one starts observing the formation of a lower dimensional invariant manifold with increasingly dissipative orthogonal dynamics indicated by the negative part of the spectrum of $\frac{\partial F(\gamma ,y)}{\partial y}$ and the evolution of the filled regions in Figure 7. As the resulting dynamics becomes more restricted and less trivial, the following observations can be made: (1) the model tends to be less flexible and more difficult to fit the data, (2) the cost of the model is increasing with ${\alpha }_{div}$ and $\theta$, and (3) the explicit method demonstrates unstable oscillatory behavior when the level of dissipation exceeds its stability threshold.

3.2. Example 2. (Stiff ODE)

In our second example, we aim to fit the model to the given stiff ODE

$$\begin{array}{c}\hfill \dot{z}=-20\left(z-cos\left(t\right)\right),t\in [0,2].\end{array}$$

(22)

The training data in Figure 6 consist of 100 trajectories with randomly sampled initial conditions and the given number of uniformly distributed points along each trajectory as shown in Figure 9.

Since the exact Lipschitz constant of the forcing term is equal to $-20$, we use $F^{-25,-15}(\gamma ,x)$ in (19) with the one-dimensional vector field $F(\gamma ,x):\mathsf{\Gamma }\times {\mathbb{R}}^1\to {\mathbb{R}}^1$ given by the multilayer perceptron with 2 hidden layers of width 4 and ReLU activation function, i.e,

$$\begin{array}{c}\hfill F(\gamma ,y)={\gamma }_{out}\sigma {\displaystyle \left({\gamma }_2\sigma {\displaystyle \left({\gamma }_1\sigma ({\gamma }_{in}y+b_0)+b_1\right)}+b_2\right)},\end{array}$$

where ${\gamma }_{in}\in {\mathbb{R}}^{4\times 1}$, ${\gamma }_{out}\in {\mathbb{R}}^{1\times 4}$, ${\gamma }_i\in {\mathbb{R}}^{4\times 4}$ and $b_i\in {\mathbb{R}}^{4\times 1}$.

We took $T=20$ residual layers initialized with Xavier uniform initializer and trained the network for 50 epochs using Adam optimizer, batch of size 1, and the regularized loss given by

$$\begin{array}{c}\hfill L\left(\gamma \right):=\frac{1}{N·\{2,4,10\}}\sum _{i=1}^N\sum _{j=1}^{\{2,4,10\}}\parallel y_j^i-z^i\left(t_0^i+\frac{j}{\{1,2,5\}}\right){\parallel }^2+\frac{0.1}{T}\sum _{t=1}^T{∥{\gamma }_t-{\gamma }_{t-1}∥}^2.\end{array}$$

(23)

The top row of Figure 9 illustrates the evolution of the single trajectory generated by three residual models with $\theta =0.0,0.5$ and $1.0$. Each model was trained using the loss function in (23) with 2, 4 and 10 points taken along the trajectories of the training dataset. One can see that the explicit ResNet is unstable as was expected for the stiff system in (22), and the generated solution becomes increasingly oscillatory along the transient part of the trajectory as we increase the number of the training points from 2 to 10. This is due to the model attempting to be increasingly expressive for the data it cannot potentially fit. Once the dynamical system relaxes to the slow manifold, the accuracy of the model improves slightly but remains susceptible to the orthogonal perturbations. In order to stabilize the model, more layers need to be taken which will lead to the increased memory footprint and computational complexity. At the same time, both implicit residual networks with $\theta =0.5$ and $1.0$ are unconditionally stable and improve their accuracy with increasing number of the training points as expected.

The bottom row of Figure 9 shows the continuous dynamics generated by the vector fields learned by three considered implicit models. In this case, the discrete-time stability is not an issue anymore. However, the corruption caused by the instability of the explicit method transfers to the inaccurate behavior of the continuous system as well. On the contrary, both implicit methods lead to the satisfactory approximation of the original vector field with the midpoint scheme ($\theta =0.5$) being observably more accurate, likely due to its higher order of convergence. This suggests the proposed implicit networks as a means for learning stiff continuous-time ODE systems.

3.3. Example 3. (Periodic ODE)

In this example, we consider the Lotka–Volterra system given by

$$\begin{array}{cc}\hfill {\dot{z}}_1& =\alpha z_1-\beta z_1{z}_2,\hfill \\ \hfill {\dot{z}}_2& =\delta z_1{z}_2-\gamma z_2\hfill \end{array}$$

(24)

with $\alpha =\frac{2}{3}$, $\beta =\frac{4}{3}$, and $\delta =\gamma =1$. These equations are used to model the time evolution of the biological systems with two interacting species one being the prey and the other being the predator. The model has two equilibrium points when neither of the two interacting populations is changing. The first equilibrium is at $z_1=z_2=0$ and the non-trivial one is at $z_1=\raisebox{1ex}{$\gamma $}\!\left/ \!\raisebox{-1ex}{$\delta $}\right.=1$, $z_2=\raisebox{1ex}{$\alpha $}\!\left/ \!\raisebox{-1ex}{$\beta $}\right.=\frac{1}{2}$. Other solutions are periodic and lie on the closed curves in the phase space. Figure 6 shows five such curves; each curve contains 51 points which are distributed uniformly on the time interval $t\in [0;10]$ and used as the training data for our example.

To fit this data, we utilized implicit residual networks with $T=50$ layers and the two-dimensional vector field approximated by the multilayer perceptron with four hidden layers of width 20 and ReLU activation functions, i.e.,

$$\begin{array}{c}\hfill F(\gamma ,y)={\gamma }_{out}\sigma {\displaystyle \left({\gamma }_4\sigma {\displaystyle \left({\gamma }_3\sigma {\displaystyle \left({\gamma }_2\sigma {\displaystyle \left({\gamma }_1\sigma ({\gamma }_{in}y+b_0)+b_1\right)}+b_2\right)}+b_3\right)}+b_4\right)},\end{array}$$

where ${\gamma }_{in}\in {\mathbb{R}}^{20\times 2}$, ${\gamma }_{out}\in {\mathbb{R}}^{2\times 20}$, ${\gamma }_i\in {\mathbb{R}}^{20\times 20}$ and $b_i\in {\mathbb{R}}^{20\times 1}$. We initialized the network with Xavier uniform initializer and trained it for 3000 epochs using Adam optimizer, full batch size, and the loss given by

$$\begin{array}{c}\hfill L\left(\gamma \right):=\frac{1}{50·N}\sum _{i=1}^N\sum _{j=1}^{50}\parallel y_j^i-z^i\left(0.2j\right){\parallel }^2.\end{array}$$

Note that we did not use any form of weight normalization or regularization.

Figure 10 shows the learned trajectories and eigenvalues of the vector field along these trajectories for three implicit networks with $\theta =0.0,0.5$ and $1.0$ on the time interval $t\in [0,10]$. At a first glance, all three networks learn very similar vector fields and can accurately fit the data on the time interval of the training dataset. Moreover, the top row of Figure 11 shows that all three methods are also successful at extrapolating the dynamics to the time interval $t\in [0,200]$. However, the bottom row of the same figure shows that the midpoint scheme ($\theta =0.5$) is the only one which produces the vector field accurate for the long-time integration of the continuous dynamics. This behavior is indeed expected since the method is a geometrical integrator for this type of system, recall Figure 2 and the discussion in Section 2. The explicit residual network ($\theta =0.0$) is strictly expanding for such conservative systems, and the learned vector field tends to compensate for this behavior resulting in the dissipative continuous dynamics. The situation is reverse for the backward Euler integrator ($\theta =1.0$) since the method is strictly dissipative and hence the learned vector field is overly expanding.

3.4. Example 4. (MNIST Classification)

For the last example, consider the problem of classifying images of handwritten digits from the subset of MNIST dataset of size 1000. For this purpose, we adopt the standard preactivation ResNet-18 architecture with eight initial channels and train the network using the Adam optimizer for 100 epochs with learning rate of ${10}^{-2}$. We used the cross entropy loss function

$$\begin{array}{c}\hfill L\left(\gamma \right):=\frac{1}{N}\sum _{i=1}^N\left[-y_{label}^i+log\left(\sum _{j=0}^{9}exp\left(y_j^i\right)\right)\right]\end{array}$$

and normalized the forcing terms $F(\gamma ,x)$ of all residual layers as $F^{-3,1}(\gamma ,x)$ using (19). Additionally, the divergence regularization in (20) with ${\alpha }_{div}=0.01$ was applied to each residual layer.

To test the robustness of the trained network, we corrupted the original dataset with the Gaussian noise of varying standard deviation.

Table 3. Classification accuracy in Example 4 for different levels of Gaussian noise corruption.

Noise	Top-1 Accuracy					Top-2 Accuracy
Intensity	$\mathbf{\theta }=\mathbf{0}$	0.25	0.50	0.75	1.00	$\mathbf{\theta }=\mathbf{0}$	0.25	0.50	0.75	1.00
0.0	100.0	100.0	100.0	100.0	100.0	100.0	100.0	100.0	100.0	100.0
0.1	98.2	99.6	99.9	99.9	99.9	99.7	100.0	100.0	100.0	100.0
0.2	89.2	95.7	96.3	98.3	98.4	96.7	99.3	99.8	100.0	100.0
0.3	74.7	86.0	89.3	93.2	93.6	89.0	95.2	97.8	99.0	98.8
0.4	59.3	74.1	77.2	81.8	84.7	75.5	89.1	91.1	94.4	95.0
0.5	47.2	60.4	64.7	69.8	73.0	65.7	79.6	83.4	87.1	87.9

illustrates the classification accuracy for different levels of noise intensity and five different implicit residual networks. The results show that the implicit architectures with a proper regularization can significantly improve the robustness properties of trained networks.

4. Conclusions and Future Work

In this work, we presented a novel implicit residual layer and provided a memory-efficient algorithm to evaluate and train deep neural networks composed of such layers. We also proposed a regularization technique to control the spectral properties of the presented layer and showed that it leads to improved stability and robustness of the trained networks. The obtained numerical results support our findings.

We see several opportunities for potential improvements to the presented architecture. For example, implementations of the forward and backward propagation algorithms can be further optimized to account for the repetitive nature of the training process. This includes the better estimation of initial guesses for nonlinear solvers and preconditioners for linear solvers and other ways to reuse available information from previous runs. It is also interesting to study other, possibly multistep, types of implicit residual layers and their combinations. In this regard and in addition to the provided examples, we plan to identify the best use cases and applications for deep residual networks containing implicit layers. In particular, we are interested in exploring the impact of adversarial training and the proposed spectral regularization on the properties of the trained implicit networks. Other applications of interest include the identification of physical systems with known properties such as energy dissipation/conservation, large horizon time series forecasting, applications with corrupted and noisy data, etc. Finally, as the proper regularization is essential for the good performance of implicit layers, new regularization approaches tailored to specific applications should also be analyzed. We intend to study these questions in our future works.