Monotone operator equilibrium networks

Introduction

Recent work in deep learning has demonstrated the power of implicit-depth networks, models where features are created not by explicitly iterating some number of nonlinear layers, but by finding a solution to some implicitly defined equation. Instances of such models include the Neural ODE , which computes hidden layers as the solution to a continuous-time dynamical system, and the Deep Equilibrium (DEQ) Model , which finds a fixed point of a nonlinear dynamical system corresponding to an effectively infinite-depth weight-tied network. These models, which trace back to some of the original work on recurrent backpropagation , have recently regained attention since they have been shown to match or even exceed to performance of traditional deep networks in domains such as sequence modeling . At the same time, these models show drastically improved memory efficiency over traditional networks since backpropagation is typically done analytically using the implicit function theorem, without needing to store the intermediate hidden layers.

However, implict-depth models that perform well require extensive tuning in order to achieve stable convergence to a solution. Obtaining convergence in DEQs requires careful initialization and regularization, which has proven difficult in practice . Moreover, solutions to these models are not guaranteed to exist or be unique, making the output of the models potentially ill-defined. While Neural ODEs do guarantee existence of a unique solution, training remains unstable since the ODE problems can become severely ill-posed . Augmented Neural ODEs improve the stability of Neural ODEs by learning ODEs with simpler flows, but neither model achieves efficient convergence nor performs well on standard benchmarks. Crucial questions remain about how models can have guaranteed, unique solutions, and what algorithms are most efficient at finding them.

In this paper, we present a new class of implicit-depth equilibrium model, the Monotone Operator Equilibrium Network (monDEQ), which guarantees stable convergence to a unique fixed point.We largely use the terms “fixed point” and “equilibrium point” interchangably in this work, using fixed point in the context of an iterative procedure, and equilibrium point to refer more broadly to the point itself. The model is based upon the theory of monotone operators , and illustrates a close connection between simple fixed-point iteration in weight-tied networks and the solution to a particular form of monotone operator splitting problem. Using this connection, this paper lays the theoretical and practical foundations for such networks. We show how to parameterize networks in a manner that ensures all operators remain monotone, which establishes the existence and uniqueness of the equilibrium point. We show how to backpropagate through such networks using the implicit function theorem; this leads to a corresponding (linear) operator splitting problem for the backward pass, which also is guaranteed to have a unique solution. We then adapt traditional operator splitting methods, such as forward-backward splitting or Peaceman-Rachford splitting, to naturally derive algorithms for efficiently computing these equilibrium points.

Finally, we demonstrate how to practically implement such models and operator splitting methods, in the cases of typical feedforward, fully convolutional, and multi-scale convolutional networks. For convolutional networks, the most efficient fixed-point solution methods require an inversion of the associated linear operator, and we illustrate how to achieve this using the fast Fourier transform. The resulting networks show strong performance on several benchmark tasks, vastly improving upon the accuracy and efficiency of Neural ODEs-based models, the other implicit-depth models where solutions are guaranteed to exist and be unique.

Related work

There has been a growing interest in recent years in implicit layers in deep learning. Instead of specifying the explicit computation to perform, a layer specifies some condition that should hold at the solution to the layer, such as a nonlinear equality, or a differential equation solution. Using the implicit function theorem allows for backpropagating through the layer solutions analytically, making these layers very memory efficient, as they do not need to maintain intermediate iterations of the solution procedure. Recent examples include layers that compute inference in graphical models , solve optimization problems , execute model-based control policies , solve two-player games , solve gradient-based optimization for meta-learning , and many others.

Stability of fixed-point models

The issue of model stability has in fact been at the heart of much work in fixed-point models. The original work on attractor-style recurrent models, trained via recurrent backpropagation , precisely attempted to ensure that the forward iteration procedure was stable. And indeed, much of the work in recurrent architectures such as LSTMs has focused on these issues of stability . Recent work has revisited recurrent backpropagation in a similar manner to DEQs, with the similar aim of speeding up the computation of fixed points . And other work has looked at the stability of implicit models , with an emphasis on guaranteeing the existence of fixed points, but focused on alternative stability conditions, and considered only relatively small-scale experiments. Other recent work has looked to use control-theoretic methods to ensure the stability of implicit models, , though again they consider only small-scale evaluations.

Monotone operators in deep learning

