Neural Deep Equilibrium Solvers
Authors: Shaojie Bai
Vladlen Koltun
J Zico Kolter
Venue: ICLR 2022 Poster (2022)
Keywords: Deep learning
Implicit models
Deep equilibrium models

Abstract: A deep equilibrium (DEQ) model abandons traditional depth by solving for the fixed point of a single nonlinear layer $f_\theta$. 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\times$ speedup in DEQ network inference without any degradation in accuracy across numerous domains and tasks.

Published as a conference paper at ICLR 2022
NEURAL DEEP EQUILIBRIUM SOLVERS
Shaojie Bai
Vladlen Koltun Apple
J. Zico Kolter
Carnegie Mellon University
Carnegie Mellon University and Bosch Center for
ABSTRACT
A deep equilibrium (DEQ) model abandons traditional depth by solving for the fixed point of single nonlinear layer fe. 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-lime efficiency), which is usually via classic techniques such as Broyden's 's method or Anderson acceleration. In this paper. we show that one can exploit such decoupling and substantially enhance this fixed point computation using custom neural solver. Specifically, our solver uses a parameterized network to both guess an initial value of the optimization and perform iterative updates, in method that generalizes a learnable form of Anderson acceleration and can be trained end-to-end in an ansupervised manner. Such solution is particularly well suited the implicit model setting, because inference in these models requires repealedly solving for 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 's training time), require few additional parameters (1-3% of the original model size), yet lead to a 2x speedup DEQ network inference without any degradation accuracy across numerous domains and tasks.
INTRODUCTION
Recent progress on implicit networks, such 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 series of operators hierarchically, implicit
models define their outputs as solutions to nonlinear dy- namical systems. For example, DEQ models (which this paper will focus on) define their outputs fixed points (a.k.a. equilibria) of a layer fo and input x; i.c.,
60
Validation perplexity
DEg DEQ (regularized)
50
HyparDEQ (reg.) (ours) Transformer XL
output z" fo x). Then, in the backward pass, a DEQ implicitly differentiates through the final fixed point (Krantz& & Parks 2012 Bai et al., 2019 Fung et al.,
40
30
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 reach the fixed points without storing intermediate activations, thus consuming constant truin- ing memory. Recent works have successfully applied the DEQ framework high- dineeniiennion tasks such 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).
20
100
200
300
400
500 600 700
Inference ms/batch (size=12)
Figure 1: Pareto curves of the same DEQ with different solvers WikiText-1033 language modeling (on 1 GPU).
However, is also well well-knownn 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 fo) even on MNIST classification; DEQS scale to realistic tasks, but

---

Published as a conference paper at ICLR 2022
the overhead of fixed-point solvers is magnified by the task scales, rendering the model 3-6x slower than state-of-the-ae (SOTA) explicit networks /Vaswani et al.) 2017]] Wang et al.) 2020) at inference.
Can make equilibrium models faster by taking advantage of their implicimess?? One benefit of DEQ's formulation is the fact that they decouple the representational capacity (determined by fo) 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. D 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 's Method (Broyden (965)). We propose tiny, learnable, and content-aware solver module that is automatically customized to specific DEQ. Our hypersolver consists of two parts. First, we introduce a learned initializer that estimates 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 fa is frozen when the hypersolver is trained, the training is very fast and does not compromise generalization.
Our experiments apply this approach diverse domains with large datasets: WikiTextt-00 103 language modeling (Merity et al. (017)) ImageNet classification (Deng et al.] 2009), and Cityscapes segmenta- tion with megapixel images (Cordts et al.] (016). 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 's training time), are extremely compact (about 1-3% of the DEQ's model size), and lead to a consistent and universal 1.6-2x accelcration 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 scage- implicit models to 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 practice.
RELATED WORK
Deep Implicit Models. Recent research on models without prescribed computation graph or
hierarchical stacking led to a new class of deep learning models where the output defined as the
solution of nonlinear systems (Duvenaud et al. |020 AAms & Kolter 2017] Chen et al. 201 8b.
Wang et al. 2019
El Ghaoui et al. 2019: Bai et al. [:19]]21)2200
Gould et al. [019]] Gu et al.) 20000]
Wang et al.
2020
Neural ODEs (NODEs) (Chen et al. 2018b; Dupont et al. 2019). for example.
model infintesimal steps of a residual layer fo by solving an initial value problem (IVP) (Coddington)
1955) parameterized by this layer; i.e. Fir = fo(z(t),t) t), z(0) = x, t= t=0,... O.,,..,7' 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 fo and input x; fe(".,
& Levinson
ie. z" = fe (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, WWinston& & Kolter (2020) propose a parameterization of fo that guarantees stability of DEQ models (i.e., unique fixed point). Fung et al. (202T) show that one can simplify the implicit differentiation of Lipschitz DEQS (Revay et al. 2020) 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 the solver rather than the layer fe. and is orthogonal and complementary to regularization methods.
2

---

Published as a conference paper at ICLR 2022
Improving Implicit Model Solvers. Of particular relevance to 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., Fuler's method) on each solution checkpoint along the ODE path. Another recent work (Kidger 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 the solution (even if one starts the same initial point ()))). Our work the first to propose a neural ixed-point solver for equilibrium models.
An important line of work has explored learnable optimization methods. Li & Malik (2016) 2017) propose 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. Andrychowicc al. (2016) intro- duce the "learning to learn" (L2L) framework, where a gradient update rule for the parameters is learned by an LSTM with pre-defined horizon T of parameter update steps. However, such ap- proaches (Andrychowicz et al. 2216;; Chen et al.) 2017; Wichrowska et al. 2017] Ravi & Larochelle 2016) have had some difficulty generalizing to larger tasks due to the need unroll for large 7 (e.g., 128 (Andrychowicz et al.) 2016)). Our work is related to these prior efforts in I.2I., 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 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 cannot know a priori what good "initial guess" of optimal parameters may be, show that is possible and reasonable to infer this in the hidden unit space with implicit models. Concurrent to our work. Venkataraman & Amos (202I) studies an RNN- learnable fixed-point acceleration scheme specifically in the application of convex cone programming.
Learning to Optimize/L.eara
3 BACKGROUND. EQUILIBRIUM MODELS AND FIXED-POINT SOLVERS
Deep Equilibrium Models. Given layer (usually shallow block; e.g., self-attention (Vaswani al., (017)) fe and an input x, a DEQ model aims solve for an "infinite-level" feature representa- tion without actually stacking the fo layer infinite times. Instead, we can solve directly for the fixed point z of the system:
go x) == (x) 00
The fixed point can be estimated by quasi-Newa Neastn (or Newton's methods, which provide superlinear (or even quadratic) convergence (Broyden 1965 Anderson 1965). Subsequently, in the backward pass, one can implicitly differentiate through the equilibrium point, even without knowledge of how it is estimated, and produce gradients with respect to the model parameters 0 by solving a cooobian-based linear equation:
0e: a6
0so Ofo(z', x) 00 Oz Jz(2') O50 ((0(x)) x)
Ol
--
(1)
00
Oz
Oz.
00
Jacobian of fo
The most important message from Eq. equation iis that the backward pass can be computed with merely the knowledge of zz, irrespective of how it is found. More recently. Fung et al. (2021) prove the fcasibility of directly replacing Jz(zz) ) with -I (i.e., Jacobian free backward pass). which significantly accelerates training.
Fixed-point Solvers for DEQs. Prior works have explored number of techniques for finding
the fixed points of DEQs. For example, Bai et al. (2019) 2020); Lu et al. (2021) used Broyden's 'ss method (Broyden 1965), the memory consumption of which grows linearly with the number of iterations since all low-rank updates are stored. Other recent work (Duvenaud et al., 2020 Gilton
let al., 2021) shifted to Anderson acceleration (AA) (Anderson) 1965). lightweight solver that
is provably equivalent to a multi-secant quasi-Newtor method (Fang & Saad 2009). We briefly introduce AA here, since approach will use it as the starting point.
3

---

