Neural Deep Equilibrium Solvers

A deep equilibrium ( DEQ ) model abandons traditional depth by solving for the fixed point of a single nonlinear layer fθ . This structure enables decoupling the internal structure of the layer ( which controls representational capacity ) from how the fixed point is actually computed ( which impacts inference-time efficiency ) , which is usually via classic techniques such as Broyden ’ s method or Anderson acceleration . In this paper , we show that one can exploit such decoupling and substantially enhance this fixed point computation using a custom neural solver . Specifically , our solver uses a parameterized network to both guess an initial value of the optimization and perform iterative updates , in a method that generalizes a learnable form of Anderson acceleration and can be trained end-to-end in an unsupervised manner . Such a solution is particularly well suited to the implicit model setting , because inference in these models requires repeatedly solving for a fixed point of the same nonlinear layer for different inputs , a task at which our network excels . Our experiments show that these neural equilibrium solvers are fast to train ( only taking an extra 0.9-1.1 % over the original DEQ ’ s training time ) , require few additional parameters ( 1-3 % of the original model size ) , yet lead to a 2× speedup in DEQ network inference without any degradation in accuracy across numerous domains and tasks . 1 INTRODUCTION . Recent progress on implicit networks , such as Neural ODEs ( NODEs ) ( Chen et al. , 2018b ; Dupont et al. , 2019 ; Rubanova et al. , 2019 ; Jia & Benson , 2019 ; Kelly et al. , 2020 ) and deep equilibrium ( DEQ ) models ( Bai et al. , 2019 ; Winston & Kolter , 2020 ; Kawaguchi , 2021 ; Bai et al. , 2020 ; Gilton et al. , 2021 ) , has motivated this novel class of networks to the forefront of deep learning research . Instead of stacking a series of operators hierarchically , implicit models define their outputs as solutions to nonlinear dynamical systems . For example , DEQ models ( which this paper will focus on ) define their outputs as fixed points ( a.k.a . equilibria ) of a layer fθ and input x ; i.e. , output z⋆ = fθ ( z⋆ , x ) . Then , in the backward pass , a DEQ implicitly differentiates through the final fixed point z⋆ ( Krantz & Parks , 2012 ; Bai et al. , 2019 ; Fung et al. , 2021 ) , regardless of how forward pass is computed in the first place . Such insulated forward and backward passes enable an equilibrium model to leverage arbitrary black-box solvers to reach the fixed points without storing intermediate activations , thus consuming constant training memory . Recent works have successfully applied the DEQ framework on high-dimensional tasks such as language modeling ( Merity et al. , 2017 ) and semantic segmentation ( Cordts et al. , 2016 ) , with performance competitive with architectures like Transformers ( Vaswani et al. , 2017 ; Dai et al. , 2019 ) . However , it is also well-known that these implicit models are slow , which is ( arguably ) their single most limiting drawback compared to traditional feedforward models ( Duvenaud et al. , 2020 ; Dupont et al. , 2019 ; Bai et al. , 2021 ) . For example , Neural ODEs could take well over 100 forward solver iterations ( i.e. , evaluations of fθ ) even on MNIST classification ; DEQs can scale to realistic tasks , but the overhead of fixed-point solvers is magnified by the task scales , rendering the model 3-6× slower than state-of-the-art ( SOTA ) explicit networks ( Vaswani et al. , 2017 ; Wang et al. , 2020 ) at inference . Can we make equilibrium models faster by taking advantage of their implicitness ? One benefit of DEQ ’ s formulation is the fact that they decouple the representational capacity ( determined by fθ ) and forward computation ( controlled by the solver ) , which is not possible in any explicit model ( e.g. , ResNet-101 ( He et al. , 2016 ) ) . Hence , given a trained DEQ , one can trade off inference time and the accuracy of the estimated fixed point by simply reducing the number of solver iterations . This yields a speed/accuracy trade-off curve , as shown in Fig . 1 . However , this trade-off ( i.e. , movements along the pareto curves ) can be highly risky : as we gradually increase inference speed by compromising the quality of fixed point estimates , model accuracy also degrades drastically . In this work , we show that we can shift the DEQ speed/accuracy trade-off curve by exploiting such decoupling to customize the fixed-point solving . Prior work on equilibrium models relies on classic solvers , which are manually designed and generic ( e.g. , Broyden ’ s Method ( Broyden , 1965 ) ) . We propose a tiny , learnable , and content-aware solver module that is automatically customized to a specific DEQ . Our hypersolver consists of two parts . First , we introduce a learned initializer that estimates a good starting point for the optimization . Second , we introduce a generalized parameterized version of Anderson mixing ( Anderson , 1965 ) that learns the iterative updates as an input-dependent temporal process . Overall , the hypersolver consumes a tiny amount of parameters . Since fθ is frozen when the hypersolver is trained , the training is very fast and does not compromise generalization . Our experiments apply this approach to diverse domains with large datasets : WikiText-103 language modeling ( Merity et al. , 2017 ) , ImageNet classification ( Deng et al. , 2009 ) , and Cityscapes segmentation with megapixel images ( Cordts et al. , 2016 ) . Our results suggest that neural deep equilibrium solvers add little overhead to training ( only taking an extra 0.9-1.1 % over the original DEQ ’ s training time ) , are extremely compact ( about 1-3 % of the DEQ ’ s model size ) , and lead to a consistent and universal 1.6-2× acceleration of inference with no compromise in accuracy . Overall , we believe this paper achieves two major objectives , both vital for the quickly growing community studying implicit models : first , we advance these large-scale implicit models to a much more practical level across architectures ( e.g. , almost as fast as Transformers ) ; and second , we formally bring up and exploit this valuable notion of how implicit layers decouple representational capacity and forward computation , opening a new door to significantly advancing the agenda of deploying implicit models in practice . 2 RELATED WORK . Deep Implicit Models . Recent research on models without a prescribed computation graph or hierarchical stacking led to a new class of deep learning models where the output is defined as the solution of nonlinear systems ( Duvenaud et al. , 2020 ; Amos & Kolter , 2017 ; Chen et al. , 2018b ; Wang et al. , 2019 ; El Ghaoui et al. , 2019 ; Bai et al. , 2019 ; 2020 ; Gould et al. , 2019 ; Gu et al. , 2020 ; Wang et al. , 2020 ) . Neural ODEs ( NODEs ) ( Chen et al. , 2018b ; Dupont et al. , 2019 ) , for example , model infintesimal steps of a residual layer fθ by solving an initial value problem ( IVP ) ( Coddington & Levinson , 1955 ) parameterized by this layer ; i.e . ∂z ∂t = fθ ( z ( t ) , t ) , z ( 0 ) = x , t = 0 , . . . , T . Deep equilibrium ( DEQ ) models ( Bai et al. , 2019 ; Winston & Kolter , 2020 ) seek to directly solve for a “ fixed-point ” representation corresponding to a ( not necessarily residual ) layer fθ and input x ; i.e . z⋆ = fθ ( z⋆ , x ) . Implicit models are appealing in part due to their analytical backward passes ( e.g. , adjoint method or implicit differentiation ) that only depend on the final output , which can dramatically reduce memory consumption during training . Regularizing Implicit Models . Implicit models are known to be slow during training and inference . To address this , recent works have developed certain regularization methods that encourage these models to be more stable and thus easier to solve . For NODEs , Dupont et al . ( 2019 ) augment the neural ODE hidden state ; Grathwohl et al . ( 2019 ) use spectral normalization ( Miyato et al. , 2018 ) to stabilize the NODE dynamics ; Kelly et al . ( 2020 ) regularize higher-order time derivatives of the ODE system . For DEQs , Winston & Kolter ( 2020 ) propose a parameterization of fθ that guarantees stability of DEQ models ( i.e. , unique fixed point ) . Fung et al . ( 2021 ) show that one can simplify the implicit differentiation of Lipschitz DEQs ( Revay et al. , 2020 ) to accelerate the backward pass . Bai et al . ( 2021 ) summarize DEQ stability issues and propose to address them by regularizing the Jacobian matrices of equilibrium layers . In comparison , our work focuses on the solver rather than the layer fθ , and is orthogonal and complementary to regularization methods . Improving Implicit Model Solvers . Of particular relevance to our work are recent advances in the Neural ODE literature that improve the ODE flow solver . Poli et al . ( 2020 ) introduce a Neural ODE formulation that adds a learnable residual fitting step to the original solver steps , aiming to approximate the higher-order terms of canonical ODE solvers ( e.g. , Euler ’ s method ) on each solution checkpoint along the ODE path . Another recent work ( Kidger et al. , 2021 ) focuses on improving the adjoint method by replacing the usual L2 norm with a more flexible seminorm to make the NODE backward solver faster . To the best of our knowledge , no such solver improvement has been explored in the equilibrium model context . Unlike Neural ODEs , DEQs do not use ODE solvers and do not have unique & well-defined trajectories to the solution ( even if one starts at the same initial point z [ 0 ] ) . Our work is the first to propose a neural fixed-point solver for equilibrium models . Learning to Optimize/Learn . An important line of work has explored learnable optimization methods . Li & Malik ( 2016 ; 2017 ) propose to use reinforcement learning ( guided policy search ) to learn a new generic unconstrained continuous optimization algorithm , where the training set consists of numerous randomly generated objective functions . Andrychowicz et al . ( 2016 ) introduce the “ learning to learn ” ( L2L ) framework , where a gradient update rule for the parameters is learned by an LSTM with a pre-defined horizon T of parameter update steps . However , such approaches ( Andrychowicz et al. , 2016 ; Chen et al. , 2017 ; Wichrowska et al. , 2017 ; Ravi & Larochelle , 2016 ) have had some difficulty in generalizing to larger tasks due to the need to unroll for a large T ( e.g. , 128 ( Andrychowicz et al. , 2016 ) ) . Our work is related to these prior efforts in L2L , but differs in important ways . First , the L2L framework aims to learn a learning algorithm that will be applied to multiple models and tasks , while we aim to fit the nonlinear dynamics of a specific implicit model . Second , the optimization we tackle is not on the parameter space , but on the hidden unit space ; this means that the RNN optimizer used in L2L would not work here , because the fixed points themselves can be of variable sizes at test time ( e.g. , sequence lengths , image sizes ) . Third , while L2L methods can not know a priori what a good “ initial guess ” of optimal parameters may be , we show that it is possible and reasonable to infer this in the hidden unit space with implicit models . Concurrent to our work , Venkataraman & Amos ( 2021 ) studies an RNN-based learnable fixed-point acceleration scheme specifically in the application of convex cone programming .

Summary: The paper proposes to speed up inference of Deep Equilibrium Models (DEQs) by replacing the classic fixed-point solvers (Broyden or Anderson Acceleration) by a learned extension of AA. Their approach operates on a pre-trained DEQ, and trains a small neural network to propose an initialization and update scheme based on ground truth fixed points. Their method is orthogonal to existing regularization approaches to speeding up DEQs.  The paper has extensive experiments across large scale tasks: Language modeling, ImageNet classification, and Semantic Segmentation. They show pareto improvements across these tasks as compared to standard DEQs, while only adding a ~1% additional training overhead of the DEQ.