Although most work in the field of monotone operators is concerned with general convex analysis, the recent work of does also highlight connections between deep networks and monotone operator problems. Unlike our current work, however, that work focused largely on the fact that many common non-linearities can be expressed via proximal operators, and analyzed traditional networks under the assumptions that certain of the operators were monotone, but did not address conditions for the networks to be monotone or algorithms for solving or backpropagating through the networks.

A monotone operator view of fixed-point networks

This section lays out our main methodological and theoretical contribution, a class of equilibrium networks based upon monotone operators. We begin with some preliminaries, then highlight the basic connection between the fixed point of an “infinite-depth” network and an associated operator splitting problem; next, we propose a parameterization that guarantees the associated operators to be maximal monotone; finally, we show how to use operator splitting methods to both compute the fixed point and backpropagate through the fixed point efficiently.

Deep equilibrium models

The monDEQ architecture is closely relate to the DEQ model, which parameterizes a “weight-tied, input-injected” network of the form $z_{i+1}=g(z_{i},x)$, where $x$ denotes the input to the network, injected at each layer; $z_{i}$ denotes the hidden layer at depth $i$; and $g$ denotes a nonlinear function which is the same for each layer (hence the network is weight-tied). The key aspect of the DEQ model is that in this weight-tied setting, instead of forward iteration, we can simply use any root-finding approach to find an equilibrium point of such an iteration $z^{*}=g(z^{*},x)$. Assuming the model is stable, this equilibrium point corresponds to an “infinite-depth fixed point” of the layer. The monDEQ architecture can be viewed as an instance of a DEQ model, but one that relies on the theory of monotone operators, and a specific paramterization of the network weights, to guarantee the existence of a unique fixed point for the network. Crucially, however, as is the case for DEQs, naive forward iteration of this model is not necessarily stable; we therefore employ operator splitting methods to develop provably (linearly) convergent methods for finding such fixed points.

2 Fixed-point networks as operator splitting

We begin by observing that it is possible to characterize this equilibrium point exactly as the solution to a certain operator splitting problem, under certain choices of operators and activation $\sigma$. This can be formalized in the following theorem, which we prove in Appendix B:

Finding a fixed point of the iteration (1) is equivalent to finding a zero of the operator splitting problem $0\in(F+G)(z^{\star})$ with the operators

and $\sigma(\cdot)=\operatorname{prox}^{1}_{f}(\cdot)$ for some convex closed proper (CCP) function $f$, where $\operatorname{prox}^{\alpha}_{f}$ denotes the proximal operator

It is also well-established that many common nonlinearities used in deep networks can be represented as proximal operators of CCP functions . For example, the ReLU nonlinearity $\sigma(x)=[x]_{+}$ is the proximal operator of the indicator of the positive orthant $f(x)=I\{x\geq 0\}$, and tanh, sigmoid, and softplus all have close correspondence with proximal operators of simple expressions .

In fact, this method establishes that some seemingly unstable iterations can actually still lead to convergent algorithms. ReLU activations, for instance, have traditionally been avoided in iterative models such as recurrent networks, due to exploding or vanishing gradient problems and nonsmoothness. Yet this iteration shows that (with input injection and the above constraint on $W$), ReLU operators are perfectly well-suited to these fixed-point iterations.

3 Enforcing existence of a unique solution

The above connection is straightforward, but also carries interesting implications for deep learning. Specifically, we can establish the existence and uniqueness of the equilibirum point $z^{\star}$ via the simple sufficient criterion that $I-W$ is strongly monotone, or in other wordsFor non-symmetric matrices, which of course is typically the case with $W$, positive definiteness is defined as the positive definiteness of the symmetric component $I-W\succeq mI\Leftrightarrow I-(W+W^{T})/2\succeq mI$. $I-W\succeq mI$ for some $m>0$ (see Appendix A). The constraint is by no means a trivial condition. Although many layers obey this condition under typical initialization schemes, during training it is normal for $W$ to move outside this regime. Thus, the first step of the monDEQ architecture is to parameterize $W$ in such a way that it always satisfies this strong monotonicity constraint.

We therefore propose to simply parameterize $W$ directly in this form, by defining the $A$ and $B$ matrices directly. While this is an overparameterized form for a dense matrix, we could avoid this issue by, e.g. constraining $A$ to be lower triangular (making it the Cholesky factor of $A^{T}A$), and by making $B$ strictly upper triangular; in practice, however, simply using general $A$ and $B$ matrices has little impact upon the performance of the method. The parameterization does notably raise additional complications when dealing with convolutional layers, but we defer this discussion to Section 4.2.