Published as a conference paper at ICLR 2022
Algorithm 1 Anderson acceleration (AA) prototype (with parameter B3 and m)
1: Input: initial point |1 E I,,, fixed-point function fo R R", max storage size
2: for k = O...., K do
3: 1) Set my min(m,k) k}
4: 2) Compute weights at for the past my Anderson steps s.t. 21-0 at = =1.
5: 3) z[k=1]_ Ei of To To(=!l=m=)))) (1 - 3) Σi-0 of m[-mm=
(AA_update step)
6: end for
Prototype algorithm T illustrates the main idea of Anderson acceleration: maintain size-m storage of the most recent steps, and update the iteration as a normalized linear combination of these steps with weights a, (step 3). In the canonical AA algorithm, the weights are computed in greedy manner at each step to minimize the linear combination:
a = arg min |G s.t. 1a=1, aeR
(2)
... ge(z [A)] are the past (up to 1) residuals; typically, 8 1 and m & 5. Eq. equation can be solved by a least-squares method. In all prior works with DEQs (Bai
where G(k] = [ge (2 [=))
al. 2019 2020 Winston & Kolter 2020 Revay et al. [2000] Fung et al.] [202T] Lu et al. 2021 the
fixed point iteration starts with an initial that is either 0 or a random sample from N (0,
4 NEURAL DEEP EQUILIBRIUM SOLVERS
While classic fixed-pont preed-point estimation algorithms, as presented in Section 3 already work well, they are generic and make minimal assumptions about the specific problem being solved. For example, while multiple papers in optimization literature have acknowledged that tuning (and my)) well as varying / = (==)k=...... K for each Anderson iteration i- could accelerate AA's convergence to the fixed point (Anderson) 1965 Fang & Saad 2009 Walker & Ni] 2011), this is rarely considered in practice because it's unclear what schedule should be applied to these parameters.
We propose make fixed-point solvers for DEQ models learnable and basten which is made possible by the unique properties of implicit models. First, unlike generic problems, the nonlinear system for each DEQ is uniquely defined by the input x (e.g., an image, etc.): z" (x) =z" = fe x). This opens the door to learning to make an informed initial guess, followed by content-based iterative updates in the solver. Second, due implicit models' disentanglement of representation capacity with forward computation, our goal of improving solvers is decoupled from the original learning goal of the DEQ model itself (i.e., the solver is not aware of the original task, such as to predict the class of an image). Hence, are able to train this neural solver in lightweight and unsupervised manner, directly with the help of groundtruth point solutions (see below).
4.1 GENERAL FORMULATION
For given DEQ layer fo and (possibly random) input x. we assume access to its exact fixed point z = z (x) fa(z zz,,, x), which can be obtained by taking a classic solver (e.g., Broyden's 's method) and running for as many iterations as needed (e.g.. 100 steps) to a high level of precision.
The overall structure of the hypersolver is shown in Fig. We use a tiny neural network parameterized by w= {} (explained below) learn the initialization and iterative solving process, and unroll the learnable solver for some K steps to yield prediction (x)) To train this neural solver, we minimize an objective L(w, K) (discussed in Sec. 4.2) by backpropagating through this K-step temporal process (Mozer 1989. Robinson & Fallside 1977.. The original DEQ parameters are frozen, and only the hypersolver parameters are trained here. We also do not need the groundtruth label y (c.g., the class of an image) that corresponds to input x. which means these neural equilibrium solvers can also be fine-tuned on the fly after deployment, at inference time.
The initial values can have a significant impact on the optimization process and its convergence speed. We propose to make an insut-based guess with a tiny network hg: z[(] = (x),. where @ are the parameters. Note that the goal of the initializer is not solve the underlying problem
Initializer.
4

---

Published as a conference paper at ICLR 2022
Algorithm 2 HyperAnderson Iterations (parameterized parts highlighted in color)
1: Input: initial point [0) his (x) E RR", (frozen) layer So, storage G 0 [m(++l)) with size m + 1, HyperAnderson network 8g.
2: Define gz(z) fz(z) z. Set G[O] (z(
3: for = 0 K do
4: Set my min(m, k} and Gl = G[O:(m=+ + 1)] & (m)++))x
5: Compute , Ag (Glk),, where at (ā am (m)++))
6: a = ==a (1-1'a*),
(normalization step)
mc+1
7: 2 = 1 GTG [
En a/
(same AA_up as in Alg, simplified)
8: Update G = concat (G[:], [=o(2(kt1))])
9: end for
10: Return [b+1]
Gl&+1), ( AA. 11 a. ,Bx)
Canonical Anderson Mixing
Input [ (go (z O.O.
20 1 R
Glk)
Solve af arg min aE
ally st. 1 aa =1
[K]]
Set 8 8
whorage
Iteration
(a) The original generic Anderson solver
Gl++1], [k+1] AA AA.Upeate
2M, ,8,))
Neural Equilibrium Solver (ours) Parameters w = {d.E}
Gl Demcnes
G.
Temporal Convelution
Predict a
(e(G)
[0] (x(x)) E R Input x _ (x (1) J.......
+2 [K]
Epuet m+ (estchunss)
(a comgresssed version of G1
8g
Iteration
(b) Our proposed (tiny but learnable) HyperAnderson solver
The canonical Anderson solver is based on local least-squares solution at each iteration, with B 3& set constant. 26 Our neural exed-point solver provides better initial guess 00 and learnable iterative updates.
Figure 2: 2a
at all (e.g.. to classify an image; we don't I even need the groundtruth label y). but only to yield a quick initial estimate. For example, in language modeling, where Rxxd is length-T sequence, we set
ha (x) ReI.U (Co.vldooodd where Convid Rxxd
Txp
and where W RPxq, with being the dimension of the fixed point of single token. We set to be very small (e.g.. 100), that ha tiny and fast. Note that this I-layer initializer by itself has very low expressivity and is usually poor model for the original task, as we verify in Sec. 5.3
HyperAnderson Iterations. We further parameterize the setting of 8g and at while following the AA prototype outlined in Alg. [ In lieu of setting Eq. 2 for a to a least-squares solution over the past few residuals G, we make both a E (m)++) and 8 E R explicit learnable functions of G with neural network sg (G) (mm=1)xx (R(m==1) x IR); see Alg. 2
challenge here is that (the dimension of 2 is typically large in practice, it is affected by the scale of the input (e.g.. in DEQ sequence models (Bai et al. 2019).. n is over 1.5 10 on a single textual sequence of length 200). This makes sę map from an extremely high-dimensional space to a cow-dimensional space (e.g., m = 5). To keep Sq fast, small, and applicable to inputs of varying dimensionalities (e.g., sequence length image size), we propose to first compress each go (x()) to form smaller yet still representative version G(A] of [[A] [2 (z(k-mu)) ......99(())) (x( For
example, when each go (x]*)) is image feature map residual of dimension C H W.
5

---

Published as a conference paper at ICLR 2022
#2. Compute precise fixed point - z (x) (with potentially many generic salver steps nreded)
#3. Compute with neural solver (K steps)
NeuralSolve(x. 8g)
-2 [k] 96 a.. PB. Do..=
#1. Define Hypersolver Mip $6 for each
Frozen (trained) fo
Update hypersolver by...
#4. Compute loss for the hypersolver
Gradients V) Caii. Vel(u, Ar AoV
L(w, K) = A1hrom Aghini AsCa where w = {o.
Figure 3: The training procedure of the neural deep equilibrium solver. With a given fo and input we optimize the hypersolver parameters {Ø, { via losses applied on the HyperAnderson iterations and the initializer (see Sec. 4.2)
perform global pooling to form C--dimensional vector Pool(ga(z!1)) as its compressed version: (go
GlRl = [
ggo (z (ml)=)))))
.... )....
IR (x((++))) and predict a² ,Bx = sg kG
(4)
Once we have this representative collection G(A), we treat it as mini time-series of length (m² 1) that encodes the latest estimates of the fixed point. We then apply a 2-layer temporal convolu- tion (van den Oord et al. 2016) to learn predict: 1) relative weight a for each of these past residuals i E [ and 2) the HyperAnderson mixing coefficient 3x for the current iteration. There- fore, 9g shall gradually learn to adjust these parameters and in light of the previous hypersolver steps, and receive gradients from later iterations. We explain the detailed design choices of in Appendix B while noting that it still completely captures the AA prototype (see Alg. D
4.2 TRAINING THE NEURAL EQUILIBRIUM SOLVERS
One benefit of training hypersolvers vers on implicit models is that they can be trained in an unsupervised manner via (x), which slower classic method can provide as many as needed, and for any given (possibly even random) input tensor x. Moreover, unlike NODE solvers (Chen et al., [0188) Poli al., 2020). DEQ model does not have a unique trajectory and thus its hypersolvers do not need trajectory fitting at all. All that we need is to drive everything to be as close to x as possible. As an example, a neural solver could learn to sacrifice progress earlier iterations if subsequently converges to the equilibrium faster. Formally, given a hypersolver sE) that yields a set of states (x(A), G(A] ;8k)k=0..... K (recall [! has (x)), we introduce objectives for its training.
Fixed-point Convergence Loss. The first loss aims to encourage convergence at all intermediate estimates J.-........ K of the HyperAnderson iterations: Lconv Σ wk |z( [b] z 2. where wk is the weight for the loss from iteration & such that Σ w/k = 1. We set 2'y to be monotonically increasing with ls such that later iterations apply heavier penalty for deviation from the fixed point.
Initializer Loss. We also train the initializer by maximizing the proximity of the initial guess the fixed point: Linit [| z" 12, We separate this objective from Leomv since the initialization is predicted directly from the input x and does not go through HyperAnderson updates.
Alpha Loss. Although we replace the generic Anderson solver (Anderson) 1965) in terms of how at 3h are computed in each iteration, empirically found it still beneficial to guide the hypersolvers' prediction of with an auxiliary loss especially at the start of the training: La Σk [| l2. In practice, we gradually decay the weight of this loss to as training progresses. We summarize the complete training procedure of neural solver top of DEQ in Fig.
4.3 DISCUSSION
Complexity of hypersolver. Note that So remains frozen during hypersolver training. This means that for given DEQ model fa and input the fixed point z (x) fz(z"; :x) also remains the same we are just trying to learn to find it faster, with a limited K-iteration budget. Moreover, we designed the initializer ha and HyperAnderson network sę to be intentionally simple (e.g. 1 layer with few hidden units), so that each hypersolver step is even faster than the original Anderson step, whose main computational overhead occurs in solving the constrained optimization in Eq. 2
These points also highlight the difference between the neural solver and techniques such as model compression (Han et al. 2015) distillation (iinoon et al. 2015), where a pruned/smaller (but still

---

Published as a conference paper at ICLR 2022
representationally rich) model is trained to match the output and performance of larger model. Specifically, in our case, as the fixed point z" is determined solely by fe and x, the hypersolver itself does not have much representational capacity, since its only goal is to produce guess and learnable iterations to facilitate the optimization process. E.g., the 1-layer Convld-based initializer Sec. 4.1 would be bad language model by itself since it is tiny and only sees the past tokens (see Sec. [.33 for empirical evidence), yet this limited capacity and context turn out sufficient to guide and substantially improve the solver.
"educated" initial
Training hypersolver via BPTT. While generic Anderson solver computes a* by optimizing locally with GLA), backpropagating through the HyperAnderson steps ensures that the iterative update network #g can receive gradient and learn from later iterations. This is appealing because, arguably, only the output of the Rthh iteration matters in the end. Indeed, we empirically verify via ablation studies in Sec. Sec.Etthat such learned a and 3 predictors already significantly accelerate the convergence process even without the presence of the initializer. Note that as DEQ models' fo layer is typically richly parameterized, the backpropagation- thrriih-tttouut-ieeeee (BPTT) might consume a lot of memory. To limit memory consumption, we use small batch sizes for hypersolver training. (This does not affect the training of the DEQ model itself, which is separate.) We have observed that hypersolver training is highly effective with small batch sizes, as reported in Sec. annd App. A As an alternative solution, since these hypersolvers are very fast to train in practice, one could also use methods such as gradient checkpointing (Chen et al. 2016).
Complementarity with DEQ regularizations. Besides tiny size and fast training, the value and usefulness of neural equilibrium solvers are highlighted by how DEQ models decouple representa- tional capacity and forward solver choice. In particular, our method is orthogonal to prior work that accelerates DEQ models by structural regularization of fe (Winston & Kolter 2020 Revay et al. 2020 Bai et al. [022)) or approximating the Jacobian of fo in the backward pass (Fung et al. 2021)). In Sec. we show evidence that our method (which is solver-based) integrates well with regula gggllariztion approaches (which are fø- fe-based) and yields broad improvements compared to canonical solvers (e.g.. Broyden or Anderson methods) regardless of how fo was trained or what structure it uses.
EXPERIMENTS
In this section, we verify the various benefits of exploiting neural solvers in implicit models. Specif- ically, as our goal is to show the superiority of the learnable solvers over generic solvers on both performance and efficiency aspects, we compare the movement of the entire speed/accuracy pareto curve rather than a single point on the curve. To achieve this purpose, we study the hypersolver on some of the largest- scale experiments that DEQs have been used on: WikiText- 103 language modeling (Merity et al. 2017). ImageNet classification (Deng et al. 2009), and Cityscapes semantic segmentation with megapixel images (Cordts et al. 2016). Overall, we show that: 1) neural solvers bring universal improvement over generic solvers on DEQ models in all scenarios, with a typically 1.6-2x speedup at inference and loss in performance (i.e., the new pareto curves strictly dominates old ones); 2) these hypersolvers can be trained very quickly; and 3) these methods complement prior methods such as regularizations on fg bring these implicit models to new competitive level. At the end of this section, we also conduct extensive ablative studies on the design of the hypersolver.
Note that since the neural solver training is independent of the DEQ training, we do not need to train the actual DEQ model fo itself (but could instead directly work on top of a pre-trained DEQ). Therefore, the major hyperparameters in our setting are only the relative weights of the loss objectives (see Sec. [4.2 Appendix A). We also clarify that the use of hypersolver does implicitly assume local stability around z for convergence which we find almost always holds empirically, and can be regularized for (Bai et al. 2021). Code is available at hhtpsssp. / / i thub. com /ocoo//oouss lab/ b/de
5.1 LARGE-SCALE EXPERIMENTS ON VISION AND LANGUAGE TASKS
To evaluate the the neural deep equilibrium solvers, apply them three largest-scale and highest- dimensional tasks the implicit models have ever been applied on, across the vision and language modalities. In contrast prior works (Chen et al. 20188; Winston & Kolter 2020 Bai al. 2021) that measure the number of function evaluations (NFEs), we directly measure wall-cloct inference speed under the exact same experimental settings g., input scale). We elaborate the detailed experimental settings and the implications of the results below.
7

---

Published as a conference paper at ICLR 2022
Wikitext-103 Language Modeling
ImageNet classifiication
60
Validation perplexity
DEQ (ne (Anderson)
Hyperbio reg) touns)
50
DEQ (reg.) (Andersan)
Hyperidd (wep) (ours)
Accuracy
rego HypoDEE tno regal (ours)
40
Transformess
DEQ (rep) (Anderson)
65
HypHDDO (reg. (ours) ResNet
30
60
20
55
100
200
400
600
700
50
100
150
200
250
300
350
400
Inference ms/batch (size=12)
Inference ms/batch (size=48)
(a) Wikitext- 103 language modeling
(b) ImageNet classification
80
Cityscapes semantic segmentation
Relative training overhead
Reletive model size
75
X X
300%
100%
70
80%
50%
of trainng
65
DEQ Ino reg.1 (Anderson)
60%
55
60
X WRNet-WBBw3
HypwBB0 (re rega (eurs)
10.1%
1922%
19.0%
cos
00.4%
96.1%
50
20%
45
Wiifee 321 Inagavet
O1ye0ass
WhlTec-333 Iragihst
Chacapes
40
200
400
800 1000 1200 1400
1600
Noural (Hyper0
(6spe model (6)
Inference ms/batch (size-3)
360 to ring time (S)
DEG
(c) Cityscapes segmentation
(d) Overhead comparison
Comparisons of DEQs with classic and neural solvers. All speed/accuracy curves within the same plot are benchmarked on the same GPU with the same experimental setting, averaged over 6 independent runs.
Figure 4: 4a
4c
4d The overhead of DEQ hypersolver extremely small.
WikiText- 103 Language Modeling, In this experiment, fa is a Transformer layer (Vaswani et al. 2017 Dai et al. 2019 Bai et al. 2019) and the fixed points z" are (embeddings of) text sequences. We train the neural solver on sequences of length 60 for 5000 steps, and demonstrate its inference-time effect in Figure 4a (where we use validation sequence length of 150). Specifically, compared with the original -rQ-Transformer (Bai et al. 2019) (Y curve), which uses generic Anderson acceleration (Anderson) 1965) Broyden's ss method (Broyden] 1965) (both have similar pareto curves; see App. this same DEQ model solved with our neural approach (dubbed HyperDEQ: see curve) achieves significantly better efficiency. Morcover, method is complementary to prior work that builds faster implicit models by Jacobian regularizations (Finlay et al. 2020 Bai et al. 2021). To demonstrate this, additionally train DEQ--Transforee
model with Jacobian regularization (Bai et al.] 2021) curve), and apply the neural solver on this regularized DEQ (* curve). This movement the speed/perplexity curves validates the DEQ property at the core of this paper: the decoupling of the representational capacity (i.e., fe) and the forward computation (i.e., the solver). With everything combined, we bring the performance of implicit Transformer based DEQs close to the explicit Transformer-XL. (Dai et al. 2019), which is the SOTA architecture on this task.
ImageNet classification. We additionally evaluate HyperDFR on ImageNet classification (224 224 images), customizing a neural solver on top of 4-resolutional multiscale DEQ models (Bai al. 2020). We train the HyperDEQ with 12 HyperAnderson iterations, and the speed/accuracy curves are shown in Figure 4b| (A and * curves). Note that while Jacobian regularization (o curve) eventually hurts the performance of a multiscale DEQ (cf. Y curve) due to the strong constraint imposes, the DEQ model with neural solver achieves faster inference without sacrifing any accuracy (since fa. and thus z are identical); e.g., we reach 75.0% accuracy while being almost 2x faster.
Cityscapes semantic segmentation. We also show that our neural solver approach works well in domains where existing eegularization-basee methods (see Sec. (2 fail. Specifically, we apply the neural equilibrium solver on Cityscapes semantic segmentation, where the task objective is to label every pixel on high-resolttion (typically 2048 x 1024) image with the class of the object that the pixel belongs to. As in the ImageNet and WikiText-108 103 tasks, we found that there is a consistent gain in using the neural solver over the generic alternative, accelerating fixed-point convergence by more than factor of 2 (see Figure 4c). In contrast, prior methods such as Jacobian regularization (Bai et al.) 2021) do not work in this setting, due to their dependence on the exact structure of fe. (Specifically, when fe is convolution- based and the image is very large, Jacobian regularization that encourages contractivity is at odds with the gradual broadening of the receptive field.) Our neural solver is orthogonal to the structure of fo (which is frozen), and only improve how the solver functions.
8

---

Published as a conference paper at ICLR 2022
5.2 TRAINING EFFICIENCY OF THE NEURAL SOLVER
We also provide extra training analysis in Fig. 4d Not only is our approach effective, but the overhead for training the neural solver is also extremely small: the neural solver module is tiny (< 4% of the DEQ model size) and requires only about 1% of the training time needed by the original DEQ model (e.g., on WikiText- 103, a DEQ requires 130 hours on 4 GPUs; the neural solver requires only about 1.2 extra hours). We believe this strong evidence that neural solvers are simple, lightweight, and effective tools that take advantage of the decoupling properties of equilibrium models to yield an almost-free acceleration at inference time. We also perform convergence analysis in App. D
Interestingly, one can also employ the neural solver to accelerate the DEQ training, but with three caveats: 1) during training the fixed point manifold also keeps changing: 2) we want to amortize the cost of computing "groudtruth" z"; and 3) we still keep the backward implicit differentiation intact. Thus, we propose to train the neural solver {ho $z} and the DEQ model fo in an alternating manner, and elaborate more in App. D We empirically observe this leads to 16-20% DEQ training speedup.
5.3 ABLATIVE STUDIES AND LIMITATIONS
Finally, we perform a series of ablation studies to understand the benefits of multiple components within our design of the neural equilibrium solvers. We use the language modeling task on WikiText-103 for this purpose (where fe is Transformer layer), while noting that we've noticed similar trends in all other settings. The results are presented in Fig. [ The HyperDEQ with everything combined (initializer, a², and 3e predictions) performs best. Making the An- derson iterations learnable generally improves convergence. Moreover, although simply adding an initializer to a generic solver (+ curve) does not help much, learning and backpropagating through the HyperAnderson iterations makes the initializer quite useful (cf. and * curves). We additionally take the learned initializer ha from Hyp and verify that this tiny module by itself still poor language model (see Table Sec. 4.3). but is valuable to our HyperAnderson iterations. More ablation studies (c.g.. how is predicted) are reported in Appendix C
Table I: Perplexity (ppl) on WikiText- 103
Model Test ppl
Gated ConvNet Dauphin et al. 2017
230M
37.2
Transformer-X XL Dupont ot at 2019
165M
24.2
HyperDEQ (reg.) w/ 12 iters (ours)
98M
23.4
Initializer h2 (Convld)
0.4M
836.94
Wikitext- D3 Language Modeling (Ablative)
50
Validation perplexity
45-
DEQ (reg.) (Anderson)
DEQ (reg.) [Rroyden)
HsporDE0 (rog. only int
40
Hypword (reg)
HsperOEO onl
35
Hlywerro (reg onny S. and in
Hyp er (rog teurs)
30
25
100
150
200
250
Inference ms/batch (size=12)
Figure 5: Ablative studies on HyperDEQ (reg.).
We also note two caveats for our approach. First, as mentioned in Sec. 4.3 backpropagating through the HyperAnderson iterations means the memory could grow with the number of steps K that run for. However, we don't 't find this to be problematic in practice, as we observed the training of these hypersolvers to be very insensitive batch size and that at inference time hypersolvers do easily generalize to iterations > K (see also App. D). Second, though our method brings consistent improvements over generic solvers, in some cases a certain amount of iterations may still be required for good performance (c.g.. fa is a 3 x 3 convolution and the input is a large image).
DISCUSSION
We introduce a neural fixed-point solver for deep equilibrium (DEQ) models. The approach is simple, customizable, and extremely lightweight. Unlike prior works that regularize the structures or parameterizations of the implicit layer design (usually at the cost of accuracy), we propose to exploit this valuable notion of how implicit models decouple the representation (i.e., fo) from the forwards computation. We directly learn model-specifh equilibrium solver that provides: 1) better-informed initial guesses; and 2) parameterized iterations that generalize Anderson acceleration and take into account future steps. Our experiments show that these modifications substantially improve the speed/accuracy trade-off across diverse large-scale tasks, while adding almost no overhead to training. We see these encouraging results as a significant step towards making implicit models more practical, and hope that this work will further motivate the application of implicit models such as Neural ODEs, DEQs, and other variants (Gu et al. 2020 Wang et al. 2020) to real, large-scale datasets.

---

Published as a conference paper at ICLR 2022
7 ACKNOWLEDGEMENT
Shaojie Bai is supported by grant from the Bosch Center for Artificial Intelligence (BCAT). We thank Brandon Amos for the helpful discussions.
REFERENCES
Brandon Amos and J, Zico Kolter. OptNet: Differentiable optimization as a layer in neural networks. In International Conference on Machine Learning (ICML), 2017.
Donald G Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM (JACM), 12(4):547-560, 1965.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew Hoffman, David Pfau, Tom Schaul, Brendan Shillingford. and Nando De Freitas. Learning to learn by gradient descent by gradient descent. In Neural Information Processing Systems, 2016.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Neural Information Processing Systems, 2019.
Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale deep equilibrium models. In Neural Information Processing Systems, 2020.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Stabilizing equilibrium models by Jacobian
regularization. In International Conference on Machine Learning (ICML), 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. arXiv:2005. /4165, 2020.
Charles G Broyden. A class of methods for solving nonlinear simultancous equations. Mathematics of Computation, 1965.
Liang-Chie Chieh Chen, George Papandreou, Iasonas Kokkinos, Kevin Murphy, and Alan L. Yuille. DeepLab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected CRFs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(4), 2018a.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differen- tial equations. In Neural Information Processing Systems, 2018b.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv: 1604.6604611100 174, 2016.
Yutian Chen, Matthew W Hoffman, Sergio Gómez Colmenarejo. Misha Denil, Timothy P Lillicrap. Matt Botvinick, and Nando Freitas. Learning to learn without gradient descent by gradient descent. In International Conference on Machine Learning (ICML), 2017.
Bowen Cheng, Maxwell Collins, Yukun Zhu, Ting Liu, Thomas S Huang. Hartwig Adam, and Liang-Chieh Chen. danoptic-deeplab: A simple, strong, and fast baseline for bottom-up panoptic segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12475-1 12485, 2020.
Earl A Coddington and Norman Levinson. Theory of ordinary differential equations. Tata McGraw-
Hill Education, 1955.
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The Cityscapes dataset for semantic urban scene understanding. In Computer Vision and Pattern Recognition (CVPR), 2016.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-F XL: Attentive language models beyond a fixed-length context. In Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
10

---

Published as a conference paper at ICLR 2022
Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated
convolutional networks. In International Conference on Machine Learning (ICML), 2017.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei Li. ImageNet: A large-scale hierarchical image database. In Computer Vision and Pattern Recognition (CVPR), 2009.
Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2019.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11999, 2020.
Emilien Dupont, Amaud Doucet, and Yee Whye Teh. Augmented neural ODEs. In Neural Information Processing Systems, 2019.
David Duvenaud, J. Zico Kolter, and Matthew Johnson. Deep implicit layers tutorial neural ODEs, deep equilibirum models, and beyond. Neural Information Processing Systems Tutorial, 2020.
Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, and Armin Askari. Implicit deep learning, a9Xiv: 06315, 2019.
Haw-ren Fang and Yousef Saad. Two classes of multisecant methods for nonlinear acceleration. Numerical Linear Algebra with Applications, 16(3):197-211, 2009.
Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M Oberman. How to train your neural ODE. arXiv:2002.02798, 2020.
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. Fixed point networks: Implicit depth models with Jacobian-free backprop. 20Xiv:2103./2803, 2021.
Davis Gilton, Gregory Ongie, and Rebecca Willett. Deep equilibrium architectures for inverse problems in imaging. arXiv:2/02.07944 2021.
Stephen Gould, Richard Hartley, and Dylan Campbell. Deep declarative networks: A new hope. 1etp::999900866666 2019.
Will Grathwohl. Ricky TQ Chen, Jesse Betterncourt, Ilya Sutskever, and David Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. In International Confer- ence on Learning Representations (ICLR), 2019.
Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui. Implicit graph neural networks. In Neural Information Processing Systems, 2020.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv:1510.00149, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition (CVPR), 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. 10Xiv:1503.02531,,0 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8), 1997.
Junteng Jia and Austin Benson. Neural jump stochastic differential equations. arXiv: 10403. 2019.
Kenji Kawaguchi. On the theory of implicit deep learning: Global convergence with implicit layers. In International Conference on Learning Representations (ICLR), 2021.
Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, and David Duvenaud. Learning differential equations that easy to solve. In Neural Information Processing Systems, 2020.

---

Published as a conference paper at ICLR 2022
Patrick Kidger, Ricky TQ Chen, and Temry Lyons. "hey, that's not ODE": Faster ODE adjoints with 12 lines of code. In International Conference on Machine Learning (ICML), 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Steven G Krantz and Harold Parks. The iuniiii theorem: History, theory, and applications. Springer, 2012.
Ke Li and Jitendra Malik. Learning to optimize. arXiv:1606.01885, 2016.
Ke Li and Jitendra Malik. Learning to optimize neural nets.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations (ICLR), 2017.
Cheng Lu, Jianfei Chen, Chongxuan Li, Qiuhao Wang. and Jun Zhu. Implicit normalizing flows. In International Conference Learning Representations (ICLR), 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations (ICLR), 2017.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations (ICLR), 2018.
Michael C Mozer. A focused back pro-prropggttion algorithm for temporal pattern recognition. Complex Systems, 3(4))349-38)) 1989.
Matthew Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations, arXiv preprint arXiv:1802.05560 2018.
Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Hypersolvers: Toward fast continuous-depth models. arxiiv2007.09601 2020.
Sachin Ravi and Hugo Larochelle. Optimization model for few-shot learning. In International Conference on Learning Representations (ICLR), 2016.
Max Revay, Ruigang Wang, and Ian R Manchester. Lipschitz bounded equilibrium networks. arXiv:2010.01732, 2020.
AJ Robinson and Frank Fallside. The utility driven dynamic error propagation network. University of Cambridge Department of Engineering Cambridge, MA, 1987.
Yulia Rubanova, Ricky TQ Chen, and David Duvenaud. Latent ODEs for iregularly-sampled time series. arXiv:1907.033: 0030999777 2019.
Aäron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W. Senior, and Koray Kavukcuoglu. WaveNet: A generative model for raw audio. 1rXiv: (09.0)99999999,.... 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems, 2017.
Shobha Venkataraman and Brandon Amos. Neural fixed-point acceleration for convex optimization. arXiv preprint arXiv:2107. 2021.
Homer Walker and Peng Ni. Anderson acceleration for fixed-point iterations. SIAM Journal on Numerical Analysis, 49(4): 2011.
J. Wang. K. Sun, T. Cheng, B. Jiang, C. Deng, Y. Zhao, D. Liu, Y. Mu. M. Tan, X. Wang, W. Liu, and B. Xiao. Deep high-resoluttin representation learning for visual recognition. IEEE Transactions Pattern Analysis and Machine Intelligence, 2020.
12

