Declarative nets that are equilibrium models

1 INTRODUCTION . Implicit layers ( Pineda , 1987 ; Almeida , 1990 ) have recently been subject to renewed attention ( Kolter et al. , 2020 ) . In contrast with explicitly defined layers , implicit layers define 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 fixed 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 . ( 2021 ) ) . 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 fixed points of an iterative scheme . A correspondence between DDNs and DEQs is expected ( but undiscovered ) , given the fundamental connection between fixed 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 benefit ? 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 . Notation . Let X ⊆ Rdx and Y ⊆ R be coordinate and target spaces . Define a Y−valued stochastic process { y ( x , ω ) } x∈X indexed by x ∈ X with outcome ω from a sample space Ω. LetX ∈ Rn×dx be a matrix such that the ith row ofX is some element xi ∈ Rdx . Similarly define X̃ ∈ Rñ×dx . Take n evaluations { yi } ni=1 = { y ( xi , ω0 ) } ni=1 from a realisation of the stochastic process for every ith row in X , and form a corresponding matrix Y . Using the same realisation of the stochastic process , take ñ evaluations { ỹi } ñi=1 = { y ( x̃i , ω0 ) } ñj=1 for every ith row in X̃ . Call X and Y the inner coordinate and target data , and X̃ and Ỹ the outer coordinate and target data . X and X̃ encode the positions of corresponding elements in Y and Ỹ , and induce kernel matrices whose utility roughly corresponds with positional embeddings ( Gehring et al. , 2017 ) in transformers ( Vaswani et al. , 2017 ) . Example implication of result . Consider image inpainting a single 32 × 32 colour image . Take X = { 1 , . . . , 32 } 2 × { 1 , 2 , 3 } , and X ∈ Zn×3 to be a matrix consisting of ( non-repeated ) triples , where n < 322 . Take X̃ ∈ Zñ×3 to be the matrix consisting of remaining triples , such that n+ ñ = 322 × 3 . An incomplete image Y ∈ Rn together with its missing values Ỹ ∈ Rñ are jointly sampled from the same distribution from which CIFAR10 images are drawn . We may use kernel ridge regression or a Gaussian process with Gaussian likelihood to produce a per-pixel prediction of Ỹ given X , Y and X̃ . These models represent special cases of regularised or Bayesian kGLMs with closed-form predictors , 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 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 fixed hyperparameters , the DEQ has a fixed 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 involving kernel logistic regression for MNIST image segmentation is given in Appendix G . 2 DEEP EQUILIBRIUM MODELS . Given some Y ∈ Rn×dy containing row elements y ∈ Y , deep equilibrium models ( DEQs ) ( Bai et al. , 2019 ) compute an embedding z∗ ∈ Z ⊂ Rn×dz as a fixed point of some function and then output an affine transformation f∗ ( Y ) of the fixed point . More concretely , letting gθ ( · ; Y ) : Rn×dz → Rn×dz denote a parameterised function ( e.g . a neural network ) with parameters θ ∈ Θ , f∗ ( Y ) = V z∗ , z∗ = gθ ( z ∗ ; Y ) , z∗ ∈ Z ⊂ Rn×dz , ( 1 ) where V ∈ Rdu× ( n×dz ) are referred to as readout parameters , and V may be thought of as a du×ndz matrix acting on a flattened vector with ndz entries . We will find it more natural to introduce a linear skip connection between Y and the output . We will also choose dy = dz = 1 . That is , we consider f∗ ( Y ) = V1z ∗ + V2Y , z ∗ = gθ ( z ∗ ; Y ) , z∗ ∈ Z ⊂ Rn ( 2 ) 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 Ỹ with respect to the DEQ and neural network parameters . Given a dataset { ( Ys , Ỹs ) } Ns=1 , this procedure may be realised by running an iterative algorithm over the equality constrained optimisation problem min γ , θ , V1 , V2 N∑ s=1 L ( Ỹs , Fγ ( f∗ ( Ys ) ) ) subject to z∗s = gθ ( z ∗ s ; Ys ) , f ∗ ( Ys ) = V1z ∗ s + V2Ys , s = 1 , . . . , N ( 3 ) 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 V1 , V2 ; these may be computed under mild conditions without backpropagating through the fixed point solution via implicit differentiation . We refer the reader to Bai et al . ( 2019 ) for details . Dealing with existence and uniqueness of fixed points . In order for ( 1 ) and ( 2 ) to specify useful computational rules , they have to admit at least one fixed point . To avoid having to choose which fixed point should be returned with additional rules , it might be desirable that they admit exactly one fixed point . We give example configurations of fixed points in Figure 2 . The classical Banach fixed point theorem ( BFPT ) gives sufficient conditions for the existence and uniqueness of fixed 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′ ∈ Z , d ( H ( z ) , H ( z′ ) ) ≤ qd ( z , z′ ) . Every contraction admits a unique fixed point . The contraction property is suggestive of a simple algorithm ( Hasselblatt & Katok , 2003 ) for finding the fixed point z∗ of a contraction H given an initial guess z0 : While some termination condition is not met , Update the current estimate zr for the fixed point to beH ( zr−1 ) . This algorithm allows one to interpret the inner problem of ( 3 ) as an infinitely 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 ‖ · ‖ of the Jacobian of H is strictly less than 1 on Z , providing a useful test for checking whether a unique fixed 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 ‖DH‖ ≤ λ < 1 on Z , then H is a contraction on Z . 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 fixed point exists ( Bai et al. , 2019 ) , constrain or monitor the parameters during learning to ensure that a unique fixed 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 fixed point exists ( Winston & Kolter , 2020 ) . Other conditions are given by El Ghaoui et al . ( 2021 ) . We do not address this problem ; we assume ( in a precise sense , see Assumption 3 ) verifiable 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 find parameter initialisations as a corollary of our main theoretical result . See Figure 3 .

Summary: The authors draw a resemblance between two classes of models, namely deep declarative networks (DDNs) and deep equilibrium models (DEQs), under mild assumptions. In particular, they show that solving an empirical risk minimization problem of a kernelized generalized linear model (kGLM) over a RKHS in a DDN admits a fixed point that, when viewed with an outer minimization, resembles the optimization of a simple DEQ with a fully connected implicit layer. They use their result to motivate an initialization scheme for DEQs that matches the solution of a trained kGLM.