4 Computing the network fixed point

Given the monDEQ formulation, the first natural question to ask is: how should we compute the equilibrium point $z^{\star}=\sigma(Wz^{\star}+Ux+b)$? Crucially, it can be the case that the simple forward iteration of the network equation (1) does not converge, i.e., the iteration may be unstable. Fortunately, monotone operator splitting leads to a number of iterative methods for finding these fixed points, which are guaranteed to converge under proper conditions. For example, the forward-backward iteration applied to the monotone operator formulation from Theorem 1 results exactly in a damped version of the forward iteration

This iteration is guaranteed to converge linearly to the fixed point $z^{\star}$ provided that $\alpha\leq 2m/L^{2}$, when the operator $I-W$ is Lipschitz and strongly monotone with parameters $L$ (which is simply the operator norm $\|I-W\|_{2}$) and $m$ .

A key advantage of the monDEQ formulation is the flexibility to employ alternative operator splitting methods that converge much more quickly to the equilibrium. One such example is Peaceman-Rachford splitting which, when applied to the formulation from Theorem 1, takes the form

where we use the explicit form of the resolvents for the two monotone operators of the model. The advantage of Peaceman-Rachford splitting over forward-backward is two-fold: 1) it typically converges in fewer iterations, which is a key bottleneck for many implicit models; and 2) it converges for any $\alpha>0$ , unlike forward-backward splitting which is dependent on the Lipschitz constant of $I-W$. The disadvantage of Peaceman-Rachford splitting, however, is that it requires an inverse involving the weight matrix $W$. It is not immediately clear how to apply such methods if the $W$ matrix involves convolutions or multi-layer models; we discuss these points in Section 4.2. A summary of these methods for computation of the forward equilibrium point is given in Algorithms 1 and 2.

5 Backprogation through the monotone operator layer

Finally, a key challenge for any implicit model is to determine how to perform backpropagation through the layer. As with most implicit models, a potential benefit of the fixed-point conditions we describe is that, by using the implicit function theorem, it is possible to perform backpropagation without storing the intermediate iterates of the operator splitting algorithm in memory, and instead backpropagating directly through the equilibrium point.

To begin, we present a standard approach to differentiating through the fixed point $z^{\star}$ using the implicit function theorem. This formulation has some compelling properties for monDEQ, namely the fact that this (sub)gradient will always exist. When training a network via gradient descent, we need to compute the gradients of the loss function

where $(\cdot)$ denotes some input to the layer or parameters, i.e. $W$, $x$, etc. The challenge here is computing (or left-multiplying by) the Jacobian $\partial z^{\star}/\partial(\cdot)$, since $z^{\star}$ is not an explicit function of the inputs. While it would be possible to simply compute gradients through the “unrolled” updates, e.g. $z^{k+1}=\sigma(Wz^{k}+Ux+b)$ for forward iteration, this would require storing each intermediate state $z^{k}$, a potentially memory-intensive operation. Instead, the following theorem gives an explicit formula for the necessary (sub)gradients. We state the theorem more directly in terms of the operators mentioned Theorem 1; that is, we use $\operatorname{prox}_{f}^{1}(\cdot)$ in place of $\sigma(\cdot)$.

For the equilibrium point $z^{\star}=\operatorname{prox}_{f}^{1}(Wz^{\star}+Ux+b)$, we have

denotes the Clarke generalized Jacobian of the nonlinearity evaluated at the point $Wz^{\star}+Ux+b$. Furthermore, for the case that $(I-W)\succeq mI$, this derivative always exists.

To apply the theorem in practice to perform reverse-mode differentiation, we need to solve the system

The above system is a linear equation and while it is typically computationally infeasible to compute the inverse $(I-JW)^{-T}$ exactly, we could compute a solution to $(I-JW)^{-T}v$ using, e.g., conjugate gradient methods. However, we present an alternative formulation to computing (11) as the solution to a (linear) monotone operator splitting problem:

where $D$ is a diagonal matrix defined by $J=(I+D)^{-1}$ (where we allow for the possibility of $D_{ii}=\infty$ for $J_{ii}=0$).