---

Published as a conference paper at ICLR 2022
Po- Wei Wang. Priya Donti, Bryan Wilder, and Zico Kolter. SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver. In Internationall Conference on Machine Learning (ICML), 2019.
Tiancai Wang, Xiangyu Zhang, and Jian Sun. Implicit feature pyramid network for object detection. arXiv preprint 20X2v:0012.35560 2020.
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In International Conference Machine Learning (ICML), 2017.
Ezra Winston and J. Zico Kolter. Monotone operator equilibrium networks. In Neural Information Processing Systems, 2020.
Zhilin Yang. Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint 17X1v:1906.00237,, 2019.
13

---

Published as a conference paper at ICLR 2022
Table 2:
Task settings (see Sec. 5) Note that in WikiText-103 and ImageNet, we train the neural solver only for a few thousand gradient steps, which is far less than complete epoch on these datasets. "In addition, we decay the loss weight X (for La) to 5e-8 on a linear schedule over the first 1.5-2K training steps (see below).
Task
Language modeling
Image classification
Semantic segmentation
Dataset
WikiText 103
Merily al. 2017
ImageNet Deng
2019
Cityscapes Cordus ct 2016
Download link
Link
Link
Link 2975/500/552 (images)
Split (trtin/val/test)
103M/218K72466 (words)
1.28M/-
TSSK (images)
Vocabulary size
267,735
Not Applicable
Not Applicable Image
Input type
Text Sequence
Image
Implicit model arch.
Multiscule-DER
Multiseale-DEs
Input scale (train)
Length=60
W= 224
W=2048 1024 H W W=1024 512 (7.86- 10") (7.86- 10))
Input scale (train)
Length=150
H x W=224 224
J, (x') size (per sample)
(1.55-0p* 10") (1.05- 10*)
(1.88 10") (1.88 10")
Batch size (train)
16
Optimizer (Ir)
Adam (0.001)
Adam (0.001)
Adam (0.001)
HyperAnderson (train)
10
HyperAnderson storage
5
Training steps 7
5000
4000
3000
Loss weight A (for Camm)
0.1
0.1
0.1
Loss weight Ay (for Lux))
0.05
0.02
0.02
Looss weight Ay (For C..)
le
1e:5
A EXPERIMENTAL DETAILS
We describe here (and summarize in Table 2 the datasets, pretrained models and the training settings (including byperparameters) for the experimental results reported in Sec. We note that overall, since our approach is orthogonal to how these DEQ models were trained in the first place, we directly downloaded the pretrained DEQ models from the publicly released DEQ repo (Bai al. (2019) (for language modeling) and the MDEQ repo (Bai et al. 2020) (for image classification and segmentation), which our code base is built upon. All of our experiments were conducted on NVIDIA RTX 2080 Ti GPUs. The pareto efficiencies were benchmarked on 1 GPU, averaged over 6 independently trained hypersolvers (i.c., 6 random seeds) for each task.
A.1 WIKITEXT-10 103 WORD-LEVEL LANGUAGE MODELING
In a language modeling task, an autoregressive network is trained to predict the next token given a sequence of the past ones. Large-scale contextualized language models have been widely studied in literature, and is the cornerstone of many recent progress in realistic natural language AI models; e.g., ELMO (Peters et 2018), BERT (Devlin et al. 2019), XLNet (Yang et al.) 2nd GPT- 3 (Brown et al. 2020). Formally, given an autoregressive network F (like a Transformer (Vaswani al. (017)) or LSTMS (Hochreiter & Schmidhuber 1997)) and input sequence X1:t, the prediction YI: = Fx IR should be ideally as close as possible to the groundtruth "next words" X2++1+ with y: R (i.e., like an N-way classification, where N is the vocabulary size). Our experiments follow the exact same settings as the TE@--Trnsformer network in Bai et al. (2019) where a deep equilibrium (DEQ) model is as sequence model and fo single Transformer block (i.e., multi-head self-attention (Vaswani et al. 2017 Dai al. 2019)): i.e.,
1 hz)j.) where zize fo x1:), and a linear layer.
(5)
We use one of the most commonly and the largest-scale textual corpus, Wikitext-109 (Merity et al. 2017), to evaluate our method. This dataset can be downloaded at this link Specifically, the Wikitext- corpus contains over 103M words in its training split, and 218K/246K words for validation/test. Morcover, the entire corpus contains N 267, 735 unique words, which retain rare words, numbers, punctuation and case (as opposed to all text being lowercased) from the original Wikipedia articles.
Note that our approach only introduces minimal new hyperparameters (as the original DEQ model parameters are frozen). For the language modeling task, use Adam optimizer (Kingma & Ba 2013) with start learning rate 0.001 and cosine learning rate annealing (Loshchilov & Hutter 2017). The neural solver is trained for 5000 steps, with sequences of length 60 and batch size 10, top of
14

---

Published as a conference paper at ICLR 2022
pretrained DEQ with word embedding dimension 700. We set the HyperAnderson iteration limit K to 10 (i.e., the HyperAnderson module unrolled for steps at training). while evaluating the hypersolver at validation time for up to K 16 steps for the pareto efficiency curve benchmarking (see Fig. 4ap.
A.2 IMAGENET CLASSIFICATION
In both the Imagenet classification and the Cityscapes segmentation tasks (see below), we use a pretrained multiscale-DEO (MDEQ-smaal, with 4 resolutions) model (Bai al. (0200) for training our neural equilibrium solvers. In this case, feature maps maintained at multiple (c.g.. n)) resolutions driven their equilibria simultaneously; i.c.,
z
[ 2 =
(6)
To verify our approach, directly train and test on the largest vision datasets trained and reported for implicit models. The ImageNet dataset (Deng et al. 2009) contains 128,1280 training images and over 150K test images of resolution 224 224, which are distributed over 1.000 classes. Notably, as the DEQ model is already trained and frozen, we do not perform data augmentations (c.g., random cropping or horizontal flipping, which are standard practice for training models) to the input images, and make sure that the train/test inputs undergo the exact same transformations. The neural solver is trained for 4000 steps, each with 10 HyperAnderson iterations.
A.3 CITYSCAPES SEGMENTATION
The other, even higher-resolution, vision dataset we apply the hypersolver on is the Cityscapes dataset (Cordis et al. 2016). This is large-scale dataset that contains a diverse set of high- stereo video sequences recorded in street scenes from 50 different cities (at different times of the day, spanning several months), with pixel-level annotations of 5,000 frames over 19 classes (e.g., road, person, rider, truck, bus, etc.). Specifically, these 5,000 frames are divided into 2.975 training, 500 validation and 1 1,525 test images. The multiseale-DEQ (MDE)) as in the ImageNet case above, with 4 resolutions driven to the fixed-point simultaneously.
follow the exact formulation
The Cityscapes pixel-level semantic segmentation task is also standard widely-used bench- mark (Chen et al., 2018a Wang et al.] 2020 Cheng et al., 2020), and is especially challenging because of its rich contents and high resolution. To train the neural solver, in ImageNet, we do perform data augmentations like horizontal flipping (but perform random cropping to 1024 x 512) on the training set, where we set the HyperAnderson step limit to 12 and the batch size to 6. The neural solver is then evaluated in the standard full-resolution setting, on Cityscapes val 2048 1024 images.
A.4 GROUNDTRUTH SOLVERS AND LOSS WEIGHTS
In of the experiments, we use the canonical Broyden's 's method (Broyden 1965) with 60 standard quasi-Newton steps as the "groundtruth" solver that yields the exact fixed point solution. However, overall we found the choice of this high-precision solver not important as long as we run it for enough number of steps to ensure proper convergence (e.g.. canonical Anderson acceleration algorithm with steps would also suffice, since it's the same groundtruth fixed point).
The major hyperparameters introduced by the neural solver are the relative loss weights. For the K fixed-point convergence loss Lcom Σk=1 wk [x] ||2 (see Sec. Sec.4], we apply weights wy 0 to all intermediate HyperAnderson steps, with Wy monotonically increasing (such as later HyperAnderson steps are applied heavier penalty for deviation from the fixed point) and sum to 1; e.g., we can simply do
A
Wk
ΣE1
Note that other similar formulations possible, such wp Σx 5 and we generally do not find them empirically make a huge difference. We further illustrate this in the ablation studies presented in Appendix C
2²
15

---

Published as a conference paper at ICLR 2022
Initializer he
Hyper Anderson Iteration (compression step)
Initial guess
Gl
(
Pact resibul
Linear()
qute
(A compressed
Lineer RelU
verion of GPM,
Seleet the bost
independent impert snale T)
d
Causal Convid (lornel 3)
(meed cuntest eumlessiiniion unlisod)) and peoject
1
"The of farmer... in 1932.
Input (cobbeldings) x R
(a) Visualization of TEQ-Transformer (for WikiText- language modeling) initializer and HyperAnderson compression module.
Initializer h&p
HyperAnderso Iteration (compression step)
G
[
Conn2d (kernel size 3) + ReLU
aeemeee [+ u Fusa
GlobelPool2di
Initial guess x0
compreseed version GIM, independent impat scale HF, W)
medittio fratums, the global pasi.
Inpot (injection) x RxxHx
(b) Visualization of multiscale-DDi -DE (for ImageNet classifications and Cityscapes segmentations) initializer and HyperAnderson compression module.
Figure 6: Visualization of the neural equilibrium solver design. We use a local context and shallow layer to build the initializer, and perform global pooling over the input (ID sequences 2D images) create compressed (but representative) versions of the iterate G(R].
We applied a simple grid search on the relative loss weights Å1 A2 and A on the three losses, finding it overall beneficial to adjust these weights that the 3 losses the start of the training be at roughly the same magnitude/scale. Moreover, we empirically find it useful to set X3 larger initially encourage Anderson-like iterations (ef. Eq. equation (2). but later gradually decay this weight to almost 0, so that the hypersolver can more flexibly exploit the a in each HyperAnderson iterations. In all experiments, we decay λ 5e-8 linear schedule over 2,000 steps.
B ADDITIONAL DISCUSSIONS ON THE NEURAL EQULIBRIUM SOLVER DESIGN
We briefly illustrate in this section the designs of the initializer and the HyperAnderson iteration, and how we make them both lightweight and expressive (of the feature map). This implies two design principles that affect the model we use: 1) the hypersolver module itself should have an minimial number of parameters possible, which implies projections to oowddimessions; 2) it should be also invariant to input dimensionality while being able to cover the entire input. The second point is especially important since the original implicit model can be applied on inputs of various scales (e.g.. sequences of length, 100, 200, etc.; same for images), and so it is preferable that the neural solver is also adaptive to such variety. We note that this rules out the learning- to-lerr (L2L; see Sec. 2 way of parameterizing optmizers by RNNs: whereas the parameter space has fixed dimensionality, hidden unit space does not.
HyperDEQ for sequences. In the sequence modeling task, as mentioned in Sec. our initializer he is simply ID causal convolution followed by non-linear projection (see Fig. 6at
(x) ReLU((onnlld==()))
where Convidk-3 IR. Txd - Txy (1D causal convolution with kernel size 3)
16

---

Published as a conference paper at ICLR 2022
Wikitoxt- 103 Language Modeling (Ablative)
Wilkbeet- 103 Language Modeling (Ablative)
36
36
DEQ rego (Andersoni
DEQ reg. (Anderson)
DE (no reg) (Broyden)
DCC (no reg) (Broyden)
34
HypurDE (ne req.) (e softmant:))
34
hyperDEO (no reg.) de - KTI
Hypebtt reg) unnarmalieed)
Hyperoro ino regó) ouniform
Validatinn perplexity
32
HypweDE (ow rng) (ours)
Validation perplexity
32
hyperDEO repà
HyperBEO ine rego the coo)
HyperDEO ine rep)) tho intialberi
30
30
HyperDEo (ours)
26
28
26
26
24
24
220
240
260
260
300
320
340
360
380
400
220
240
260
280
300
320
340
360
360
400
Inference ma/batch (sizo=12)
Inference m5/batch (siro122)
(a) Ablations on the importance of a normalization and learning. A curve is our proposed method.
(b) Ablations on loss components and relative loss weights.
curve is our proposed method.
Figure 7: Further ablations on the neural solver design (in terms of the prediction and loss components, Note that we use an unregularized DEQ here (in contrast to Fig. 5 to better demonstrate the curve differences, which could sometimes be small.)
In particular, we set p < d (c.g., in the WikiText- 103 experiment, = 700 and p 100) so that this module requires little parameters and is fast to evaluate. We reiterate here that this is possible because the goal of the initializer is not to solve the original problem (e.g.. high-dimensional language modeling), but only to give "reasonable" initial guess to the fixed point solver. In other words, the initial guess itself may be bad in a task perspective (as we verified in Table D but good in an optimization perspective (see Fig. 5)
Moreover, compress the past and current residuals G(k] (a. ........' [x)
R(m+1)xTxq by simply taking the last and most contextualized token of each sequence; i.e., (z)(z)))) R7. This is reasonable because the design of a sequence model like Transformer layers (which parameterizes our So for DEQ) already ensures that the last token contains most of the contextual information of the entire sequence (and this technique is frequently used for textual classification tasks, or in vision transformers (Dosovitskiy et al., 2020)). We therefore propose to compress the past residuals of the fixed point estimates by directly using the (a projection of the) very last tokens (see Fig. 6a which are collected and treated as a mini time-series as mentioned in Sec. F
HyperDEQ (multiscale) for vision. The initializer for multiscale DEQ models is similarly de- signed, with series of small-conter 1-layer convolutions (e.g., kemel size 3) to generate the initial guess for each resolution stream (see Fig. (b)) To compress the past residuals, we leverage the fact that all resolutions are fused in each layer of the multiscale DEQ design, and only take the lowest-
resolution feature map zm of the past fixed point estimate z [z| ..... =I:] (in practice, use n = 4). As shown in Fig. 66, we directly apply global pooling to pass into the later HyperAnderson
steps (by treating as a mini time-series, like in the sequence case).
C ADDITIONAL EFFICIENCY AND ABLATIVE STUDIES
We further study some ablative settings not covered in Sec. 5.3 here. We specifically focus on two aspects: 1) how is predicted; and 2) the effects of different loss components.
In the original AA prototype, the values of a is determined by solving a lcast-squared solution over the past m steps, as described by the constrained optimization problem equation 2 In particular, . values can be any real value, but has to sum to 1. In our hypersolver, we
these a
propose to "normalize" the predicted a by a shift: a = k² + (1-1"a") 1. We here compare two I
alternatives: no normalization at all (i.e., have a learned freely; see in Fig. 7a). and softmax-bar based normalization (i.e., requiring a to be all- all-positive; see in Fig. 7 Specifically, we find that the exact choice of a normalization does not affect the overall substantial pareto curve improvement (e.g., compare them with the Y curve baseline), but ensuring that the a values can be negative while
17

---

Published as a conference paper at ICLR 2022
Fixed-point Convergence (WT103,, Jacobian reg.)
Training WT103 Perplexity Time
0.40
50
Canorical Andarson
Validation Set Perplexity
DEQ Ireg.)
Relative residual (ert -1-
Neural Solver (rained on K 6)
45
HyperDER (reg. (ours)
0.30
0.20
35
0.15
0.10
30
0.05
25
0.00
8
10
50
100
150
200
Fixed-politt solver step
Training Time w/ GPUs (Hours)
(a) The neural equilibrium solver improves upon the convergence path of canonical Anderson acceleration. The speedup HyperDEQ is result of this faster convergence and the more lightweight solver step.
(b) Hypersolver can also be used at training time to further accelerate the DEQ model training by alter- nating their training (in addition to existing methods like Jacobian stabilization).
Figure 8: Left: Convergence analysis of the hypersolver at inference time. Right: Although trained with a frozen fo, the neural equilibrium solver can also be used to accelerate DEQ training.
still normalized (i.e., sum to I) overall benefit the performance. For instance, with the same inference speed, we get a 1-1.5 perplexity improvement.
We also analyze the effect of different loss components in Fig. 7b For the main fixed-point conver- gence loss Lcomv. we compare two alternatives of the relative weights applied on the intermediate HyperAnderson steps: only apply w/ 1 at the final output K (see + curve), or set w/ to be uniform for all & (see + curve). We empirically find that the hypersolver trained in both scenarios perform well, but the monotonically incresing wi perform best. Morcover, removing either the initializer loss Linit (note that we still keep the initializer he itself, just don't 't supervise it; see the 1) or the alpha loss (the curve) La impacts the performance. Interestingly, removing the Liit loss on the initializer yields even worse performance than even removing the initializer itself (the curve). However, we note that all of these ablative settings still significantly outperform the speed/accuracy efficiency than the generic solvers, which suggests the overall benefit of using customized learnable solvers for implicit models.
(i.e., putting larger emphasis on later iterations)
Finally, we provide additional empirical evidence on the convergence and generalizability of the neural solvers in Fig. Sa where we compare the convergence of a pre-trained, regularized Transformer- based DEQ model under 1) canonical Anderson acceleration; and 2) a neural solver (with initializer) trained to unroll for K = 6 HyperAnderson steps. The shaded areas around the curves correspond to the standard deviation of the relative residual at each solver step, measured on 1000 randomly selected textual sequences of length 100. From Figure Sa] we first note that a hypersolver trained with K unrolling steps is able generalize (i.e., converge well) beyond K, although both solvers eventually plateau at some of e residual value. Second, with the parameterized fixed-point solving, we see that the neural solver consistently improves over the canonical solver's 's convergence, while being more lightweight to compute (i.e., each solver neural solver step is cheaper than canonical Anderson step), which account for the 2x speedup we observe across the board for DEQ model inference efficiency in Sec. 5
D USE HYPERSOLVER FOR DEQ TRAINING
turns out we could also leverage its acceleration power to improve the not just the inference but also training of the original DEQ model (i.e., replace the classical solver at DEQ training with a hypersolver). However, we note three major caveats. First, as the DEQ model is being trained (e.g.. by SGD), parameterization of the layer fo also changes, which means the fixed-point manifold of this implicit model constantly gets updated. Therefore, a neural deep equilibrium solver (ha, $e) trained on parameter state fe, may gradually deviate from accurate and stable convergence as fa also deviates from fe, as a result of more and more parameter updates. Second, precisely computing the groundtruth fixed points z" with classical solver at every training step would add heavy burden to the training overhead. Third, the usage of hypersolver to solve forward pass fixed point does not affect the backward pass implicit differentiation (notably, decoupled forward and backward trajectories is
18

