D ECLARATIVE NETS THAT ARE EQUILIBRIUM MODELS
Artiﬁcial Intelligence
Future Science Platform ¶
A BSTRACT
Implicit layers are computational modules that output the solution to some problem depending on the input and the layer parameters.
Deep equilibrium models (DEQs) output a solution to a ﬁxed point equation.
Deep declarative networks (DDNs) solve an optimisation problem in their forward pass, an arguably more intuitive, interpretable problem than ﬁnding a ﬁxed point.
We show that solving a kernelised regularised maximum likelihood estimate as an inner problem in a
DDN yields a large class of DEQ architectures.
Our proof uses the exponential family in canonical form and provides a closed-form expression for the DEQ parameters in terms of the kernel.
The activation functions have interpretations in terms of the derivative of the log partition function.
Building on existing literature, we interpret DEQs as ﬁne-tuned, unrolled classical algorithms, giving an intuitive justiﬁcation for why DEQ models are sensible.
We use our theoretical result to devise an initialisation scheme for DEQs that allows them to solve kGLMs in their forward pass at initialisation.
We empirically show that this initialisation scheme improves training stability and performance over random initialisation
I NTRODUCTION
Implicit layers (Pineda, 1987; Almeida, 1990) have recently been subject to renewed attention (Kolter et al., 2020).
In contrast with explicitly deﬁned layers, implicit layers deﬁne a mapping in terms of a solution to some problem depending on the input and problem parameters.
For example, deep equilibrium models (DEQs) consist of layers that output ﬁxed points of parameterised functions (Bai et al., 2019).
Deep declarative networks (DDNs) use declarative nodes which output solutions to optimisation problems (Gould et al., 2016;
Amos & Kolter, 2017).
Neural ordinary differential equations (NODEs) output solutions to ODEs (Chen et al., 2018; Dupont et al., 2019).
Traditional explicit layers can always be represented as implicit layers (for example, see Proposition 4.10 of Gould et al.
Also, solutions to certain convex optimisation problems may be obtained via an iterative optimisation procedure such as Newton’s method or gradient descent, and as such, may be represented as ﬁxed points of an iterative scheme.
A correspondence between DDNs and DEQs is expected (but undiscovered), given the fundamental connection between ﬁxed points of iterative maps and critical points of optimisation problems (Ryu & Boyd, 2016).
This leads to two natural questions: (1) For a given optimisation problem, what is the corresponding DEQ architecture?
(2) Can this correspondence be exploited for theoretical, conceptual, or practical beneﬁt?
Contributions.
(1) We prove an equivalence between a DEQ and a DDN with a classical statistical model — a kernelised generalised linear model (kGLM) — as the declarative node, as illustrated in Figure 1 and formally stated in Theorem 3 and Corollary 6.
The weights of the DEQ layer have closed-form expressions in terms of the kernel.
The surprise in our result is that the feature mapping involved in this correspondence is exactly the class of hidden layers that are most commonly used in practice.
(2) We empirically demonstrate that initialising a DEQ as a DDN using the derived expression for the weights improves performance and stability over random parameter initialisation.
Implicit network module α ∗ = arg min kGLM ( α, X, Y )
f ∗ ( Y ) =
k ( (cid:101) X,X )
α ∗ α ∗ =
arg min kGLM ( α, X, Y ) f ∗ ( Y ) =
k ( (cid:101) X,X ) α ∗ DDN z ∗ =
σ (cid:0) W 1 z ∗ + W 2 T ( Y ) + W 3 ρ ( z ∗ ) (cid:1)
f ∗ ( Y ) =
V 1 ( z ∗ + ρ ( z ∗ ))
+ V 2 T ( Y ) DEQ z ∗ = σ (cid:0) W 1 z ∗ + W 2 T ( Y ) + W 3 ρ ( z ∗ ) (cid:1) is equivalent to f ∗ ( Y ) =
V 1 ( z ∗ + ρ ( z ∗ ))
+ V 2 T ( Y )
F γ
L ( F γ ( Y ) , (cid:101) Y )
Y
X, (cid:101) X (cid:101) Y
Y
X, (cid:101) X (cid:101) Y
Notation.
Let X ⊆
R d x and Y ⊆
R be coordinate and target spaces.
Deﬁne a Y − valued stochastic process { y ( x, ω ) } x ∈ X indexed by x ∈ X with outcome ω from a sample space Ω .
Let X ∈ R n × d x be a matrix such that the i th row of X is some element x i ∈ X .
Similarly deﬁne (cid:101) X ∈ R (cid:101) n × d x .
Take n evaluations { y i } ni =
= { y ( x i , ω 0 ) } ni =1 from a realisation of the stochastic process for every i th row in X , and form a corresponding matrix Y .
Using the same realisation of the stochastic process, form (cid:101) Y with evaluations { y ( (cid:101) x
i , ω 0 ) } (cid:101) nj =1 for every i th row in (cid:101) X .
Call X
and Y the inner coordinate and target data, and (cid:101)
X and (cid:101) Y the outer coordinate and target data.
Example implication of result.
Consider image inpainting a single 32 × 32 colour image.
Take X = { 1 , . .
. , 32 } 2 × { 1 , 2 , 3 } , and X ∈ Z n × 3 to be a matrix consisting of (non-repeated) triples, where n < 32 2 × 3 .
Take (cid:101) X ∈ Z (cid:101) n × 3 to be the matrix consisting of remaining triples, such that n + (cid:101) n = 32 2 × 3 .
An incomplete image Y ∈ R n together with its missing values (cid:101) Y ∈ R (cid:101) n are jointly sampled from an image distribution.
We may use kernel ridge regression to produce a perpixel prediction of (cid:101) Y given X, Y and (cid:101) X .
This model represents a special case of a regularised kGLM with a closed-form predictor, but more generally an algorithm is required to compute the predictor— see Appendix A.
We now move from a single image to a set of images.
We wrap the kGLM with an outer minimisation loop over the loss L between a neural network F γ output applied to the kGLM predictor and target (cid:101) Y to obtain a DDN, as sketched in Figure 1.
We view the DDN as producing a per-image prediction .
We show that such a DDN is equivalent to a DEQ with a closed-form expression for the parameters depending on the kernel.
When the kGLM has ﬁxed hyperparameters, the DEQ has a ﬁxed set of parameters (made precise in Theorem 3 and Corollary 6).
Initialising the
DEQ with these parameters improves training stability and performance, as shown in § 4.
Another example using kernel logistic regression for image segmentation is given in Appendix G.
[Figure]
D EEP EQUILIBRIUM MODELS
Given some Y ∈ R n × d y containing row elements y ∈ Y , d eep eq uilibrium models (DEQs) (Bai et al., 2019) compute an embedding z ∗ ∈
Z ⊂ R n × d z as a ﬁxed point of some function and then output an afﬁne transformation f ∗ ( Y ) of the ﬁxed point.
More concretely, letting g θ ( · ; Y ) :
R n × d z → R n × d z denote a parameterised function (e.g. a neural network) with parameters θ ∈ Θ ,
[Equation]
where V ∈ R d u × ( n × d z ) are referred to as readout parameters, and V may be thought of as a d u × nd z matrix acting on a ﬂattened vector with nd z entries.
We will ﬁnd it more natural to introduce a linear skip connection between Y and the output.
We will also choose d y =
d z = 1 .
That is, we consider
[Equation]
Typically DEQs are trained in a supervised manner by running some iterative algorithm such as gradient descent on a loss that aims to match a neural network involving f ∗ ( Y ) to some target (cid:101) Y with respect to the DEQ and neural network parameters.
Given a dataset { ( Y s , (cid:101) Y s ) } Ns =1 , this procedure may be realised by running an iterative algorithm over the equality constrained optimisation problem
[Equation]
subject to
[Equation]
as if the algorithm were to solve the problem, even though it may be highly non-convex.
Here
L is some loss function and F γ is a neural network with parameters γ .
A large class of iterative algorithms requires the derivative with respect to the network parameters θ and V 1 , V 2 ; these may be computed under mild conditions without backpropagating through the ﬁxed point solution via implicit differentiation.
We refer the reader to Bai et al.
for details.
Dealing with existence and uniqueness of ﬁxed points.
In order for (1) and (2) to specify useful computational rules, they have to admit at least one ﬁxed point.
To avoid having to choose which ﬁxed point should be returned with additional rules, it might be desirable that they admit exactly one ﬁxed point.
We give example conﬁgurations of ﬁxed points in Figure 2.
The classical Banach ﬁxed point theorem (BFPT) gives sufﬁcient conditions for the existence and uniqueness of ﬁxed points.
Theorem 1 (BFPT) .
Let ( Z , d ) be a complete metric space.
A function H : Z → Z is said to be a contraction if there exists some 0 ≤ q < 1 such that for all z, z
(cid:48) ∈
Z , d ( H ( z ) , H ( z (cid:48) )) ≤ qd ( z, z (cid:48) ) .
Every contraction admits a unique ﬁxed point.
The contraction property is suggestive of a simple algorithm (Hasselblatt & Katok, 2003) for ﬁnding the ﬁxed point z ∗ of a contraction H given an initial guess z 0 : While some termination condition is not met, Update the current estimate z r for the ﬁxed point to be H ( z r − 1 ) .
This algorithm allows one to interpret the inner problem of (3) as an inﬁnitely deep neural network with weights shared between each layer.
More advanced algorithms are available, which we do not discuss here.
H is a contraction if the induced matrix norm (cid:107) · (cid:107) of the Jacobian of H is strictly less than 1 on Z , providing a useful test for checking whether a unique ﬁxed point of H exists.
More generally,
Proposition 2.
Let Z be an open strictly convex set, Z its closure, H : Z → Z differentiable on Z and continuous on
Z .
If (cid:107) DH (cid:107) ≤ λ < 1 on Z , then H is a contraction on Z .
[Figure]
We provide a proof in Appendix B. While elegant, contractions are stronger than necessary and using them limits the set of admissible networks.
It is therefore desirable to work with tighter conditions, or otherwise sidestep the issue.
A consensus on how to deal with existence and uniqueness does not yet appear to have been reached.
One may run the algorithm as if a unique ﬁxed point exists (Bai et al., 2019), constrain or monitor the parameters during learning to ensure that a unique ﬁxed point exists, add a penalisation term on an estimate of the spectral norm of the Jacobian (Roosta-Khorasani & Ascher, 2015) to encourage a contractive map (Bai et al., 2021), or design variants of DEQs (through appropriate restriction) that explicitly ensure that a unique ﬁxed point exists (Winston &
Kolter, 2020; Revay et al., 2020; El Ghaoui et al., 2021).
We do not address this problem; we assume (in a precise sense, see Assumption 2) veriﬁable conditions for the BFPT.
Initialisation.
For the vast majority of practical neural network architectures, learning involves applying an optimisation procedure (e.g. stochastic gradient descent or quasi-Newton’s method) to an empirical risk minimisation problem in the overparameterised network’s parameters outside of conditions that guarantee convergence to a minimum.
Debate continues as to why such procedures lead to good generalisation performance (Belkin et al., 2019; Zhang et al., 2021).
However, there are intuitive, empirical and theoretical reasons to suggest that appropriate parameter initialisation plays an important role (Glorot & Bengio, 2010;
He et al., 2015; Poole et al., 2016; Hu et al., 2020).
For DEQs, initialisations have not yet been studied in great detail.
By interpreting DEQs as unrolled kGLMs, we ﬁnd parameter initialisations as a corollary of our main theoretical result.
See Figure 3.
E XACT FIXED PARAMETER EQUIVALENCE BETWEEN DDN S AND DEQ S
We present our main result here and give a special case of a tractable Ising model in Appendix C.2.
Our main result is a derivation of a fully connected or convolutional DEQ architecture as a model that is equivalent to a DDN that solves a kGLM as its inner problem.
3. S ETTING : A K GLM INSIDE A DDN
Inner problem.
We work with the regular (Wainwright & Jordan, 2008) exponential family in canonical form (Deisenroth et al., 2020, Section 6.6), with log partition function A : R → R and sufﬁcient statistic T :
Y → R .
In our setting, A is both inﬁnitely differentiable due to regularity, and strictly convex due to minimality (see Appendix A).
We choose a prior p λ over the predictor f , which generalises the commonly used Gaussian process prior 1 (Sch¨olkopf et al., 2001; Canu & Smola, 2006).
For log-concave and differentiable q , deﬁne
[Equation]
A MAP f ∗ ( · ; X, Y ) =
k ( · , X ) α ∗ with representer coefﬁcients α ∗ satisﬁes
[Equation]
where we have used strict convexity of A to ensure at most one minimum exists.
For convenience, we deﬁne ρ : ≡ ( − log q ) (cid:48) ◦
( A (cid:48) )
− 1 .
When q is constant, p λ is a Gaussian process prior and ρ ≡ 0 .
Deep declarative network.
We will use (5) as a declarative node (Gould et al., 2021) by wrapping it with an outer optimisation task in order to learn the parameters of a neural network by empirical risk minimisation.
Let L denote a loss function (for example, mean squared error), and let F γ deﬁne a generic neural network with parameters γ .
Form a set of N quartets { ( X s , Y s , (cid:101)
X s , (cid:101) Y s ) }
Ns =1 , for example a set of N images.
Deﬁne K s =
k ( X s , X s ) .
We consider the optimisation problem
[Equation]
subject to
[Equation]
The outer problem is in general non-convex and highly nonlinear in γ .
3. M AIN RESULT : DERIVING A FIXED - PARAMETER DEQ
Under our construction, the equivalence between the embeddings z ∗ of kGLMs and DEQs always holds.
Under the additional Assumption 1, the DEQ and the kGLM are equivalent predictive models.
Assumption 1.
The kernel matrix K is strictly positive deﬁnite.
We now state our main result, the proof of which is given in Appendix D.
Theorem 3.
Let α ∗ be the minimiser of one of the inner optimisation problems (5), inducing a function f ∗ ( · ; X, Y )
=
k ( · , X ) α ∗ .
Deﬁne z ∗ :=
A (cid:48)
( Kα ∗ ) .
Then z ∗ is a ﬁxed point satisfying
[Equation]
where W 1 = W 3 = − λ − 1 K and W 2 = λ − 1 K are parameter matrices, σ ≡ A (cid:48) and ρ ≡ ( − log q ) (cid:48) ◦
( A (cid:48) )
− 1 are monotone non-decreasing.
Furthermore, under Assumption 1, for any test index (cid:101) X ∈
R (cid:101) n × d x and for at least one z ∗ satisfying (7),
[Equation]
where V 1 = − λ − 1 k ( (cid:101) X, X ) and V 2 = λ − 1 k ( (cid:101) X, X ) .
Further conditions (Assumption 2) force the ﬁxed point iteration to have a unique solution.
Assumption 2.
Let Z be an open strictly convex set and Z its closure.
Suppose
[Equation]
Note that when ρ is zero, that is p λ is Gaussian, Assumption 2 reduces to (cid:107) K/λ (cid:107) sup z ∈Z | A (cid:48)(cid:48)
( z ) | < 1 .
The second derivative of the log partition is equal to 1 for the Gaussian distribution with known variance and is bounded by r/ 4 for the Binomial distribution with r trials, further simplifying Assumption 2.
Recall that the spectral radius (the largest eigenvalue) of K/λ being less than some constant is sufﬁcient to ensure the existence of an operator norm (cid:107) · (cid:107) such that (cid:107) K/λ (cid:107) is less than the same constant.
However, this guarantee does not identify the operator norm.
In our experiments, we instead ﬁx the operator norm as the spectral norm and check Assumption 2 against this norm.
Proposition 4.
Under Assumption 2, (7) admits a unique ﬁxed point on Z .
The outer problem depends on the solution α ∗ s to the inner problem only through f ∗ ( (cid:101)
X s ; X s , Y s ) , which under Assumption 1 in turn only depends on z ∗ s .
We therefore have the following.
Corollary 5.
Under Assumptions 1 and 2, the optimisation problem (6) is equivalent to the constrained optimisation problem
[Equation]
subject to
[Equation]
where V 1 s = − λ
− 1 k ( (cid:101)
X s , X s ) , V 2 s = λ − 1 k ( (cid:101)
X s , X s ) , W 1 s =
W 3 s = − λ − 1 K s and W 2 s =
λ
− 1 K
s take the role of DEQ parameters and σ = A (cid:48) and ρ are monotone non-decreasing and take the role of DEQ activation functions.
The ﬁxed point condition for z ∗ s is met by exactly one element of Z .
The model in Corollary 5 allows the parameters W 1 s , W 2 s , V 1 s and V 2 s to vary with s whereas (3) does not.
An additional assumption ensures these parameters are constant in s .
Corollary 6.
In the same setting as Corollary 5, if X s = X and (cid:101) X s = (cid:101)
X are constant in s , the parameters W 1 , W 2 , W 3 , V 1 , V 2 are constant in s and we may write f ∗ ( Y s )
≡ f ∗ s ( (cid:101)
X s ; X s , Y s ) .
Note that in (3), the outer problem is with respect to γ, θ, V 1 , V 2 whereas in (9) the outer problem is only with respect to γ .
The parameters V 1 , V 2 and W 1 , W 2 (which take the role of θ ), are automatically determined according to X ,
(cid:101) X and the kGLM kernel k .
Corollary 5 is suggestive of a DEQ model that we are yet to implement whose weights are a dynamic kernel function of the coordinate X s .
One parameterisation of such a model is to have an auxillary network that predicts a ﬁnite-dimensional feature mapping of the coordinates X s .
Finally, we note one more special case when X s =
(cid:101)
X s .
In this case, since K s = k ( (cid:101)
X s , X s ) , predictions may be formed as f ∗ ( (cid:101) X s ; X s , Y
[Equation]
) =
[Equation]
α ∗
[Equation]
= σ − 1 ( z ∗
[Equation]
=
W 1 s z ∗ s + W 2 s T ( Y s )
+
W 3 s ρ ( z ∗ s ) , without the need to invert K s and therefore without the need for Assumption 1.
3. R EMARKS ON MATCHING DDN S AND DEQ S
Special cases.
Choosing p K/λ to be Gaussian, ρ ≡ 0 .
Additionally, if T is the identity and σ ≡ A (cid:48) ≡ tanh , we obtain a fully connected layer with tanh activations, as derived in Appendix C.2.
Logistic sigmoidal activations are obtained when the kGLM is specialised as kernel logistic regression.
When A (cid:48) and T are the identity, the kGLM becomes kernel ridge regression, with Gaussian process regression as the corresponding Bayesian model.
Certain inﬁnitely wide (nonlinear) neural networks are Gaussian processes (Neal, 1995).
Our result shows that one may also construct Gaussian processes whose posterior predictive mean is an inﬁnitely deep neural network with linear activations.
The role of the coordinates X and (cid:101) X differs between the two models.
Symbol matching.
Theorem 3 says that an embedding (2) computed by a DEQ given Y is the same as an embedding (7) computed by a kGLM on coordinates X for a training set ( X, Y ) if the hidden parameters W 1 , W 2 and W 3 are scalar multiples of the kernel matrix k
( X, X ) .
Initialisation via optimisation warm-start.
In a stronger setting, Corollary 6 shows a connection between the two optimisation tasks (3) and (9).
Given a ﬁxed DEQ architecture with a ﬁxed point condition that may be expressed as (7) (and noting that this equation also includes convolutional architectures), we may initialise parameter values that ensure that the prediction is equivalent to a trained kGLM and identify the hyperparameters of the kGLM.
This initialisation may be naive (i.e. assume a “reasonable” kernel function, coordinates and regularisation parameter for a wide range of tasks) or more informed (i.e. leverage some information about the task to suit the kernel and regularisation to the task being performed).
Note that even without Assumptions 1 and 2, we may initialise the DEQ such that the feature embeddings z ∗ are equivalent for at least one of the ﬁxed points.
We demonstrate both naive and informed approaches in § 4.
On augmentation and assumptions.
Assumption 2 allows one to compute without concern using (7) since a unique ﬁxed point is guaranteed to exist.
DEQs are sometimes implemented without checking that a unique ﬁxed point exists.
This assumption may be replaced by any other assumption that is sufﬁcient to ensure a unique ﬁxed point.
Using the derivative test on the open set Z instead of Z allows us to handle important edge cases, namely when A (cid:48) ≡ tanh (with derivative 1 at the origin) or A (cid:48) is Leaky ReLU with gradients in { 0 , 1 } on both sides of the origin (with undeﬁned derivative at the origin).
Assumption 1 allows one to map the kGLM embedding to the prediction.
The skip connection present in (2) but not (1) is not a severe limitation from a practical perspective.
The augmented model considers a coordinate X absent from the original DEQ.
Note that any DEQ may be written as an augmented model by simply ignoring the coordinate X .
Under our construction but without Assumptions 1 or 2, the equivalence between the embeddings produced by kGLMs and DEQs is still valid for at least one of their respective ﬁxed points.
[Figure]
Convolutional layers.
The derived network architecture involves only matrix multiplication, however as is well-known, we may represent a convolution through a sparse matrix equation (Figure 4).
Conveniently, the local structure expected in images may be encoded by a particular informed kernel choice in order to recover a convolutional architecture.
We formally discuss how this may be done in 2 D RGB images in Corollary 11, Appendix E through a particular choice of kernel k .
Other priors.
One may replace the strongly log-concave prior p K/λ with, for example, the conjugate prior of the exponential family (which is only log-concave).
Here we focus on the strongly log-concave prior to recover a form more similar to DEQs that are implemented in practice.
Spectral normalisation.
Supposing (cid:107) · (cid:107) is the spectral norm, Assumption 2 is suggestive of a particular weight normalisation scheme: if λ is less than the spectral norm of K , then
[Equation]
The technique of using layers with unit spectral norm is explored by Miyato et al.
(2018) in the context of improving training stability of GANs.
Our results indicate that such spectral normalisation also leads to improved stability in DEQs at initialisation, in the sense that it guarantees a unique ﬁxed point by Proposition 4.
Abusing terminology, we refer to the spectral norm of the linear transformation corresponding with a convolutional layer as the spectral norm of the convolutional layer.
An easy to implement method for calculating the spectral norm of convolutional layers is available (Sedghi et al., 2018), which we use in our implementation and experiments.
E XPERIMENTS
DEQ parameters may be partitioned into two sets: the parameters of the layer that deﬁnes the ﬁxed point iteration ( θ =
( W 1 , W 2 , W 3 , V 1 , V 2 ) ), and the remainder of the parameters.
Leveraging our theoretical result, we give DEQs a warm-start by initialising θ such that the DEQ solves a kGLM.
We empirically demonstrate that such initialisation is superior to random intialisation.
We stress that our comparison is between randomly initialised DEQs and kGLM initialised DEQs only .
We are not interested in demonstrating that DEQs can achieve results competitive with state-of-the-art models, as this has already been demonstrated in previous studies (Bai et al., 2019; 2021).
Our limited scope allows us to perform a more deﬁnitive empirical study with meaningful statistics.
We give a detailed explanation of two of the tasks in the main text, and refer the reader to Appendix F for more experiments and Appendix G for an empirical demonstration of Theorem 3.
For random initialisation, we initialise weights from a symmetric uniform distribution (the Pytorch default).
Smooth sequence-to-sequence task.
Let X =
[ − 2 π − 2 , 2 π ] and deﬁne a function h ( x ) =
e − x 2 / 10 sin( x ) + e − ( x +9) 2 .
Let X s = X
and (cid:101)
X s = (cid:101)
X be 100 points on a uniform grid in [ − 2 π, 2 π ] and [ − 2 π − 2 , 2 π − 2] respectively.
Deﬁne sequences Y s and (cid:101) Y s to be joint evaluations of the s th realisation of a Gaussian process with covariance function k true ( x, x (cid:48) )
= 0 .
1 exp −(cid:107) x − x (cid:48) (cid:107) 2exp (cid:0) − sin 2 (cid:107) x − x (cid:48) (cid:107) (cid:1) , having means
h ( X ) and h ( (cid:101) X ) respectively.
[Figure]
For informed kGLM initialisation, we pretend that we know the kernel and coordinate space.
We choose k ≡ k true and form the required kernel matrices by taking X s = X and (cid:101)
X s = (cid:101) X .
For naive GLM initialisation, pretending we do not know the true underlying kernel or coordinate space, we default to k naive =
exp −(cid:107) x − x (cid:48) (cid:107) 2and choose X s =
X to be a uniform grid over the interval [ − 2 π, 2 π ] .
We initialise the hidden parameters W 1 and W 2 according to Theorem 3 and initialise V 1 and Vaccording to the same rule as random initialisation.
When training using naive GLM , we update only the readout parameters for the ﬁrst 10 epochs.
We generate a training and test set of size N = 20 , 000 and 2 , 000 .
Each sequence is evaluated on an n = 100 dimensional uniform grid.
We choose g θ ( z ∗ )
=
tanh ( W 1 z ∗ + W 2 Y + b ) and train for 400 epochs using Adam with default hyperparameters.
We repeat this for 100 trials using seeds 0 to 99 , see Figure 5(a) and (b).
Image denoising.
We consider a convolutional architecture as described in Appendix E, with σ ≡
ReLU , T the identity, and ρ ≡ 0 applied to an image denoising task using the CIFAR10 dataset.
Technically the ReLU is not admissible, but we are interested in examining its empirical properties.
The coordinate space X is the space of all 2 D pixel and channel triples ( i, j, c ) ∈ { 1 , . . .
, 32 } 2 × { 1 , 2 , 3 } .
Elements Y s are CIFAR10 images corrupted by i.i.d.
additive N (0 , 0 . 2) noise, clipped to take values between 0 and 1 .
Elements (cid:101)
Y s are the corresponding uncorrupted CIFAR10 images.
For random initialisation, we sample the ﬁlter weights i.i.d.
∼ N (0 , Var ) .
For kGLM initialisation, we let λ
=
(cid:107) K (cid:107) /C , where (cid:107) · (cid:107) denotes the spectral norm and C >
(see Assumption 2).
Spectral norms are calculated using the fast and exact method described by Sedghi et al.
The kernel is constructed randomly from the squared exponential kernel and described in Appendix F.
We are interested in separating the effects of scaling and initialisation scheme.
For random initialisation, the scale is chosen through the variance, inducing a convolutional layer with random spectral norm.
For kGLM initialisation, the scale C is equal to the spectral norm.
We try all variances in logspace (10 − 3 , 1 , 25) and all C in logspace (10 − 2 , 10 , 25) , resulting in a grid of spectral norms covering roughly the same space for both schemes.
We try 100 random seeds for each conﬁguration, resulting in 2 × 100 × 25 = 5000 models.
In Figure 6, we plot individually after training for 0 , . .
, 5 epochs the test MSE against the spectral norm for both initialisation schemes and all random seeds.
kGLM initialisation offers improved test error, training stability and reduced variance in test error.
The curves for kGLM initialised models adhere to the shape predicted by classical generalisation theory, where the spectral norm measures model complexity.
For kGLM initialisation, spectral norms of 10 0 appear to represent a critical point; smaller values lead to stable training and larger values lead to unstable training.
This contrasts with randomly initialised models, most of which have increasing test error with increasing spectral norm in their later epochs.
This means it is easy to select an appropriate scale for kGLM
[Figure]
initialisation, but not for random initialisation.
In Figure 3, we plot sample targets, inputs and outputs for random and kGLM initialised DEQs without any training, showing that samples using kGLM initialisation are visually more similar to the expected output than random initialisation.
Figureand epoch 0 of Figure 6 provide a sanity check of Theorem 3.
More plots are given in Appendix F.
D ISCUSSION AND CONCLUSION
Related work.
A review on connections between other optimisation-based iterative procedures and neural networks generally is given by Monga et al.
Repeated applications of weight-tied layers can sometimes be seen as solving model-based problems, if the weights are chosen appropriately.
Since the derived network implicitly minimises a conceptually tractable and principled model, it possesses a certain degree of interpretability.
However, only the minimisation problems and not the forward pass computations of the network are interpretable, and as such the network does not satisfy the requirements of interpretability in many settings (Rudin, 2019).
Monga et al.
(2021) also discuss how disabling weight-tying and allowing parameters to be ﬁne-tuned, thereby escaping the optimisation-based iterative procedure, can sometimes lead to improved performance, perhaps at the cost of interpretability, generalisation performance and theoretical guarantees.
Notable models analysed under this framework include ISTA (Beck & Teboulle, 2009) for solving
LASSO, which when unrolled and ﬁne-tuned yields LISTA (Gregor & LeCun, 2010), and ADMM for compressive sensing (Boyd et al., 2011) with an unrolled version called ADMM-CSNet (Yang et al., 2018).
Our work considers a different class of iterative models to those surveyed by Monga et al.
(2021), and recovers exactly the DEQ model.
In these works, the architectures, tasks, model classes and activations are speciﬁc to a particular setting.
More recently, Ramsauer et al.
(2020) introduce a class of networks whose predictions minimise a continuous analogue of a Hopﬁeld energy.
This can be used to justify transformer architectures.
Conclusion.
DDNs solve optimisation problems in their forward pass.
We considered the problem of regularised maximum likelihood in an RKHS.
Using this DDN layer, we derived a DEQ with a convolutional or fully connected implicit layer and a ﬁxed closed-form expression for the weights.
Such a result gives intuitive justiﬁcation for computing with DEQ models: DEQs can solve classical statistical problems in their forward pass.
Using this theoretical connection also allowed us to initialise DEQs as DDNs.
Such initialisation scheme offered performance beneﬁts over random initialisation both at initialisation and during training.
Future extensions of our work include considering other statistical estimators, learning the kernel, and obtaining Bayesian rather than point estimates of the canonical parameter.
E THICS STATEMENT
Our work is mostly theoretical in nature and is unlikely to directly cause harm, discrimination or privacy violations.
However, as our work may be used in future in applications, we encourage users of our work to adhere to ICLR’s general ethical principles.
We welcome open enquiry and scrutiny of the theoretical results presented in this paper.
R EPRODUCIBILITY STATEMENT
We have included code and the instructions required to reproduce our results.
The code is publically available at https://github.com/RussellTsuchida/deq-glm .
A CKNOWLEDGEMENTS
Russell, Suk Yee, Lars and Cheng Soon gratefully acknowledge the support of CSIRO’s Machine
Learning and Artiﬁcial Intelligence Future Science Platform.
A K ERNELISED GENERALISED LINEAR MODELS A. K ERNEL TRICK AND REPRESENTER THEOREM
The predictions of some classical algorithms depend on the training input data X ∈ R n × d x only through the matrix product XX (cid:62) .
A widely exploited technique in machine learning, the kernel trick, replaces X with ψ ( X ) ∈ R n × d ψ , where d ψ may be large or in fact inﬁnite.
Here ψ : R d x →
R d ψ and with an abuse of notation, when applied to a matrix it is understood to apply individually to each row so that ψ ( X ) ∈ R n × d ψ .
It is easy to change the resulting predictor since the type of ψ ( X ) ψ ( X ) (cid:62) ∈ R n × n matches that of XX (cid:62) , and scales with n and not d ψ .
One example is in kernel ridge regression (Saunders et al., 1998).
Following our abuse of notation, we will use k ( X, X ) ∈ R n × n to denote the matrix with ij th entry k ( x i , x j ) :
= ψ ( x i ) (cid:62) ψ ( x j ) .
Any function of the form k together with the matrix X characterises the matrix ψ
( X ) ψ ( X ) (cid:62) .
Such a function is called a kernel function, and is a valid choice if and only if it corresponds with an inner product in feature space or equivalently, if it is positive semi-deﬁnite (PSD).
The space of functions H k that the kernel k characterises is called a reproducing kernel Hilbert space (RKHS).
Let k be a positive semi-deﬁnite (PSD) function inducing a reproducing kernel Hilbert space (RKHS)
H k .
We will use K = k
( X, X ) ∈ R n × n to denote the matrix with ij th entry k ( x i , x j ) .
We will make use of the representer theorem, which for appropriately regularised empirical risk minimisation problems, allows one to solve optimisation problems in nonparametric function space.
Theorem 7 (Representer Theorem (Sch ¨ olkopf et al., 2002)) .
Denote by R : [0 , ∞ ) → R a strictly monotonic increasing function, by X a set and by L : ( X × R 2 ) n → R (cid:83) {∞} an arbitrary loss function.
Let H k be an RKHS with kernel k .
Then
[Equation]
A. K ERNELISED G ENERALISED LINEAR MODELS
We work with a special case of the exponential family in canonical form (Deisenroth et al., 2020,
Section 6.6), as shown in (10).
As has also been noted by others (O’sullivan et al., 1986; Canu & Smola, 2006; Cawley et al., 2007), the representer theorem allows one to extend appropriately regularised GLMs (McCullagh & Nelder, 1989) to kGLMs by replacing the linear predictor with an arbitrary element of an RKHS.
Here we demonstrate this fact.
Furnish Y with a σ -algebra B , forming a measurable space ( Y , B ) .
Let T be a measurable function
T : Y → R , called the sufﬁcient statistic .
Let ν be some reference measure and h : Y → R ≥ 0 be any function such that (cid:82) Y dH ( y ) < ∞ , where dH ( y )
=
h ( y ) dν ( y ) .
We call H the base measure and h the base density .
H is absolutely continuous with respect to ν , which may be the Lebesgue measure or the counting measure.
Deﬁne F to be the set of all η ∈ R such that (cid:82) Y exp( T ( y ) η ) dH ( y ) < ∞ .
Deﬁne the log-partition function A : F → R by A ( η ) = log (cid:82) Y exp
( ηT ( y )) dH ( y ) .
We may deﬁne a probability measure P ( dy | η ) = p ( y | η ) ν ( dy ) , where
[Equation]
[Equation]
is with an abuse of terminology called a probability density function for both discrete and continuous cases.
Finally, we form a joint probability density function as a product of marginals via p ( Y | η ) =
(cid:81) ni
=1 p ( y i | η i ) .
We will assume that F is an open set, so that following the terminology of Wainwright & Jordan (2008), (10) is called a member of the regular exponential family.
Regularity ensures that the logpartition function is inﬁnitely differentiable, and
that ∂A∂η =
E
[ T ( y )]
(Wainwright & Jordan, 2008,
Proposition 3.1).
Since T ( y ) and η are scalars, (10) satisﬁes a technical deﬁnition of a minimal exponential family.
Convexity of A is guaranteed for members of any regular exponential family, but minimality ensures that A is strictly convex (Wainwright & Jordan, 2008, Proposition 3.1).
Generalised linear models use an exponential family for the likelihood of observing response y i .
We work exclusively with univariate responses y i and canonical parameters η i , but GLMs are constructed more generally.
For GLMs with a canonical link function, one chooses η i = x (cid:62) i β
= f
( x i ) to be a linear function of the predictor variables x i with parameters β ∈ R d x .
It is possible to extend GLMs to accommodate a non-linear predictor f in an RKHS H k with a small modiﬁcation.
This modiﬁcation can be motivated by either a Bayesian or frequentist perspective by placing a Gaussian prior over the function values or regularising respectively, and then solving the resulting functional optimisation problem over an RKHS using the representer theorem.
Following the former philosophy, week seek the maximum a posteriori (MAP) estimate f ∗ of p (cid:0) f | X, Y (cid:1) with a prior density (Sch¨olkopf et al., 2001; Canu & Smola, 2006)
[Equation]
(11) where λ is a parameter and Z is a normalisation constant.
If f is an inﬁnite dimensional function, such a density is not deﬁned (recall that Gaussian processes are deﬁned in terms of their ﬁnite marginal distributions, which are jointly Gaussian).
However, if we restrict the underlying coordinate space X to be ﬁnite, then f is simply a vector evaluated over all coordinates and such a notation for the prior may be employed without concern.
The restriction of X to a ﬁnite space is without practical restriction in machine learning, since we only ever condition on and predict at a ﬁnite number of data points.
Using (11) as the prior and (10) as the likelihood, the MAP satisﬁes
[Equation]
so that by the representer theorem,
[Equation]
where
[Equation]
As an example, in binary kernel logistic regression, using the representer theorem we seek to minimise
[Equation]
where K = K ( X, X ) ∈ R n × n is the kernel matrix with pq th entry K pq =
k ( x p , x q ) .
The solution can be found by applying Newton’s method, resulting in an iterative ﬁxed-point numerical scheme for α ∗ .
Details of this method can be found in a number of references (Zhu & Hastie, 2005).
A detailed example of the Ising model is given in Appendix C.
B D ERIVATIVE TEST
We will require the mean value theorem.
Theorem 8 (Mean value theorem) .
Let S :
[ a, b ] → R n be continuous on [ a, b ] and differentiable on ( a, b ) .
Then there exists some t ∈ ( a, b ) such that
[Equation]
We now restate the result to be proven.
We follow the approach in theorem 2.2.16 of Hasselblatt &
Katok (2003).
Proposition 2.
Let Z be an open strictly convex set, Z its closure, H : Z → Z differentiable on Z and continuous on
Z .
If (cid:107) DH (cid:107) ≤ λ < 1 on Z , then H is a contraction on Z .
Proof.
Let z 1 , z 2 ∈ Z and deﬁne c ( t ) :
= z 1 + t ( z 2 − z 1 ) for t ∈ [0 , 1] and S ( t ) :
= H (cid:0) S ( t ) (cid:1) .
Then
S : (0 , 1) →
Z by strict convexity.
By the mean value theorem, there exists some t ∈ (0 , 1) such that
[Equation]
and therefore H is a contraction mapping.
C W ARM - UP : A TRACTABLE I SING MODEL
We demonstrate our main result for a special Ising model case.
The general case is given in § 3.
C. A S A K GLM
Let x and x (cid:48) be two dimensional indices on the ﬁnite lattice, x ∈ { 1 , . .
. , a }×{ 1 , . .
, b } for integers a, b ∈ Z .
The Ising model prescribes that
[Equation]
where Z is the partition function, some normalising constant ensuring that p ( Y | f
( X )) is a valid probability mass function.
Here the parameters J pq determine interactions between the spin states y ip and y iq , f x ip controls an interaction between the spin state y ip and an external magnetic ﬁeld at lattice position x ip and t is a temperature parameter.
We let f x = f ( x ) belong to an RKHS H k with
PSD kernel k .
With a standard Gaussian prior over f , we have that
[Equation]
Supposing no interactions, J pq = 0 and Z = 2 cosh ( t − 1 f ( x ip ))
(Nguyen et al., 2017).
By Theorem 7, letting K = t − 1 k ( X, X ) denote the matrix with pq th entry t − 1 k ( x ip , x iq ) ,
[Equation]
The problem is strictly convex so we may ﬁnd the solution by applying Newton’s method.
Using that the derivative of log (2 cosh( · )) is tanh( · ) , stationarity of the solution implies that
[Equation]
The Jacobian ∂F∂α ∗ is given by
[Equation]
where
[Equation]
, so that the c th Newton iteration requires solving a linear system,
[Equation]
We employ the shorthand M ( α ∗ ( c ) ) =
M ( c ) for any matrix function of α ∗ ( c ) .
The Newton update is typically rearranged as an Iteratively Re-Weighted Least Squares (IRWLS) form by substituting α ∗ ( c − 1) =
J − 1 ( c
[Equation]
J ( c − 1) α ∗ ( c − 1) , yielding
[Equation]
where z :=
Kα + D − 1 ( Y − µ ) and µ = tanh ( Kα ∗ ) .
Once α ∗ has been approximated to some tolerance using this numerical scheme, the predictor f ∗ may be evaluated at points (cid:101) X ∈ R (cid:101) n × d x using the representer theorem,
[Equation]
C. A S A DEQ
Rearranging (12), and applying tanh to both sides of the equation, we obtain
[Equation]
Choosing z ∗ =
tanh ( Kα ∗ ) , W 1 = − ( tλ ) − 1 K , W 2 = ( tλ ) − 1 K , we have
[Equation]
Note that α ∗ being a unique ﬁxed point of (12) implies that z ∗ is a ﬁxed point of (14), but not that the ﬁxed point of (14) is necessarily unique.
One condition that guarantees the uniqueness of the ﬁxed point of (14) is that the induced matrix norm of K is less than or equal to 1 .
Without this assumption, (14) is still true but when used as a computational model, the output will in general depend on the ﬁxed-point solver and initial conditions passed to that solver (and even still, might never return unstable ﬁxed points).
Finally by assuming that K is invertible and using (13), we obtain
[Equation]
where V 1 = − ( λ ) − 1 (cid:101) K and V
= ( λ ) − 1 (cid:101) K .
D P ROOF OF MAIN RESULT
Theorem 3.
Let α ∗ be the minimiser of one of the inner optimisation problems (5), inducing a function f ∗ ( · ; X, Y )
=
k ( · , X ) α ∗ .
Deﬁne z ∗ :=
A (cid:48)
( Kα ∗ ) .
Then z ∗ is a ﬁxed point satisfying
[Equation]
where W 1 = W 3 = − λ − 1 K and W 2 = λ − 1 K are parameter matrices, σ ≡ A (cid:48) and ρ ≡ ( − log q ) (cid:48) ◦
( A (cid:48) )
− 1 are monotone non-decreasing.
Furthermore, under Assumption 1, for any test index (cid:101) X ∈
R (cid:101) n × d x and for at least one z ∗ satisfying (7),
[Equation]
where V 1 = − λ − 1 k ( (cid:101) X, X ) and V 2 = λ − 1 k ( (cid:101) X, X ) .
Proof.
Since A is strictly convex, A (cid:48)
= σ strictly increasing and invertible.
Q := − log q ( · ) is a convex function, and its derivative Q (cid:48) is monotone nondecreasing.
Also note that Q (cid:48) ◦
( A (cid:48) )
− 1 is monotone nondecreasing.
Let K
= k
( X, X ) .
From (5), stationarity at the minimum, differentiability of A and Q , and the representer theorem implies
[Equation]
Note that z ∗ is only a (not necessarily unique) ﬁxed point of (16).
Suppose Assumption 1.
Since K is invertible, and σ is invertible (since A is strictly convex), then again using the representer theorem, new predictions may be formed as
[Equation]
where (cid:101) K
= k ( (cid:101) X, X ) .
Proposition 4.
Under Assumption 2, (7) admits a unique ﬁxed point on Z .
Proof.
Let H be the mapping deﬁned by iteration (16),
[Equation]
H is a contraction since the induced matrix norm of the Jacobian of H satisﬁes
[Equation]
<
[Equation]
The result follows from Theorem
[Equation]
[Table]
[Figure]
[Equation]
[Figure]
empirical risk minimisation
[Figure]
[Equation]
[Figure]
t.
[Figure]
Feature embedding
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
as inner
[Figure]
[Equation]
E C ONVOLUTION AS MATRIX MULTIPLICATION
Let w, h ∈ Z > 1 denote the pixel width and height of a space of images with c channels.
Let
[Equation]
denote the corresponding pixel and pixel-channel spaces.
Let X ∈ R ( whc ) ×
c denote a matrix, where every row is one of the pixel coordinates ( i, j, c )
∈ X whc .
Choose X s = X for some set of N images { Y s } Ns =1 .
The images Y s are of size whc , and contain pixel intensities for the corresponding pixel coordinates in X .
Let x 1 , x 2 and x 3 denote the ﬁrst, second and third coordinates of an element x ∈ X whc .
Let
[Equation]
where 1 is the indicator function for some truncating parameters r 1 , r 2 > 0 .
Let κ 0 :
X whc ×
X whc
→ R denote any stationary PSD kernel, for example κ 0 ( x, x ) =
e − 12 (cid:107) x − x (cid:48)
(cid:107) 22 .
Note that these functions each deﬁne a PSD kernel.
Using the fact that kernels are closed under multiplication, deﬁne
[Equation]
Here h 1 and h 2 ensure that only close pixels have non-zero kernel values.
Note that the matrix κ ( X, X ) is sparse when r 1 or r 2 are small and has an off-diagonal structure.
More concretely choosing κ ( x, x (cid:48) )
=
e
[Equation]
(cid:107) [ x 1 ,x 2 ] − [ x (cid:48) 1 ,x (cid:48) 2 ] (cid:107) 22
(cid:107) x − x (cid:48) (cid:107) ∞ ≤ 5) where (cid:107)·(cid:107) ∞ denotes the supremum norm, we obtain a whc × whc matrix.
Each block of size wh is repeated horizontally.
This matrix represents c convolutional ﬁlters of ﬁlter length 5 × 2 + 1 = 11 in both directions.
We plot the induced 32 2 × 32 2 kernel matrix when w = h = 32 and c = 1 in Figure 4, as well as the corresponding operation as a convolution.
When c > 1 , this matrix is simply repeated vertically and horizontally c times.
For rows in the middle of the matrix, there are 11 2 non-zero entries, corresponding to the indices included in the convolution operation when centered on the pixel of the corresponding row.
The proof is tedious and mostly centres around the notation of appropriate reshaping of arrays and zero-padding.
In order to state our result, we ﬁrst need to deﬁne the neural network convolution operation.
Deﬁnition 9 (Single-channel convolution) .
Fix X
=
X wh (in the sense of (17)) and let the rows of
X be equal to the elements of X .
Let ∗ 1 denote the single-channel image convolution operation (in the sense of convolutional networks).
More concretely, C ∗ 1 Y applies a ﬁlter C ∈ R (2 r 1 +1) ×
(2 r 2 +1) to a ﬂattened image Y ∈ R wh as follows.
First ﬂatten C to be an element ˜ C of R (2 r 1 +1)(2 r 2 +1) and zero-pad Y to be an element (cid:101) Y of R (2 r 1 + w )(2 r 2 + h ) .
Then deﬁne for each j ∈ { 1 , wh } a vector ˜ C j with the same dimensionality as (cid:101) Y i , such that ˜ C j, [ j : j +(2 r 1 +1)(2 r 2 +1)]
= ˜ C and all other elements of ˜ C
j are zero.
Then for each j ∈ { 1 , wh } , deﬁne
[Equation]
so that C ∗ 1 Y ∈ R wh , where C mat ∈ R wh × (2 r 1 + w )(2 r 2 + h ) is a matrix with j th row equal to ˜ C j .
Deﬁnition 10 (Multi-channel convolution) .
Let ∗ denote the muti-channel image convolution operation (in the sense of convolutional networks).
More concretely, C ∗ Y applies a set of ﬁlters C ∈ R c 1 × c 2 × (2 r 1 +1) ×
(2 r 2 +1) to a ﬂattened image
Y ∈ R whc 2 as
[Equation]
so that C ∗ Y ∈ R c 1 wh , where Y r ∈ R wh represents the r th channel of Y .
Generally, we have the following corollary of Theorem 3.
Corollary 11.
In the same setting as Theorem 3, under Deﬁnition 9, let k ≡ κ (in the sense of (18)) and let f ∗ ( · ; Y, X ) =
κ ( · , X ) α ∗ be a global minimiser of (5).
Then z ∗ is a ﬁxed point satisfying
[Equation]
(19) where z ∗ =
A (cid:48) ( κ ( X, X ) α ∗ ) .
C 1 = C 3 and C 2 are convolutional ﬁlters with pq th entry − λ − 1 κ 0 ( X (cid:62) i, [ r 1 +1 ,r 2 +1] , X i, [ p,q ] ) and λ − 1 κ 0 ( X (cid:62) i, [ r 1 +1 ,r 2 +1] , X i, [ p,q ] ) respectively.
σ = A (cid:48) and ρ = ( − log q ) (cid:48) ◦
( A (cid:48) )
− 1 are monotone non-decreasing.
Furthermore, 1.
Under Assumption 2, (7) admits a unique ﬁxed point on Z .
Under Assumption 1, for any test index (cid:101) X ∈ R (cid:101) n × d x and for at least one z ∗ satisfying (7),
[Equation]
20) where V 1 = − λ − 1 κ ( (cid:101) X, X ) and V 2 = λ − 1 κ ( (cid:101) X, X ) .
Proof.
We begin with the case c = 1 .
By Deﬁnition 9, we have that
[Equation]
where ˜ z ∗ ∈ R (2 r 1 + w )(2 r 2 + h ) is a zero-padded z ∗ and C mat ∈
R wh × (2 r 1 + w )(2 r 2 + h ) .
The j th row of
C mat is equal to ˜ C j , where ˜ C j, [ j : j +(2 r 1 +1)(2 r 2 +1)]
= ˜ C and all other elements of ˜ C
j are zero.
Here ˜
C is a ﬂattened version of C 1 .
By removing the zero entries from ˜ z ∗ and the corresponding columns from C mat , we may write
C mat ˜ z ∗ =
W mat z ∗ for some W 1 mat ∈ R wh × wh .
The same holds for W 2 and W 3 .
We may therefore write
[Equation]
On the other hand, Theorem 3 says that
[Equation]
Choosing W 1 mat =
W 1 = − λ − 1 κ ( X, X ) and similarly for W 2 , W 3 , we obtain (19).
This implies that the ﬁlter C 1 has pq th entry − λ − 1 κ 0 ( X (cid:62) i, [ r 1 +1 ,r 2 +1] , X i, [ p,q ] ) .
The rest of the corollary follows from Theorem 3.
[Figure]
F D ETAILED EXPERIMENTS
Our experiments may be reproduced using the code provided in the supplementary material by following the README.md ﬁle.
Choice of kernel for initialisation.
To intialise a convolutional network, we need to choose r 1 , r 2 and κ 0 according to (18).
We choose r 1 = r 2 = 3 and κ 0 ( x, x (cid:48) )
= (cid:80) c
( c +1) / 2 r =
[Equation] [Equation]
2 , where g r ( a, b ) =
if a == a r and b ==
b r or a ==
b r and b =
= a r and 0 otherwise.
We sample the squared lengthscales from a uniform distribution between 0 and 4 .
The cost of kGLM initialisation is minuscule, and often even faster than random initialisation.
For a fully connected layer with hidden and readout weight matrices W and V of size n × n and (cid:101) n × n respectively, we require n 2 + n (cid:101) n evaluations of the kernel function and then a normalisation by the spectral norm (the default option for this calculation for a Python library such as Numpy will be the LAPACK divide and conquer algorithm, which will have a worst case of O ( n 3 ) , but in practice the computation will be very fast).
In contrast, random numbers (the usual method for initialising neural networks) requires n 2 + n (cid:101) n random number generations, which depending on software platform, is usually done through calls to the inverse transform sampling method, perhaps using a stochastic collocation Monte Carlo sampler.
Either way, empirically we ﬁnd that kGLM initialisation is actually much faster than random Gaussian initialisation for fully connected layers.
See Figure 8.
For convolutional layers, which are basically just extremely sparse large fully connected layers, the kernel method induces an additional overhead associated with the various reshaping operations required to put the elements of the kernel matrix in the correct position in the convolutional ﬁlter layer.
The calculation of the spectral norm is done through the method of Sedghi et al.
(2018) at a cost of O (cid:0) w 2 c 2 ( c + logw )) , where c is the number of channels and w is the maximum of the pixel width/height of the image.
Our implementation is far from optimised, but we ﬁnd that even with this overhead, the initialisation costs less than 1 second on a laptop PC.
This is a tiny fraction of the cost typically associated with training the network, which depends on the problem but will usually be much less than 1% .
See Figure 8.
Scaling up parameter counts.
The number of rows of the square kernel matrix is equal to the number of rows in X .
This means, for example, in the setting of CIFAR10 where X ∈ R 32 ∗ 32 ∗ 3 × 3 , there are roughly 9 .
4 ∗ 10 6 elements in the kernel matrix.
Using a convolutional representation, the sparse kernel matrix may be represented as a convolutional layer of size (3 , 3 , 32 , 32) .
Crucially, the parameter count scales quadratically with the number of input channels, which in this case is 3 .
One way to experiment on larger models is to scale up the size of the input data.
To this effect, we take a dataset containing 3799 hyperspectral ( 25 channel) images (HSI) of roads (Lu et al., 2020) of total size (3799 , 25 , 192 , 384) .
We consider a sequence of denoising tasks when the input data is a
10101010 1 Average spectral norm 1.00e1.00e1.00e+
T e s t M S E
Epoch =10101010 1 Average spectral norm 2×103×104×106×
T e s t M S E
Epoch =10101010 1 Average spectral norm 2×103×104×106×
T e s t M S E
Epoch =10101010 1 Average spectral norm 2×103×104×106×
T e s t M S E
Epoch =10101010 1 Average spectral norm 2×103×104×106×
T e s t M S E
Epoch =10101010 1 Average spectral norm 2×103×104×106×
T e s t M S E
Epoch = subset of the 25 hyperspectral channels.
The resulting plots are given in Figures 9, 10, 11 12.
The results are summarised in Table observed over the full grid.
Due to training instability causing NaN and Infs data points, especially in randomly initialised models, the average and “sample standard deviation” are not true unbiased estimators.
kGLM initialisation shows superior performance, both in terms of average performance and number of failed runs
[Table]
[Figure] [Figure]
[Figure]
G E XAMPLE EMPIRICAL INSTANTIATION OF EQUIVALENCE
We take binarised MNIST and for each image, independently corrupt 30% of bits by randomly ﬂipping them.
We are interested in obtaining a predictive model that can segment the original binarised image (cid:101) Y s from its noisy corrupt input Y s , together with the constant coordinate locations
X s = (cid:101)
X s =
X that may be represented as 784 × 2 coordinate indices.
The index 1 ≤ s ≤ 60 ,runs over different images in the training dataset.
We evaluate model performance on an independent testing dataset of size 10 , 000 .
We empirically compare ﬁve related models.
Kernel logistic regression ( KLR ).
Since each element of Y s ∈ R 784 is binary, a natural kGLM model to use is kernel logistic regression (KLR) (Green & Yandell, 1985; Zhu & Hastie, 2005).
With respect to the exponential family, such a model prescribes that the derivative of the log partition function A (cid:48) ( · ) =
σ ( · ) = (1 + e −· ) − 1 is the logistic sigmoid and the sufﬁcient statistic T is the identity.
As is common in the literature, we employ a zero-mean Gaussian prior so that ρ ≡ 0 .
The only remaining choices are the kernel function k and the regularisation parameter λ .
In order to ensure that Assumption 1 is satisﬁed, we choose k to be the sum of a squared exponential kernel with lengthscale 1 and a white noise kernel with variance 10 − 8 .
This forces the resulting kernel matrix K to be strictly positive deﬁnite.
In order to satisfy Assumption 2, we need to ensure that (cid:107) K/λ (cid:107) ≤(assumption A3.1 is automatically satisﬁed due to the setting of A in KLR).
Accordingly, we set λ = 2 (cid:107) K (cid:107) .
Note that we do not actually need to ensure that Assumptions 1 and Assumptionsare satisﬁed for Theorem 3 to hold, but we do need these assumptions so that all of our theoretical results hold.
The KLR model is trained individually using every data point in { Y s } 60 , 000 s = as a data set and a manually implemented iterated reweighted least squares algorithm consisting of updates according to (exact) Newton’s method.
We allow Newton’s method to run for 5 steps, at which point the norm of the difference between subsequent values of α is typically on the order of 10 − 14 .
The output of the
KLR model is a kernelised predictor f ∗ that has been trained using data X s and Y s , evaluated on X s and converted to the maximum a posteriori of the Bernoulli distribution through 1 (cid:8) σ ( f ∗ ) > 0 .
5 (cid:9) .
DEQ with kernel initial parameters ( DEQ-kernel ).
Theorem 3 implies an equivalent DEQ model with ﬁxed parameters.
We use this model, namely
[Equation]
where W 1 = V 1 = − λ − 1 K and W 2 = V 2 = λ − 1 K are parameter matrices and σ is the logistic sigmoid function.
To convert the kernelised predictor f ∗ ( Y ) to the mean of the Bernoulli distribution, we must pass it again through the inverse link function,
[Equation]
and the resulting maximum a posteriori prediction is determined by examining whether the mean of the Bernoulli distribution is greater than 0 .
[Equation]
DEQ with kernel initialisation and ﬁne-tuned parameters ( DEQ-kernel-trained ).
We interpret the DEQ with initial parameters model as a neural network, and update the parameters
W 1 , W 2 , V 1 and V 2 using Adam to minimise the squared loss between the greyscale prediction and binarised ground truth.
We also allow biases in both the hidden and readout layers.
We train forepochs using a batch size of 100 , leading to 5 × (
= 3000 gradient updates.
DEQ with random initialisation ( DEQ-random ).
We take the same neural network architecture as with kernel initialisation, but randomly initialse the parameters from a Gaussian distribution using the Pytorch default.
DEQ with random initialisation and ﬁne-tuned parameters ( DEQ-random-trained ).
We take the same neural network architecture as with kernel initialisation, but randomly initialse the parameters from a Gaussian distribution using the Pytorch default.
We also allow biases and train using the same optimiser and loss.
Observations.
Remarkably (but not surprisingly, due to Theorem 3), the binarised prediction agrees exactly with KLR (an entirely independent piece of software, that uses an entirely different solver) over all but 11 of
, 000 MNIST datapoints each containing 784 pixels.
For those 11 where disagreements were observed, only 1 of 784 pixels were different.
The very small differences arise due to the tolerance of the solvers and the way tied predictions (i
[Table]
e. those whose logistic sigmoid output is 0 . 5 ) are handled.
In all 11 cases, the DEQ output a value of 0 .
5 and KLR output some value less than but very close to 0 .
See Table 2 for posterior probabilities, and Figure 13 for the corresponding binarised predictions.
Figure 14 shows the test error as a function of training epoch for DEQ-kernel-trained and DEQ-random-trained
[Table]
the 11 examples above.
In these cases, the predicted posterior probabilities were very close to borderline.
[Figure]
H E XTENSION TO NONCENTERED CANONICAL PARAMETERS
It is straightforward to extend Theorem 3 (and related corrolaries) to the case where the canonical parameters
η are the sum of evaluations of some deterministic function m and evaluations of the predictor f .
We sketch this now and leave details to the reader to suit their application.
Starting from (10) and using prior (4), we obtain a new MAP objective
[Equation]
so that by the representer theorem, f ∗ ( · ) =
k ( · , X ) α ∗ is given by
[Equation]
where M := m ( X ) .
The same reasoning in the proof of Theorem 3 then still follows if we deﬁne z ∗ =
A (cid:48)
( Kα ∗ + M ) .
That is, the equilibrium condition implies that
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
and so on.
A similar result holds when an additional linear scaling of f ( X ) is introduced.
I S ENSITIVITY OF FIXED POINT EQUATIONS TO NON -PSD MATRIX PERTURBATIONS
The DEQ layers of the models we studied in the main text have parameter matrices that are symmetric PSD kernel matrices.
This departs from the established practice, where one randomly initialises the parameter matrices (say from a zero-mean Gaussian distribution with small variance), and then trains the parameters using gradient descent.
Here we investigate connections between our restricted parameter setting and more general parameters.
We begin by bounding the difference between ﬁxed points of DEQ layers with the same activation σ and input Y , but different parameter matrices.
Proposition 12.
Let Z be an open strictly convex set and Z its closure.
Denote by (cid:107) · (cid:107) some vector norm on Z and its induced matrix norm.
Let g 1 : Z → Z and g 2 : Z → Z be deﬁned by
[Equation]
for some Y ∈ R n , function σ with Lipschitz constant L σ ≤ 1 (with respect to (cid:107) · (cid:107) on Z ) and matrices A, B with matrix norms L A
= (cid:107) A (cid:107) < 1 and L B = (cid:107) B (cid:107) < 1 .
Deﬁne L 1 =
L σ L A and L 2 =
L σ L B .
Then for ﬁxed points z 1 = g 1 ( z 1 ) and z 2 = g 2 ( z 2 ) (each guaranteed to exist and be unique),
[Equation]
Proof.
Existence and uniqueness of ﬁxed points follows directly from Theorem 1.
Exploiting only the triangle inequality and Lipschitz properties, we have
[Equation]
Using g 2 ( Y )
=
σ ( 0 )
[Equation]
and again the triangle inequality and Lipschitz properties
[Equation]
note that
[Equation]
so that
[Equation]
We will now use this result to bound the difference between ﬁxed points found by arbitrary parameter matrices and ﬁxed points found by PSD kernel matrices, which exactly solve kGLMs as in the main text.
It helps to think of real-valued PSD matrices as being different from general matrices in two respects.
Firstly, real-valued PSD matrices are always symmetric (so their eigenvalues are real).
Secondly, the eigenvalues of PSD matrices are always nonnegative reals.
The matrix norm of the difference between any matrix and its “closest” PSD counterpart can be decomposed into these two distinct effects when the matrix norm is the Frobenius norm (Higham, 1988, Theorem 2.1).
However, the
Frobenius norm is difﬁcult to work with as it is not induced by a norm over the vector space Z , and only provides loose bounds to the spectral norm.
We instead operate on (cid:107) · (cid:107) 2 directly, and obtain a bound that contains a similar decomposition of two effects.
Proposition 13.
In the same setting as Proposition 12, choose (cid:107)· (cid:107) to be the Euclidean and spectral norms (cid:107) · (cid:107) 2 .
Let A be any n × n matrix.
A admits a decomposition into symmetric and skew symmetric parts
A =
A 1 + A 2 where
A 1 = 12 (cid:0) A + A (cid:62) (cid:1) and A 2 = 12 (cid:0) A − A (cid:62) (cid:1) .
Let { B (cid:23) 0 } denote the space of PSD matrices.
Then
[Equation]
where (cid:102) λ
2 be the largest absolute value of the negative eigenvalues of A 1 .
Proof.
From the symmetric, skew-symmetric decomposition, A = 12 Q Λ Q (cid:62) + A 2 , where Q is (real) orthogonal and Λ is a real diagonal matrix.
Choose B = 12 Q Λ + Q (cid:62) , where Λ + is the elementwise maximum of Λ and zero.
Then
[Equation]
The ﬁrst term is the operator norm of A 2 .
The second term is the largest absolute value of the negative eigenvalues of A 1 .
Finally note that L B ≤ (cid:107) A 1 (cid:107) 2 ≤ (cid:107) A (cid:107) 2 , so that 1 1 − L 2 ≤ 1 1 − L 1 , and apply Proposition 12.
This result helps us understand what happens when parameter matrices that are almost PSD kernel matrices are used.
We note that if A is PSD, the bound evaluates to zero.
If A is symmetric and real, then the bound is a linear scale of the most negative eigenvalue of A .
If A is any matrix, we have to additionally account for the operator norm of the skew-symmetric part of A , which contains purely imaginary eigenvalues.
A more blunt observation can be made when we do not know anything about the norm of the skew symmetric part or the most negative eigenvalue of the symmetric part of A , by relating both of these quantities to the norm of A .
Here rather than using relative spectral properties of A , we can just use the scale of A to conclude that if A contains small entries, there is a PSD matrix B that obtains a close ﬁxed point, and we can quantify this closeness.
This makes use of the coarse inequality that
[Equation]
A natural question to ask is when such a blunt device might be applied in practical scenarios.
It is noted in the tutorial material (Kolter et al., 2020) that DEQs parameters are typically initialised with smaller weight variance than standard feedforward networks.
Recall that Xavier Glorot initialisation for standard feedforward networks chooses a parameter standard deviation of 1 √ n .
If we slightly alter this initialisation scheme to a achieve a smaller weight variance (that is still O ( n − 1 / 2 ) ), we can apply the conclusion of the rough reasoning above by applying results from random matrix theory.
Proposition 14.
In the same setting as Proposition 13, let A be a random matrix with Gaussian entries having zero mean and standard deviation ν = C/L σ 2 √ n + t .
Then with probability 1 − 2 exp − t 2 /for
t ≥ 0 , the smallest sum of squared errors is bounded by
[Equation]
In particular if (cid:107) Y − σ ( 0 ) (cid:107) 2 2 ≤ Dn for some D ≥ 0 , the smallest mean squared error is bounded by
[Equation]
Proof.
(cid:101) λ is bounded by the largest absolute value of all of the eigenvalues of the symmetric part of
A , which is the operator norm of the symmetric part of A .
Thus
[Equation]
By Proposition 13 we therefore have
[Equation]
Choosing C such that 0 < C/L σ < 1 and letting ν = C/L σ 2 √ n + t in Theorem 15 below, we have that
[Equation]
probability exceeding
[Equation]
Theorem 15 (Corollary of Theorem 2.6 in Rudelson & Vershynin (2010)) .
Let A be an n 1 × nmatrix with independent zero-mean Gaussian random entries each having standard deviation ν .
Then for all t ≥ 0 ,
[Equation]
When C is small, we expect there to exist kernel matrices that perform roughly the same (as poorly) as random Gaussian initialisation.
This result is consistent with our empirical results in the ﬁgures in Appendix F, where we observe that randomly initialised parameter matrices with spectral norm of around 10 − 1 or less obtain roughly the same performance as kernel initialised models with the same norm, and both outperform randomly initialised parameter matrices with large spectral norms.
However, this experimental setup is slightly complicated by the use of convolutional layers, which when converted to parameter matrices, are very large, sparse and contained shared elements.