An advantage of this approach when using Peaceman-Rachford splitting is that it allows us to reuse a fast method for multiplying by $(I+\alpha(I-W))^{-1}$ which is required by Peaceman-Rachford during both the forward pass (equilibrium solving) and backward pass (backpropagation) of training a monDEQ. Algorithms detailing both the Peaceman-Rachford and forward-backward solvers for the backpropagation problem (14) are given in Algorithms 3 and 4.

Example monotone operator networks

With the basic foundations from the previous section, we now highlight several different instantiations of the monDEQ architecture. In each of these settings, as in Theorem 1, we will formulate the objective as one of finding a solution to the operator splitting problem $0\in(F+G)(z^{\star})$ for

or equivalently as computing an equilibrium point $z^{\star}=\operatorname{prox}_{f}^{1}(Wz^{\star}+Ux+b)$.

In each of these settings we need to define what the input and hidden state $x$ and $z$ correspond to, what the $W$ and $U$ operators consist of, and what is the function $f$ which determines the network nonlinearity. Key to the application of monotone operator methods are that 1) we need to constrain the $W$ matrix such that $I-W\succeq mI$ as described in the previous section and 2) we need a method to compute (or solve) the inverse $(I+\alpha(I-W))^{-1}$, needed e.g. for Peaceman-Rachford; while this would not be needed if using only forward-backward splitting, we believe that the full power of the monotone operator view is realized precisely when these more involved methods are possible.

We can form an inverse directly by simply forming and inverting the matrix $I+\alpha(I-W)$, which has cost $O(n^{3})$. Note that this inverse needs to be formed only once, and can be reused over all iterations of the operator splitting method and over an entire batch of examples (but recomputed, of course, when $W$ changes). Any proximal function can be used as the activation: for example the ReLU, though as mentioned there are also close approximations to the sigmoid, tanh, and softplus.

2 Convolutional networks

The benefit of convolutional operators in this setting is the ability to perform efficient inversion via the fast Fourier transform. Specifically, in the case that $A$ and $B$ represent circular convolutions, we can reduce the matrices to block-diagonal form via the discrete Fourier transform (DFT) matrix

The inner term here is itself a block diagonal matrix with complex $n\times n$ blocks (each block is also guaranteed to be invertible by the same logic as for the full matrix). Thus, we can multiply a set of hidden units $z$ by the inverse of this matrix by simply inverting each $n\times n$ block, taking the fast Fourier transform (FFT) of $z$, multiplying each corresponding block of $F_{s}z$ by the corresponding inverse, then taking the inverse FFT. The details are given in Appendix C.

The computational cost of multiplying by this inverse is $O(n^{2}s^{2}\log s+n^{3}s^{2})$ to compute the FFT of each convolutional filter and precompute the inverses, and then $O(bns^{2}\log s+bn^{2}s^{2})$ to multiply by the inverses for a set of hidden units with a minibatch of size $b$. Note that just computing the convolutions in a normal manner has cost $O(bn^{2}s^{2})$, so that these computations are on the same order as performing typical forward passes through a network, though empirically 2-3 times slower owing to the relative complexity of performing the necessary FFTs.

One drawback of using the FFT in this manner is that it requires that all convolutions be circular; however, this circular dependence can be avoided using zero-padding, as detailed in Section C.2.

3 Forward multi-tier networks

where $z_{i}$ denotes the hidden units at level $i$, an $s_{i}\times s_{i}$ resolution hidden unit with $n_{i}$ channels, and where $W_{ii}$ denotes an $n_{i}$ channel to $n_{i}$ channel convolution, and $W_{i+1,i}$ denotes an $n_{i}$ to $n_{i+1}$ channel, strided convolution. This structure of $W$ allows for both inter- and intra-tier influence.

One challenge is to ensure that we can represent $W$ with the form $(1-m)I-A^{T}A+B-B^{T}$ while still maintaining the above structure, which we achieve by parameterizing each $W_{ij}$ block appropriately. Another consideration is the inversion of the multi-tier operator, which can be achieved via the FFT similarly as for single-convolutional $W$, but with additional complexity arising from the fact that the $A_{i+1,i}$ convolutions are strided. These details are described in Appendix D.

Experiments

To test the expressive power and training stability of monDEQs, we evaluate the monDEQ instantiations described in Section 4 on several image classification benchmarks. We take as a point of comparison the Neural ODE (NODE) and Augmented Neural ODE (ANODE) models, the only other implicit-depth models which guarantee the existence and uniqueness of a solution. We also assess the stability of training standard DEQs of the same form as our monDEQs.