---

Published as a conference paper at ICLR 2022
also major feature of implicit models), implying the training efficiency improvement will be much smaller than that at the inference time (about 2x, as mentioned in Sec. (5)
Despite these challenges, exploiting the fact that model parameters 0 drift only gradually with training iterations (i.e., at training step t and + 1. we assume |0;- 0.+1 is small), we propose to train the neural solver (he st) (which are tiny, lightweight, and have little representational capacity) and the implicit layer fe (which are huge, representationally rich, and requires many training steps) in an alternating way. Specifically,
adopt the following procedure:
1. Warmup and train the DEQ model for some small number of steps.
2. Take snapshot of the current implicit layer fo and train from scratch hypersolver ( st} customized on top of it.
3. Use the current hypersolver (hr, se) to solve for the fixed point and train the DEQ model for the next M steps.
4. Take a snapshot of the current implicit layer fe. and fine-tune the existing hypersolver {hg, top of it by minimizing [(w, K) for some T hypersolver training steps.
5. Repeat step 3 & 4 until we reach max training steps for the DEQ model.
Empirically, we find that the proposed neural solver generalizes surprisingly well for values of M that are not too large (i.e., even though fa gets updated, the hypersolver can still reliably converge to the correct fixed points during the next M DEQ training steps). Moreover, as previously mentioned, training neural solver is extremely lightweight, and even more so for fine-tuning it, suggesting that T is very small in pratice. For example. on WikiText- 103 language modeling with a transformer-basee fo (which has 300K DEQ training steps), we set M 5000 and T 200; other words, we fine-tune the hypersolver for merely 200 steps (again, on very small batch size and sequence length; see Table 2 after every 5000 DEQ training steps, in order to keep the neural solver "up to date" with the implicit layer. This extra (but minor) cost of constantly fine-tuning hypersolver is counterbalanced by the huge gain in using hypersolver over a canonical generic root solver.
Fig. 8b shows the training efficiency improvement result of adopting hypersolver in DEQ training. With the hypersolver fine-tuned intermittently (green curve), we are able to reduce the total training time over that of the classical-solver solea----lvr-bbeeed DEQ counterpart by about 20% (purple curve). Note that just like at inference time, this improvement via hypersolver is orthogonal (and thus additional) to the acceleration effect of regularization-basel methods, like Jacobian regularizations (cf. purple and pink curves).
E ADDITIONAL EXPERIMENT ON TINY EQUILIBRIUM MODELS
While Fig. 4 provides strong evidence on how the hypersolver could improve the speed/accuracy
pareto efficiency of these large-scale DEQ mod- els (e.g., DEQ-Transformer contains almost 100M parameters). we also perform additional exper- iments to validate its efficacy on much smaller networks. Specifically, we train hypersolver on a
CIFAR-10 classification (no augmentation)
88
86
84
82
DEG treg. (Anderson))
80
HypwrDEC (reg)) (ours) ResMet-18
tiny MDEQ model that is pretrained on CIFAR-10 10 classification task (the same setting as in Bai et al.
78
76
(2000), but with Jacobian regularization) which contains only 170K parameters. The hypersolver contains about 8K parameters (including the ini- tializer ht), and is trained for only 1000 training steps. The results are shown in Fig. [ The im- provement in pareto frontier again suggests that the neural solvers provide lightweight and yet effective solution to improving these implicit models.
74
72
10
20
30
50
Inference ms/batch (size=16)
Figure 9: The benefit of hypersolver holds on tiny implicit models. Inference speed is bench- marked with batch size 16, averaged over 5 runs.
19


Reviews:
[{"review_id": "iMDXiIG4sbz", "reviewer": "Reviewer_eZVG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The technique is intuitively motivated as a learnable version of Anderson acceleration, which while lacking theoretical basis, is easy to follow and a natural method to try.\n- The experimental evaluation is convincing. The extensive empirical evaluation of the proposed approach gives the reader the sense that the hypersolver is at least worth trying on their DEQ model. \n\nWeaknesses:\n1. Being a largely empirically/intuitively motivated work, there are many user settings and hyperparameters (and even hardcoded arbitrary choices, e.g. the 3 losses introduced in section 4.2) that have to be chosen within any theoretical guidance and are not well-motivated. Still, this work opens the door for future mathemtical analysis of the proposed method, although perhaps on a more restricted or smaller practical/empirical scale.\n2. In L_conv, is z^\\ast necessarily a true solution, or just one provided by another solver? None of the losses in section 4.2 actually minimise the ostensible goal of the equilibrium solver. Why is there no loss that directly minimises the absolute value of the residual |g|? It seems like the neural solver is learning to imitate the behaviour of a provided solver (given access to z^\\ast, which may not actually be a root, nor does a root necessarily exist, nor is it guaranteed to be unique), rather than solve the problem directly. Did you try to minimise |g| instead? Any reason or intuition as to why you chose this method? \n\nComments/suggestions/questions:\n\n3. One of the really nice features of your technique is that (as you mention) the fixed point solver does not need training labels in order to update its parameters. This ability to do unsupervised training could be highlighted in the abstract.\n4. Typo. \"...we treat the it as a mini time-series of length...\"\n5. Typo. \"...2) these hyprtsolvers can be trained very quickly...\"\n6. Figure 4d. All of the interesting part of the graph occupies 99.6% to 100%. Remove the bottom 99.6% of the vertical axis. Maybe even a logarithmic scale is appropriate.\n7. It may be present in the appendix but I was not able to easily find it. How many random seeds do the curves in Figure 4 represent? (hopefully not one...!) Is each point independent from each other (does each point use different random seeds), both within the same colour and outside of the same colour?\n8. Can you try your method on some smaller networks (e.g. around 1,000-100,000 parameters) and simpler tasks to see if an advantage over Anderson acceleration still persists? I am interested in understanding whether the success of this method is tied to the difficulty of the fixed point problem, which should intuitively be greater in larger models.\n", "summary_of_the_paper": "The authors introduce a neural network approach for solving the fixed point equations arising in deep equilibrium models. This consists of a tiny network that provides an initial guess for the fixed point, as well as a small network that computes coefficients inside an algorithm inspired by Anderson iteration. The approach is intuitive and empirical. Although no theory is given, the authors demonstrate the strength of their proposed solver in large scale experimental evaluations. Specifically, the new solver is fast to train, has a small parameter count, and appears to drastically shift the pareto front of the inference speed/performance curve for all DEQ models.", "summary_of_the_review": "An empirically motivated neural network replacing the role of a traditional root finder (e.g. Anderson acceleration) in DEQ models appears to improve performance and inference time. However, this root finding network introduces a new set of hyperparameters and all the baggage usually associated with deep learning (no theory of convergence, optimisation, generalisation, hyperparameter choice) (weakness 1), and it is not clear whether the choices made by the authors are universally applicable (weakness 1,2). All in all, a good paper that is likely to be adopted by the community and spark future research.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "_tOOUJ7RVWX", "reviewer": "Reviewer_CYHx", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Strength] The paper has the following strength.\n- The paper researches the very interesting, important, and difficult problem of accelerating implicit model inference. The paper gives a working solution using neural networks to assist the inference fixed point calculation with impressive results.\n- The method offers significant speedup (2x) in forward steps over the current solvers for fixed point calcuation at inference. Bringing the inference speed of the model close to explicit forward feeding models.\n- Thorough empirical analysis on the property of the solver is well conducted in the paper, showing the overhead for training the solver is minimal compared with the training of the implicit models. It is also demonstrated that the solver generalizes and scales well to large experiments in natrual language processing and computer vision.\n- Further experiments on creative usage of the solver at training time shows that the merits of the speedup when incorporated in training is larger than the overhead on training the solver. This shows that the application of the method is \"free\".\n\n[Weekness] The paper is well written but has a few glitches. I think the paper would be a good addition to ICLR if the authors can carefully address them.\n- The notations of the paper has some typos. Is G^{[k]} in (2) a matrix? If not, the L2 norm can be defined clearly. Is the dimension of \\hat{G}^{[k]} of dimension C by (m_k + 1) instead of (m_k + 1) by C?\n- The use of Hyper Anderson Iterations can be more clearly justified. After reading the Sec. 4.1, I am still curious how much worse it would be to do the naive thing by solving problem (2) (maybe approximately). Explicit explanation with some numbers could help.\n- The paper does not offer significant theroetical contribution but this is not the goal of the paper and the reveiwer believes that it should not undermine its empirical contributions.\n- The merits of the proposed solver should be highly general and go beyond deep equilibrium models (DEMs) to other implicit models including implicit archetectures for graph-structured data (Gu, 2020) and implicit Feature Pyramid Network for object detection (Wang, 2020). Discussions and potentially experiments about how the method would work on them can only improve the merits of the work.\n\nRef:\n\nGu, F., Chang, H., Zhu, W., Sojoudi, S., El Ghaoui, L. (2020). Implicit graph neural networks. Advances in Neural Information Processing Systems 33, 11984--11995. https://papers.nips.cc/paper/2020/hash/8b5c8441a8ff8e151b191c53c1842a38-Abstract.html\n\nWang, T., Zhang, X., & Sun, J. (2020). Implicit feature pyramid network for object detection. arXiv preprint arXiv:2012.13563. https://arxiv.org/abs/2012.13563", "summary_of_the_paper": "The paper presents a method called neural deep equilibrium solver to increase the efficiency in the inference stage for implicit deep models by initializing the equilibrium states using neural network. The authors start with the traditional Anderson Acceleration scheme for fixed point calculation and extend it using neural network initialization and Anderson steps to improve the inference efficiency. The authors conduct comprehensive experiments to demonstrate that the speed up in inference is significant and general with little overhead at training time. The further experiments shows that the proposed method can be incorporated in the training procedure to give faster training while introducing the speedup at inference time.", "summary_of_the_review": "This paper presents a method to significantly improve the inference of implicit models. Although the paper does not focus on theoretical contribution, it demonstrates empirical merits very well and the results are impressive. The paper is written clearly and is a solid work overall.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "ZsKzqCnRFdZ", "reviewer": "Reviewer_YWYD", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The introduction, motivation, and overview of DEQs and their contribution is very strong. They provide extensive experimentation proving the benefits of their method across many large-scale tasks, pushing the state of the art of DEQs closer to practical deployment.\n\nSome questions:\n- Does the learned extension of AA still guarantee fixed-point convergence? Is it enough to enforce that the $\\alpha_k$ weights sum to $1$?\n- Thanks to this paper, DEQ models seem to be approaching practicality during inference time. How does their training time compare to explicit models? (I understand this is orthogonal to this paper. I'm just curious.)\n- One caveat you mention is that BPTT to train HyperAnderson network could use a lot of memory. I think this could be substantially reduced with gradient checkpointing, if it ever becomes a bottleneck.", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "I believe this is a very strong paper that significantly pushes the state of the art of DEQs, and should get accepted without reservation (barring something substantial I may have missed).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce a neural network approach for solving the fixed point equations arising in deep equilibrium models. This consists of a tiny network that provides an initial guess for the fixed point, as well as a small network that computes coefficients inside an algorithm inspired by Anderson iteration.

Overall, there is consensus among the reviewers that the paper is well written and is a strong empirical study.

I recommend acceptance as a poster.

Additional remarks:

- The authors argue the DEQs / implicit deep learning models allow a decoupling between representational capacity and inference-time efficiency. Yet, in the "Regularizing Implicit Models" paragraph, they write "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.", which seems like a contradiction to me. So while in theory I agree with this decoupling, in practice, it seems not completely true.

- Section 3 should include some discussion on conditions on f_theta for the existence of a fixed point.

- Since the initialization and HyperAnderson networks are trained using unrolling, there is some memory overhead compared to vanilla DEQs, that are differentiated purely using implicit differentiation. It would be great to clarify the amount of extra memory needed by these networks. It is necessary to justify that the initialization and HyperAnderson networks are smaller than usual neural networks.
Decision: Accept (Poster)