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 paper discusses a new perspective that establishes the connection between optimization-based layers (i.e., DDNs) and the fixed-point forward computations of the deep equilibrium networks (DEQs). In particular, the paper shows that a kGLM optimization layer, under certain regularity and contractivity assumptions, can be written as an (arguably) simpler form of deep equilibrium model layer. The paper provides solid proof and a thorough discussion of the implication of this connection, such as 1) how we might want to parameterize the matrices; and 2) how this may induce a good initialization scheme for (this particular kind of) DEQ models which improves training stability and convergence.