The training process relies upon the operator splitting algorithms derived in Sections 3.4 and 3.5; for each batch of examples, the forward pass of the network involves finding the network fixed point (Algorithm 1 or 2), and the backward pass involves backpropagating the loss gradient through the fixed point (Algorithm 3 or 4). We analyze the convergence properties of both the forward-backward and Peaceman-Rachford operator splitting methods, and use the more efficient Peaceman-Rachford splitting for our model training. For further training details and model architectures see Appendix E. Experiment code can be found at http://github.com/locuslab/monotone_op_net.

We train small monDEQs on CIFAR-10 , SVHN , and MNIST , with a similar number of parameters as the ODE-based models reported in and . The results (averages over three runs) are shown in Table 2. Training curves for monDEQs, NODE, and ANODE on CIFAR-10 are show in Figure (2) and additional training curves are shown in Figure F1. Notably, except for the fully-connected model on MNIST, all monDEQs significantly outperform the ODE-based models across datasets. We highlight the performance of the small single convolution monDEQ on CIFAR-10 which outperforms Augmented Neural ODE by 15.1%.

We also attempt to train standard DEQs of the same structure as our small multi-tier convolutional monDEQ. We train DEQs both with unconstrained $W$ and with $W$ having the monotone parameterization (5), and solve for the fixed point using Broyden’s method as in . All models quickly diverge during the first few epochs of training, even when allowed 300 iterations of Broyden’s method.

Additionally, we train two larger monDEQs on CIFAR-10 with data augmentation. The strong performance (89% test accuracy) of the multi-tier network, in particular, goes a long way towards closing the performance gap with traditional deep networks. For comparison, we train larger NODE and ANODE models with a comparable number of parameters (~1M). These attain higher test accuracy than the smaller models during training, but diverge after 10-30 epochs (see Figure F1).

Efficiency of operator splitting methods

We compare the convergence rates of Peaceman-Rachford and forward-backward splitting on a fully trained model, using a large multi-tier monDEQ trained on CIFAR-10. Figure 3 shows convergence for both methods during the forward pass, for a range of $\alpha$. As the theory suggests, the convergence rates depend strongly on the choice of $\alpha$. Forward-backward does not converge for $\alpha>0.125$, but convergence speed varies inversely with $\alpha$ for $\alpha<0.125$. In contrast, Peaceman-Rachford is guaranteed to converge for any $\alpha>0$ but the dependence is non-monotonic. We see that, for the optimal choice of $\alpha$, Peaceman-Rachford can converge much more quickly than forward-backward. The convergence rate also depends on the Lipschitz parameter $L$ of $I-W$, which we observe increases during training. Peaceman-Rachford therefore requires an increasing number of iterations during both the forward pass (Figure 2) and backward pass (Figure F2).

Finally, we compare the efficiency of monDEQ to that of the ODE-based models. We report the time and number of function evaluations (OED solver steps or operator splitting iterations) required by the ~170k-parameter models to train on CIFAR-10 for 40 epochs. The monDEQ, neural ODE, and ANODE training takes respectively 1.4, 4.4, and 3.3 hours, with an average of 20, 96, and 90 function evals per minibatch. Note however that training the larger 1M-parameter monDEQ on CIFAR-10 requires 65 epochs and takes 16 hours. All experiments are run on a single RTX 2080 Ti GPU.

Conclusion

The connection between monotone operator splitting and implicit network equilibria brings a new suite of tools to the study of implicit-depth networks. The strong performance, efficiency, and guaranteed stability of monDEQ indicate that such networks could become practical alternatives to deep networks, while the flexibility of the framework means that performance can likely be further improved by, e.g. imposing additional structure on $W$ or employing other operator splitting methods. At the same time, we see potential for the study of monDEQs to inform traditional deep learning itself. The guarantees we can derive about what architectures and algorithms work for implicit-depth networks may give us insights into what will work for explicit deep networks.

Broader impact statement

While the main thrust of our work is foundational in nature, we do demonstrate the potential for implicit models to become practical alternatives to traditional deep networks. Owing to their improved memory efficiency, these networks have the potential to further applications of AI methods on edge devices, where they are currently largely impractical. However, the work is still largely algorithmic in nature, and thus it is much less clear the immediate societal-level benefits (or harms) that could result from the specific tehniques we propose and demonstrate in this paper.

Acknowledgements

Ezra Winston is supported by a grant from the Bosch Center for Artificial Intelligence.

References

Appendix A Monotone operator theory

In the case of scalar-valued functions, this corresponds to our common notion of a monotonic function. The operator $F$ is strongly monotone with parameter $m$ if

The resolvent and Cayley operators for an operator $F$ are denoted $R_{F}$ and $C_{F}$ and respectively defined as

for any $\alpha>0$. The resolvent and Cayley operators are non-expansive (i.e., have Lipschitz constant $L\leq 1$) for any maximal monotone $F$, and are contractive (i.e. $L<1$) for strongly monotone $F$.

We will mainly use two well-known properties of these operators. First, when $F(x)=Gx+h$ is linear, then

and when $F=\partial f$ for some CCP function $f$, then the resolvent is given by a proximal operator

Operator splitting approaches refer to methods to find a zero in a sum of operators (assumed here to be maximal monotone), i.e., find $x$ such that

There are many such operator splitting methods, which lead to different approaches in their application to our subsequent implicit networks, but the two we use mainly in this work are 1) forward-backward splitting, given by the update

and 2) Peaceman-Rachford splitting, which is given by the iteration

Both methods will converge linearly to an $x$ that is a zero of the operator sum under certain conditions: a sufficient condition for forward-backward to converge is that $F$ be strongly monotone with parameter $m$ and Lipschitz with constant $L$ and $\alpha<2m/L^{2}$; for Peaceman-Rachford, the method will converge for any choice of $\alpha$ for strongly monotone $F$, though the convergence speed will often vary substantially based upon $\alpha$.

Appendix B Proofs

The proof here is immediate: the forward-backward algorithm applied to the above operators with $\alpha=1$ corresponds exactly to the network’s fixed-point iteration:

B.2 Proof of Proposition 5

First assume $W$ is of this form. Then clearly

Alternatively, if $I-W\succeq mI\Longleftrightarrow(1-m)I\succeq(W+W^{T})/2$, then

B.3 Proof of Theorem 2

Differentiating both sides of the fixed-point equation $z^{\star}=\sigma(Wz^{\star}+Ux+b)$ we have

for $J$ defined in (10) (we require the Clarke generalized Jacobian owing to the fact that the nonlinearity need not be a smooth function). Rearranging we get

To show that this derivative always exists, we need to show that the $I-JW$ matrix is nonsingular. Owing to the fact that proximal operators are monotone and non-expansive, we have $0\leq J_{ii}\leq 1$. First, letting $\lambda(\cdot)$ denote the set of eigenvalues of a matrix, note that

This follows from the similarity transform $\lambda(I-JW)=\lambda(J^{-1/2}(I-JW)J^{1/2})$ for $J>0$ and the case of $J_{ii}=0$ follows via the continuity of eigenvalues taking $\lim J_{ii}\rightarrow 0$. Now, using the fact that $0\preceq J\preceq I$, we have

since $I-W\succeq mI$ and $I-J\succeq 0$. ∎

B.4 Proof of Theorem 3

We begin with the case where $J_{ii}\neq 0$ and thus $D_{ii}<\infty$. As above, because proximal operators are themselves monotone non-expansive operators, we always $0\leq J_{ii}\leq 1$, so that $D_{ii}\geq 0$. Now, first assuming that $J_{ii}>0$, and hence $D_{ii}<\infty$, we have

Thus, we can always solve the above equation with the $v$ term of the form $W^{T}Jv$, giving

To handle the case where $J_{ii}=0\Leftrightarrow D_{ii}=\infty$, we can simply take the limit $D_{ii}\rightarrow\infty$, and note that all the operators are well-defined for this case. For instance, the resolvent operator

Appendix C Convolutional monDEQs

It is more efficient to consider the permuted form of $D$

To perform the required inversion of the operator

we use the fact that $\mathscr{F}_{s,n}$ is unitary and obtain

The inner term here itself has the blockwise-diagonal form (C2). Thus, we can multiply a set of hidden units $z$ by the inverse of this matrix by considering the permuted form (C3), inverting each block $\hat{D}^{i}$, taking the FFT of $z$, multiplying each corresponding block of $\mathscr{F}_{s,n}z$ by the corresponding inverse, then taking the inverse FFT.

C.2 Zero padding

One drawback to the above method is that using the FFT in this manner requires that all convolutions be circular. While empirically there is little drawback to simply replacing traditional convolutions with their circular variants, in some cases it may be desirable to avoid this setting, where information about the image may wrap around the borders. If it is desirable to avoid this, we explicitly remove any circular dependence by zero-padding the hidden units with $(k-1)/2$ border pixels, where $k$ denotes the receptive field size of the convolution. This zero padding can then be enforced by simply setting all the border entries to zero within the nonlinearity of the network; because setting an element to zero is equivalent to the proximal operator for the indicator of the zero set, such operations still fit within the monotone operator setting.

Appendix D Multi-tier monDEQs

To ensure $W$ has the form $(1-m)I-A^{T}A+B-B^{T}$, we restrict both $A$ and $B$ to have the same bidiagonal structure as $W$. Then the diagonal terms $W_{ii}$ have the form

To compute the off-diagonal terms $W_{i+1,i}$ note that restricting $W$ to be bidiagonal makes the off-diagonal terms of $B$ redundant. E.g. since $W_{12}=0$, then

D.2 Inversion via the discrete Fourier transform

Consider $W$ of the form (D1) with convolutions

Here the $A_{ii}$ and $B_{ii}$ terms are unstrided convolutions with $n_{i}$ input and $n_{i}$ output channels, while the $A_{i,i+1}$ are strided convolutions with $n_{i}$ input channels and $n_{i+1}$ output channels.

In order to multiply by $(I+\alpha(I-W))^{-1}$, we use back substitution to solve for $x$ in

Let $W^{\prime}=(I+\alpha(I-W))$. The back substitution proceeds by tiers, i.e.

Therefore only the diagonal blocks $W^{\prime}_{ii}$ need be inverted. The inversion of e.g.

is complicated by the fact that $A_{21}$ is strided, so that it is no longer diagonalized by the DFT. Instead, we perform inversion using the following proposition.

The desired result then follows by applying the Woodbury matrix idenetity.

We start by breaking $B$ into an unstrided convolution $B^{\prime}$ which can be diagonalized by the DFT and a matrix $U_{r,s}$ which performs the striding on each channel:

We want to show that $\mathscr{F}_{s}U_{r,s}^{T}U_{r,s}\mathscr{F}_{s}^{*}=\frac{1}{s^{2}r^{2}}S^{T}JJ^{T}S$. Observe that

Then by the properties of Kronecker product

We now show that $(F_{s}T_{r,s}F_{s}^{*})=L$ where

To do so we use several properties of the roots of unity $z^{k}=\exp(2\pi\iota k/s)$.

If $a\equiv b~{}(\text{mod }s)$ then $z^{a}=z^{b}$.

If $z$ is a primitive $s$th root of unity then $z^{m}$ is a primitive $a$th root of unity where $a=\frac{s}{\text{gcd}(m,s)}$.

The sum of the $s$th roots of unity $\sum_{k=0}^{s-1}z^{k}=0$ if $s>1$.

We first compute $L_{ij}$ for the case when $i\equiv j~{}(\text{mod }s/r)$, or in other words $i=j+\frac{ks}{r}$ for some integer $k$. We have

For the case when $i\not\equiv j~{}(\text{mod }s/r)$, or in other words $i=j+\frac{ks}{r}+m$ for some integers $k$ and $m$ with $-\frac{s}{r}<m<\frac{s}{r}$, we have

By property (2), since $\exp(2\pi\iota r/s)$ is a primitive $\frac{s}{r}$th root of unity, then $\exp(2\pi\iota mr/s)$ is a primitive $d$th root of unity where $d=\frac{s/r}{\text{gcd}(m,s/r)}$. Since $d$ divides $s/r$, we can split the sum into several sums of $d$th roots of unity using property (1), each of which will sum to zero by property (3).

where the second equality follows from property (1) since $p=p+qd~{}(\text{mod }d)$ and each sum in the third line is zero by property (3) since $\exp(2\pi\iota mr/s)$ is a primitive $d$th root of unity.

where $S_{n,m}$ is the perfect shuffle matrix. Note that $L=\frac{1}{sr}1_{r\times r}\otimes I_{s/r}$ where $1_{r\times r}$ is the $r\times r$ matrix of all ones. Then

Appendix E Experiment details

Recall that a monDEQ is defined by a choice of linear operators $W$ and $U$, bias $b$, and nonlinearity $\sigma$, and that we parameterize $W$ via linear operators $A$ and $B$. For all experiments we use $\sigma=\text{ReLU}$. In the fully-connected network $A,B$ and $U$ are dense matrices; in the single-convolution network they are unstrided convolutions with kernel size 3. The structure of the multi-tier network is as described in (D1) and (D5); we use three tiers with unstrided convolutions for $U$ and $A_{ii},B_{ii}$ and stride-2 convolutions for the subdiagonal terms $A_{i,i+1}$, all with kernels of size 3. The number of channels for single and multi-tier convolutional models varies by dataset, as shown in Table E1.

For all models, the fixed point $z^{\star}$ is mapped to logits $\hat{y}$ via a dense output layer, and the single convolution model first applies 4$\times$4 average pooling:

E.2 Training details

Because $W=(1-m)I-A^{T}A+B-B^{T}$ contains both linear and quadratic terms, we find that a variant of weight normalization helps to keep the gradients of the different parameters on the same scale. For example, when $W$ is a dense matrix, we reparameterize $A^{T}A$ as $g\frac{A^{T}A}{\|A\|^{2}}$ and $B$ as $h\frac{B}{\|B\|}$, where $g$ and $h$ are learned scalars. When $W$ consists of a single or multi-tiered convolutions, we reparameterize each convolution kernel analogously.

All models are trained by running Peaceman-Rachford with error tolerence $\epsilon=$1e-2, which reduces the number of iterations without impacting performance. The monotonicity parameter $m$ also affects convergence speed since it controls the contraction factor of the relevant operators; consistent with this, we find that Peaceman-Rachford takes longer to converge for smaller $m$, and use $m=1$ for all models since model performance is not sensitive to $m\in[0.01,1]$. We also find that the Lipschitz parameter $L$ of $I-W$ increases during training, changing the optimal $\alpha$ value. We therefore tune $\alpha\in\{1,1/2,1/4,\ldots\}$ over the course of training so as to minimize forward-pass iterations.

One detail about stopping criteria for the splitting method: computing the residual $\|z^{k+1}-f(z^{k+1})\|/\|z^{k+1}\|$ requires an additional call to the function $f$. Therefore during training we instead use the criterion $\|z^{k+1}-z^{k}\|/\|z^{k+1}\|\leq\epsilon$. The error shown in Figure 3 is the former, while the stopping criterion used in Figures 2 and F2 is the latter. Technically this latter criterion itself depends on both $\alpha$ and $L$; for different $\alpha$ and $L$ values, having $\|z^{k+1}-z^{k}\|/\|z^{k+1}\|\leq\epsilon$ implies different bounds on the residual. However, we find that this effect is minimal, so that both stopping criteria work equally well in practice.

Table E1 gives details of the training hyperparameters used for each model. All models are trained with ADAM , using batch size of 128. For all but the large CIFAR-10 models, the initial learning rate is chosen from {1e-2, 1e-3} and decayed by a factor of 10 after every 10 or 25 epochs, and the default ADAM momentum parameters are used. All training data is normalized to mean $\mu=0$, standard deviation $\sigma=1$.

When training large models on CIFAR-10 we use standard data augmentation, consisting of zero-padding the 32$\times$32 images to 40$\times$40, then randomly cropping back to 32x32, and finally performing random horizontal flips. In order to reduce the number of training epochs, we use a single cycle of increasing and decreasing learning rate to achieve super-convergence . The learning rate is increased from 1e-3 to the max learning rate (see Table E1) over 30 epochs, then decreased back to 1e-3 over 30 epochs, then held at 1e-3 for 5 epochs. (The max learning rate is chosen by training for a single epoch while increasing the learning rate until the loss diverges.) The momentum is also decreased from 0.95 to 0.85 over 30 epochs, then back to 0.95 over 30 epochs, then held at 0.95 for 5 epochs. However, we note that the model obtains the same performance when trained with constant learning rate of 1e-3 for around 200 epochs.

E.3 Dataset statistics

MNIST consists of black and white examples of handwritten digits 0-9. SVHN consists of color images of digits 0-9 extracted from house numbers captured by Google Stree View. CIFAR-10 consists of small images from 10 object classes. Dataset statistics are shown in Table E2.

Appendix F Additional results and figures