N EURAL D EEP E QUILIBRIUM S OLVERS
Anonymous authors
Paper under double-blind review
A BSTRACT
A deep equilibrium (DEQ) model abandons traditional depth by solving for the ﬁxed point of a single nonlinear layer f θ .
This structure enables decoupling the internal structure of the layer (which controls representational capacity ) from how the ﬁxed point is actually computed (which impacts inference-time efﬁciency ), 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 ﬁxed point computation using a custom neural solver.
Speciﬁcally, 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 ﬁxed 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.
I NTRODUCTION
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
[Figure]
models deﬁne their outputs as solutions to nonlinear dynamical systems.
For example, DEQ models (which this paper will focus on) deﬁne their outputs as ﬁxed 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 ﬁnal ﬁxed point z ⋆ (Krantz & Parks, 2012; Bai et al., 2019;
Fung et al., 2021), regardless of how forward pass is computed in the ﬁrst place.
Such insulated forward and backward passes enable an equilibrium model to leverage arbitrary black-box solvers to reach the ﬁxed 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 classiﬁcation; DEQs can scale to realistic tasks, but the overhead of ﬁxed-point solvers is magniﬁed 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 beneﬁt 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 ﬁxed point by simply reducing the number of solver iterations.
This yields a speed/accuracy trade-off curve, as shown in Fig.
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 ﬁxed 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 ﬁxed-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 speciﬁc 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 classiﬁcation (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: ﬁrst, 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 signiﬁcantly advancing the agenda of deploying implicit models in practice.
R ELATED W ORK
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 deﬁned 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 inﬁntesimal 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 “ﬁxed-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 ﬁnal 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 ﬁxed 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 ﬂow solver.
Poli et al.
(2020) introduce a Neural
ODE formulation that adds a learnable residual ﬁtting 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 ﬂexible 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-deﬁned trajectories to the solution (even if one starts at the same initial point z
Our work is the ﬁrst to propose a neural ﬁxed-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-deﬁned 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 difﬁculty 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 ﬁt the nonlinear dynamics of a speciﬁc 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 ﬁxed 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 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 ﬁxed-point acceleration scheme speciﬁcally in the application of convex cone programming.
B ACKGROUND : E QUILIBRIUM M ODELS AND F IXED - POINT S OLVERS
Deep Equilibrium Models.
Given a layer (usually a shallow block; e.g., self-attention (Vaswani et al., 2017))
f θ and an input x , a DEQ model aims to solve for an “inﬁnite-level” feature representation without actually stacking the f θ layer inﬁnite times.
Instead, we can solve directly for the ﬁxed point z ⋆ of the system:
[Equation]
The ﬁxed point can be estimated by quasi-Newton (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 θ by solving a
Jacobian-based linear equation:
[Equation]
[Equation]
The most important message from Eq.
equation 1 is that the backward pass can be computed with merely the knowledge of z ⋆ , irrespective of how it is found.
More recently, Fung et al.
(2021) prove the feasibility of directly replacing J g ( z ⋆ ) with − I (i.e., Jacobian-free backward pass), which signiﬁcantly accelerates training.
Fixed-point Solvers for DEQs.
Prior works have explored a number of techniques for ﬁnding the ﬁxed points of DEQs.
For example, Bai et al.
; Lu et al.
(2021) used Broyden’s 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 et al., 2021) shifted to Anderson acceleration (AA) (Anderson, 1965), a lightweight solver that is provably equivalent to a multi-secant quasi-Newton method (Fang & Saad, 2009).
We brieﬂy introduce AA here, since our approach will use it as the starting point.
Prototype algorithm 1 illustrates the main idea of Anderson acceleration: we maintain a size- m storage of the most recent steps, and update the iteration as a normalized linear combination of these steps with weights α i (step 3).
In the canonical AA algorithm, the weights are computed in a greedy manner at each step to minimize the linear combination:
[Equation]
where G [ k ] =
[ g θ ( z [ k − m k ] ) ...
g θ
( z [ k ] )] are the past (up to m + 1 ) residuals; typically, β = 1 and m ≤ 5 .
Eq.
equation 2 can be solved by a least-squares method.
In all prior works with DEQs (Bai et al., 2019; 2020; Winston & Kolter, 2020; Revay et al., 2020; Fung et al., 2021; Lu et al., 2021), the ﬁxed point iteration starts with an initial z [ 0 ] that is either 0 or a random sample from N( 0 ,I ) .
N EURAL D EEP E QUILIBRIUM S OLVERS
While classic ﬁxed-point estimation algorithms, as presented in Section 3, already work well, they are generic and make minimal assumptions about the speciﬁc problem being solved.
For example, while multiple papers in optimization literature have acknowledged that tuning m (and m k ) as well as varying β
=
( β k ) k
= 0 ,...,K for each Anderson iteration k could accelerate AA’s convergence to the ﬁxed 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 to make ﬁxed-point solvers for DEQ models learnable and content-based, which is made possible by the unique properties of implicit models.
First, unlike generic problems, the nonlinear system for each DEQ is uniquely deﬁned by the input x (e.g., an image, etc.):
z ⋆ ( x )
=
z ⋆ = f
θ ( z ⋆ , x ) .
This opens the door to learning to make an informed initial guess, followed by content-based iterative updates in the solver.
Second, due to 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, we are able to train this neural solver in a lightweight and unsupervised manner, directly with the help of groundtruth ﬁxed-point solutions (see below).
4. G ENERAL F ORMULATION
For a given DEQ layer f θ and (possibly random) input x , we assume access to its exact ﬁxed point z
⋆ =
z ⋆ ( x )
=
f θ ( z ⋆ , x ) , which can be obtained by taking a classic solver (e.g., Broyden’s method) and running it 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 ω = { φ,ξ } (explained below) to learn the initialization and iterative solving process, and unroll the learnable solver for some K steps to yield a prediction z
[ K ] ( x ) .
To train this neural solver, we minimize an objective L( ω,K ) (discussed in Sec. 4.2) by backpropagating through this K -step temporal process (Mozer, 1989; Robinson & Fallside, 1987).
The original DEQ parameters θ are frozen, and only the hypersolver parameters ω are trained here.
We also do not need the groundtruth label y
(e.g., the class of an image) that corresponds to input x , which means these neural equilibrium solvers can also be ﬁne-tuned on the ﬂy after deployment, at inference time.
Initializer.
The initial values can have a signiﬁcant impact on the optimization process and its convergence speed.
We propose to make an input-based guess with a tiny network h φ : z [ 0 ] =
h φ ( x ) , where φ are the parameters.
Note that the goal of the initializer is not to solve the underlying problem
[Table]
⊺
[Table]
m i
[Table]
k
[Table]
m i ]
[Table]
G [0] = [ g ✓ ( z [0] ) ] | {z }
[Table]
z [ k ] , ↵ k ,  k ) z
[ K ]
Iteration k
G [ k ]
Solve ↵ k = arg min ↵ 2 R mk +1 k G [ k ] ↵ k 2 s.t. 1 > ↵ =
z [ K ] G [0] =
[ g ✓ ( z [0] ) , 0 ,..., 0] | {z } storage size m + z [0] = [Equation] R n
Set  =  k =
[Table]
Mixing (past m +1 residuals)
Parameters !
= { ,⇠ }
G [ k +1] , z [ k +1] =
AA Update ( G [ k ] , z [ k ] , ↵ k ,  k )
Input x
Input x (a) The original [Table] solver
Iteration k
G [0] =
[ g ✓ ( z [0] ) , 0 ,..., 0] | {z } storage size m +
G [ k ] ˆ G [ k ] “Compress” z [0] =
h  ( x ) 2 R n
TemporalConvolution m + (a compressed version of G [ k ] ) ↵ k = ( ↵ k 0 ,...,↵ km )  k n p
Predict [Table] k , [Table]
= s [Table] [Table] [Table]
] ) s ⇠
G [ [Table] [Table] = AA Update ( G [ k ] , z [ k ] ,
[Table] [Table] k ) z [ K ]
Iteration k
G [ k ]
Solve
[Table]
=
[Table] min [Table] 2 R mk +1 k G [ k ] ↵ k 2 s.t. 1 >
[Table]
= z [ K ] G [0] =
[ g ✓ ( z [0] ) , 0 ,..., 0] | {z } storage size m + z [0] = 0 2 R n
Set  =  k =
[Table]
Parameters
[Table] G [ [Equation]
+1] [Equation] [Equation] +1]
= AA
[Figure]
at all (e.g., to classify an image; we don’t even need the groundtruth label y ), but only to yield a quick initial estimate.
For example, in language modeling, where x ∈ R T × d is a length- T sequence, we set h φ ( x ) = ReLU ( Conv1d k = 3 ( x ))
W , where Conv1d k = 3 ∶ R T ×
d → R T × p (3) and where W ∈ R p × q , with q being the dimension of the ﬁxed point of a single token.
We set p to be very small (e.g., 100), so that h φ is tiny and fast.
Note that this 1-layer initializer by itself has very low expressivity and is usually a poor model for the original task, as we verify in Sec.
HyperAnderson Iterations.
We further parameterize the setting of β k and α ki while following the
AA prototype outlined in Alg.
In lieu of setting Eq.
2 for α to a least-squares solution over the past few residuals G , we make both α ∈ R ( m k + 1 ) and β ∈ R explicit learnable functions of G with a neural network s ξ ( G ) ∶ R ( m k + 1 )× n → ( R ( m k + 1 ) × R ) ; see Alg. 2.
A challenge here is that n (the dimension of z ⋆ ) is typically large in practice, as 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 5 on a single textual sequence of length 200).
This makes s ξ map from an extremely high-dimensional space to a low-dimensional space (e.g., m = 5 ).
To keep s ξ fast, small, and applicable to inputs of varying dimensionalities (e.g., sequence length or image size), we propose to ﬁrst compress each g θ ( z [ k ] ) to form a smaller yet still representative version ˆ G [ k ] of G [ k ] =
[ g θ ( z [ k − m k ] ) ,...,g θ ( z [ k ] )] .
For example, when each g θ ( z [ k ] ) is a image feature map residual of dimension n =
C
× H × W ,
we can
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
perform global pooling to form a C -dimensional vector Pool ( g θ
( z [ k ] )) as its compressed version: ˆ G [ k ] =
[ Pool ( g θ ( z [ k − m k ] )) ,..., Pool ( g θ ( z [ k ] ))]
∈ R C ×( m k + 1 ) , and predict α k ,
β k
= s
[Equation]
( ˆ G [ k ] ) (4)
Once we have this representative collection ˆ G [ k ] , we treat it as a mini time-series of length ( m k + 1 ) that encodes the latest estimates of the ﬁxed point.
We then apply a 2-layer temporal convolution (van den Oord et al., 2016) to learn to predict: 1) a relative weight α ki for each of these past residuals i ∈
[ m k ] ; and 2) the HyperAnderson mixing coefﬁcient β k for the current iteration.
Therefore, s ξ 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 s ξ in Appendix B, while noting that it still completely captures the AA prototype (see Alg. 1).
4. T RAINING THE N EURAL E QUILIBRIUM S OLVERS
One beneﬁt of training hypersolvers on implicit models is that they can be trained in an unsupervised manner via z ⋆ ( x ) , which a 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., 2018b; Poli et al., 2020), a DEQ model does not have a unique trajectory and thus its hypersolvers do not need trajectory ﬁtting at all.
All that we need is to drive everything to be as close to z ⋆ as possible.
As an example, a neural solver could learn to sacriﬁce progress in earlier iterations if it subsequently converges to the equilibrium faster.
Formally, given a hypersolver { h φ ,s ξ } that yields a set of states ( z [ k ] ,G [ k ] ,α k ,β k ) k
= 0 ,...,K (recall z [ 0 ] =
h φ ( x ) ), we introduce 3 objectives for its training.
Fixed-point Convergence Loss.
The ﬁrst loss aims to encourage convergence at all intermediate estimates [ z [ k ] ] k = 1 ,...,K of the HyperAnderson iterations: L conv =
∑ Kk = 1 w k ∥
z [ k ] − z ⋆ ∥ 2 , where w k is the weight for the loss from iteration k
such that ∑ Kk = 1 w k = 1 .
We set w k to be monotonically increasing with k such that later iterations apply a heavier penalty for deviation from the ﬁxed point.
Initializer Loss.
We also train the initializer by maximizing the proximity of the initial guess to the ﬁxed point: L init =
∥
h φ ( x ) − z ⋆ ∥ 2 , We separate this objective from L conv 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 α k ,β k are computed in each iteration, we empirically found it still beneﬁcial to guide the hypersolvers’ prediction of α with an auxiliary loss especially at the start of the training:
L α =
∑ Kk =
∥ G [ k ] α k ∥ 2 .
In practice, we gradually decay the weight of this loss to 0 as training progresses.
We summarize the complete training procedure of a neural solver on top of a DEQ in Fig.
4. D ISCUSSION
Complexity of hypersolver.
Note that f θ remains frozen during hypersolver training.
This means that for a given DEQ model f θ and input x , the ﬁxed point
z
⋆ ( x ) =
f θ ( z ⋆ ; x ) also remains the same – we are just trying to learn to ﬁnd it faster, with a limited K -iteration budget.
Moreover, we designed the initializer h φ 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) or distillation (Hinton et al., 2015), where a pruned/smaller (but still representationally rich) model is trained to match the output and performance of a larger model.
Speciﬁcally, in our case, as the ﬁxed point z ⋆ is determined solely by f θ and x , the hypersolver itself does not have much representational capacity, since its only goal is to produce an “educated” initial guess and learnable iterations to facilitate the optimization process.
E.g., the 1-layer Conv1d -based initializer Sec. 4.1 would be a bad language model by itself since it is tiny and only sees the pasttokens (see Sec.
5.3 for empirical evidence), yet this limited capacity and context turn out sufﬁcient to guide and substantially improve the solver.
Training hypersolver via BPTT.
While a generic Anderson solver computes α k by optimizing locally with G [ k ] , backpropagating through the HyperAnderson steps ensures that the iterative update network s ξ can receive gradient and learn from later iterations.
This is appealing because, arguably, only the output of the K th iteration matters in the end.
Indeed, we empirically verify via ablation studies in Sec.
5 that such learned α and β predictors already signiﬁcantly accelerate the convergence process even without the presence of the initializer.
Note that as DEQ models’ f θ layer is typically richly parameterized, the backpropagation-through-time (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. 5 and 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 representational capacity and forward solver choice.
In particular, our method is orthogonal to prior work that accelerates DEQ models by structural regularization of f θ (Winston & Kolter, 2020; Revay et al., 2020; Bai et al., 2021) or approximating the Jacobian of f θ in the backward pass (Fung et al., 2021).
In
Sec. 5, we show evidence that our method (which is solver-based) integrates well with regularization approaches (which are f θ -based) and yields broad improvements compared to canonical solvers (e.g.,
Broyden or Anderson methods) regardless of how f θ was trained or what structure it uses.
E XPERIMENTS
In this section, we verify the various beneﬁts of exploiting neural solvers in implicit models.
Specifically, as our goal is to show the superiority of the learnable solvers over generic solvers on both performance and efﬁciency 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 classiﬁcation (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 - 2 × speedup at inference and no 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 f θ to bring these implicit models to a 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 f θ 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), which we discuss further in Appendix A. We also clarify that the use of hypersolver does implicitly assume local stability around z ⋆ for convergence – which we ﬁnd almost always holds empirically, and can be regularized for (Bai et al., 2021).
Our code is provided in the supplement.
5. L ARGE - SCALE E XPERIMENTS ON V ISION AND L ANGUAGE T ASKS
To evaluate the the neural deep equilibrium solvers, we apply them on three largest-scale and highestdimensional tasks the implicit models have ever been applied on, across the vision and language modalities.
In contrast to prior works (Chen et al., 2018b; Winston & Kolter, 2020; Bai et al., 2021) that measure the number of function evaluations (NFEs), we directly measure wall-clock inference speed under the exact same experimental settings (e.g., input scale).
We elaborate on the detailed experimental settings and the implications of the results below.
[Figure] [Figure]
WikiText-103 Language Modeling.
In this experiment, f θ is a Transformer layer (Vaswani et al., 2017; Dai et al., 2019; Bai et al., 2019) and the ﬁxed 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 a validation sequence length of 150).
Speciﬁcally, compared with the original DEQ-Transformer (Bai et al., 2019)
( Y curve), which uses generic Anderson acceleration (Anderson, 1965) or Broyden’s method (Broyden, 1965) (both have similar pareto curves; see App.
C), this same DEQ model solved with our neural approach (dubbed HyperDEQ; see curve) achieves signiﬁcantly better efﬁciency.
Moreover, our 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, we additionally train a DEQ-Transformer model with Jacobian regularization (Bai et al., 2021)
( curve), and apply the neural solver on this regularized DEQ ( ⋆ curve).
This movement of the speed/perplexity curves validates the DEQ property at the core of this paper: the decoupling of the representational capacity (i.e., f θ ) 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 classiﬁcation.
We additionally evaluate HyperDEQ on ImageNet classiﬁcation ( 224 × 224 images), customizing a neural solver on top of a 4-resolutional multiscale DEQ models (Bai et al., 2020).
We train the HyperDEQ with 12 HyperAnderson iterations, and the speed/accuracy curves are shown in Figure 4b ( and ⋆ curves).
Note that while Jacobian regularization ( curve) eventually hurts the performance of a multiscale DEQ (cf.
Y curve)
due to the strong constraint it imposes, the DEQ model with neural solver achieves faster inference without sacriﬁng any accuracy (since f θ , and thus z ⋆ , are identical); e.g., we reach 75.0% accuracy while being almost 2 × faster.
Cityscapes semantic segmentation.
We also show that our neural solver approach works well in domains where existing regularization-based methods (see Sec. 2) fail.
Speciﬁcally, we apply the neural equilibrium solver on Cityscapes semantic segmentation, where the task objective is to label every pixel on a high-resolution (typically 2048 × 1024 ) image with the class of the object that the pixel belongs to.
As in the ImageNet and WikiText-103 tasks, we found that there is a consistent gain in using the neural solver over the generic alternative, accelerating ﬁxed-point convergence by more than a 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 f θ .
(Speciﬁcally, when f θ is convolution-based and the image is very large, Jacobian regularization that encourages contractivity is at odds with the gradual broadening of the receptive ﬁeld.)
Our neural solver is orthogonal to the structure of f θ (which is frozen), and we only improve how the solver functions.
5. T RAINING E FFICIENCY OF THE N EURAL S OLVER
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 is 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 ﬁxed 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 { h φ ,s ξ } and the DEQ model f θ in an alternating manner, and elaborate more in App.
D. We empirically observe this leads to a 16-20% DEQ training speedup.
5. A BLATIVE S TUDIES AND L IMITATIONS
Finally, we perform a series of ablation studies to understand the beneﬁts of multiple components
[Table]
[Figure]
within our design of the neural equilibrium solvers.
We use the language modeling task on WikiText-103 for this purpose (where f θ is a 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, α k , and β k predictions) performs best.
Making the Anderson iterations learnable generally improves convergence.
Moreover, although simply adding an initializer to a generic solver ( (cid:58) 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 h φ from
HyperDEQ and verify that this tiny module is by itself still a poor language model (see Table 1 and
Sec. 4.3), but is valuable to our HyperAnderson iterations.
More ablation studies (e.g., how α is predicted) are reported in Appendix C.
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 we run for.
However, we don’t ﬁnd this to be problematic in practice, as we observed the training of these hypersolvers to be very insensitive to 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 (e.g., f θ is a 3 × 3 convolution and the input is a large image).
D ISCUSSION
We introduce a neural ﬁxed-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., f θ ) from the forwards computation.
We directly learn a model-speciﬁc 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 modiﬁcations 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 signiﬁcant 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.
A E XPERIMENTAL D ETAILS
We describe here (and summarize in Table 2) the datasets, pretrained models and the training settings (including hyperparameters) for the experimental results reported in Sec. 5.
We note that overall, since our approach is orthogonal to how these DEQ models were trained in the ﬁrst place, we directly downloaded the pretrained DEQ models from the publicly released DEQ repo (Bai et al., 2019) (for language modeling) and the MDEQ repo (Bai et al., 2020) (for image classiﬁcation and segmentation), which our code base is built upon.
All of our experiments were conducted on NVIDIA RTX
Ti
GPUs.
The pareto efﬁciencies were benchmarked on 1 GPU, averaged over 6 independently trained hypersolvers (i.e., 6 random seeds) for each task.
A. W IKI T EXT -103 W ORD - LEVEL L ANGUAGE M ODELING
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 al., 2018), BERT (Devlin et al., 2019), XLNet (Yang et al., 2019),and GPT3 (Brown et al., 2020).
Formally, given an autoregressive network F (like a Transformer (Vaswani et al., 2017) or LSTMs (Hochreiter & Schmidhuber, 1997)) and input sequence x 1 ∶ t , the prediction ˆ y 1 ∶ t = F ( x 1 ∶ t ) ∈ R t × N should be ideally as close as possible to the groundtruth “next words” x 2 ∶ t + 1 , with y t ∈ R N (i.e., like an N -way classiﬁcation, where N is the vocabulary size).
Our experiments follow the exact same settings as the DEQ-Transformer network in Bai et al.
(2019), where a deep equilibrium (DEQ) model is as a sequence model and f θ a single Transformer block (i.e., multi-head self-attention (Vaswani et al., 2017; Dai et al., 2019)) ; i.e., ˆ y 1 ∶ t = h ( z ⋆ 1 ∶ t ) where z ⋆ 1 ∶ t = f θ ( z ⋆ 1 ∶ t , x 1 ∶ t ) , and h a linear layer. [Equation]
We use one of the most commonly-used and the largest-scale textual corpus, Wikitext-103 (Merity et al., 2017), to evaluate our method.
This dataset can be downloaded at this link.
Speciﬁcally, the Wikitext-103 corpus contains over 103M words in its training split, and 218K/246K words for validation/test.
Moreover, 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, we use Adam optimizer (Kingma & Ba, 2015) 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, on top of a pretrained DEQ with word embedding dimension 700.
We set the HyperAnderson iteration limit
K to 10 (i.e., the HyperAnderson module is unrolled for 10 steps at training), while evaluating the hypersolver at validation time for up to K = 16 steps for the pareto efﬁciency curve benchmarking (see Fig. 4a).
A. I MAGE N ET C LASSIFICATION
In both the Imagenet classiﬁcation and the Cityscapes segmentation tasks (see below), we use a pretrained multiscale-DEQ (MDEQ-small, with 4 resolutions) model (Bai et al., 2020) for training our neural equilibrium solvers.
In this case, feature maps maintained at multiple (e.g., n ) resolutions are driven to their equilibria simultaneously; i.e.,
[Equation]
To verify our approach, we 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 (e.g., random cropping or horizontal ﬂipping, 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. C ITYSCAPES S EGMENTATION
The other, even higher-resolution, vision dataset we apply the hypersolver on is the Cityscapes dataset (Cordts et al., 2016).
This is a large-scale dataset that contains a diverse set of high-quality 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.).
Speciﬁcally, these 5,000 frames are divided into 2,975 training,validation and 1,525 test images.
The multiscale-DEQ (MDEQ) we use follow the exact formulation as in the ImageNet case above, with 4 resolutions driven to the ﬁxed-point simultaneously.
The Cityscapes pixel-level semantic segmentation task is also a standard widely-used benchmark (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, as in ImageNet, we do not perform data augmentations like horizontal ﬂipping (but perform random cropping to 1024 × 512 ) on the training set, where we set the HyperAnderson step limit to K = 12 and the batch size to 6.
The neural solver is then evaluated in the standard full-resolution setting, on Cityscapes val 2048 ×images.
A. G ROUNDTRUTH S OLVERS AND L OSS W EIGHTS
In all of the experiments, we use the canonical Broyden’s method (Broyden, 1965) with 60 standard quasi-Newton steps as the “groundtruth” solver that yields the exact ﬁxed 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., a canonical Anderson acceleration algorithm with 60 steps would also sufﬁce, since it’s the same groundtruth ﬁxed point).
The major hyperparameters introduced by the neural solver are the relative loss weights.
For the ﬁxed-point convergence loss L conv =
∑ Kk = 1 w k ∥
z [ k ] − z ⋆ ∥ 2 (see Sec. 4), we apply weights w k > 0 to all intermediate HyperAnderson steps, with w k monotonically increasing (such as later HyperAnderson steps are applied a heavier penalty for deviation from the ﬁxed point) and sum to 1; e.g., we can simply do
[Equation]
Note that other similar formulations are possible, such as w k = k∑
Ki
[Equation] 1 i 2 , and we generally do not ﬁnd them empirically make a huge difference.
We further illustrate this in the ablation studies presented in Appendix C.
= Conv2d
(kernel size 3) + ReLU . . . . . .
x
Initializer h
Initial guess z [0]
= ... ... ... . . .
P a s t m + 1 r e
s i du a l s g ✓ ( z [ k +1] ) g ✓ ( z [ k  m ] ) g ✓ ( z [ k  m +1] )
GlobalPool2d ( · ) Select lowest resolution feature, then global pool G [ k ] ˆ G [ k ] (A compressed version of G [ k ] , independent of input scale H,W )
HyperAnderson Iteration (compression step) . . .
[Table]
h
[Table]
[Table]
P a s t m +r e s i du a l s g ✓ ( z [ k +1] ) g ✓ ( z [ k  m ] ) g ✓ ( z [ k  m +1] )
G [ k ] ˆ G [ k ]
HyperAnderson Iteration (compression step)
[Table]
[Table] [Table]
T
[Table]
d . . .
[Table]
d
[Table]
p ...
[Table] [Table]
) m + m +
[Table]
Select
[Table]
Linear(
[Table] [Table]
[Table] [Table]
P a s t m +r e s i du a l s g ✓ ( z [ k +1] ) g ✓ ( z [ k  m ] ) g ✓ ( z [ k  m +1] )
GlobalPool2d ( · )
Select lowest resolution feature, then global pool
G [ k ] ˆ G [ k ] (A compressed version of G [ k ] , independent of input scale H,W )
HyperAnderson Iteration (compression step) . . .
Initializer h
Initial guess z
[0] = . . .
P a s t m +r e s i du a l s g ✓ ( z [ k +1] ) g ✓ ( z [ k  m ] ) g ✓ ( z [ k  m +1] )
G [ k ] ˆ G [ k ]
HyperAnderson Iteration (compression step) "
The son of farmer… in 1932 .”
[Figure]
(embeddings) x 2 R T ⇥ d . . . Linear + ReLU d p
[Figure]
p ... GlobalPool1d + Linear
[Figure]
all
[Figure]
across the temporal dimension (A compressed version of G [ k ] , independent of input scale T ) m + m +
[Figure]
[Figure]
[Figure]
We applied a simple grid search on the relative loss weights λ 1 ,λ 2 and λ 3 on the three losses, ﬁnding it overall beneﬁcial to adjust these weights so that the 3 losses at the start of the training be at roughly the same magnitude/scale.
Moreover, we empirically ﬁnd it useful to set λ 3 larger initially to encourage Anderson-like iterations (cf.
Eq.
equation 2), but later gradually decay this weight to almost 0, so that the hypersolver can more ﬂexibly exploit the α k in each HyperAnderson iterations.
In all experiments, we decay λ 3 to 5e-8 on a linear schedule over 2,000 steps.
B A DDITIONAL D ISCUSSIONS ON THE N EURAL E QULIBRIUM S OLVER D ESIGN
We brieﬂy 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 as minimial number of parameters as possible, which implies projections to low-dimensions; 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-learn (L2L; see Sec. 2) way of parameterizing optmizers by RNNs: whereas the parameter space has a ﬁxed dimensionality, hidden unit space does not.
HyperDEQ for sequences.
In the sequence modeling task, as mentioned in Sec. 4, our initializer h φ is simply an 1D causal convolution followed by a non-linear projection (see Fig.
6a):
[Equation]
where Conv1d k = 3 ∶ R T × d → R T × p (1D causal convolution with kernel size 3) [Figure] [Figure] [Figure]
In particular, we set p ≪
d (e.g., in the WikiText-103 experiment, d = 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
[Figure]
g., high-dimensional language modeling), but only to give a “reasonable” initial guess to the ﬁxed point solver. In other words, the initial guess itself may be bad in a task perspective (as we veriﬁed in Table 1) but good in an optimization perspective (see Fig. 5).
Moreover, we compress the past and current residuals G [ k ] =
[ g θ ( z [ k − m ] ) ,...,g θ ( z [ k + 1 ] )] ∈
R ( m + 1 )× T × q by simply taking the last and most contextualized token of each sequence; i.e., ( g θ ( z [ i ] ))
T ∈ R q .
This is reasonable because the design of a sequence model like Transformer layers (which parameterizes our f θ 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 classiﬁcation tasks, or in vision transformers (Dosovitskiy et al., 2020)).
We therefore propose to compress the past residuals of the ﬁxed 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. 4.
HyperDEQ (multiscale) for vision.
The initializer for multiscale DEQ models is similarly designed, with a series of small-context 1-layer convolutions (e.g., kernel size 3) to generate the initial guess for each resolution stream (see Fig. 6b).
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 lowestresolution feature map
z
[ i ] n of the past ﬁxed point estimate z
[ i ] =
[ z [ i ] 1 ,...,z [ i ] n ] (in practice, we use n = 4 ).
As shown in Fig.
6b, we directly apply global pooling to pass into the later HyperAnderson steps (by treating it as a mini time-series, like in the sequence case).
C A DDITIONAL E FFICIENCY AND A BLATIVE S TUDIES
We further study some ablative settings not covered in Sec.
5.3 here.
We speciﬁcally focus on two aspects: 1) how α is predicted; and 2) the effects of different loss components.
In the original AA prototype, the values of α is determined by solving a least-squared solution over the past m steps, as described by the constrained optimization problem equation 2.
In particular, these α k = ( α k 0 ,...,α k m ) values can be any real value, but has to sum to 1.
In our hypersolver, we propose to “normalize” the predicted ˆ α k by a shift: α k
= ˆ α k
+ ( 1 − 1 ⊺ ˆ α k ) m k
+ 1 ⋅ 1 .
We here compare two alternatives: no normalization at all (i.e., have α learned freely; see in Fig. 7a), and softmax-based normalization (i.e., requiring α k to be all-positive; see in Fig. 7a).
Speciﬁcally, we ﬁnd that the exact choice of α normalization does not affect the overall substantial pareto curve improvement (e.g., compare them with the Y curve baseline), but ensuring that the α k values can be negative while
[Figure]
l
[Figure]
t i v e r
[Figure]
s i
[Figure]
u a l | f (
[Figure]
z | | z |
[Figure]
The speedup of HyperDEQ is a result of this faster convergence and the more lightweight solver step.
[Figure]
still normalized (i.e., sum to 1) overall beneﬁt 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 ﬁxed-point convergence loss L conv , we compare two alternatives of the relative weights applied on the intermediate
HyperAnderson steps: only apply w k = 1 at the ﬁnal output K (see (cid:58) curve), or set w k to be uniform for all k (see (cid:58) curve).
We empirically ﬁnd that the hypersolver trained in both scenarios perform well, but the monotonically incresing w k we use (i.e., putting larger emphasis on later iterations) perform best.
Moreover, removing either the initializer loss L init (note that we still keep the initializer h φ itself, just don’t supervise it; see the ) or the alpha loss (the curve)
L α impacts the performance.
Interestingly, removing the L init 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 signiﬁcantly outperform the speed/accuracy efﬁciency than the generic solvers, which suggests the overall beneﬁt of using customized learnable solvers for implicit models.
Finally, we provide additional empirical evidence on the convergence and generalizability of the neural solvers in Fig.
8a, where we compare the convergence of a pre-trained, regularized Transformerbased 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 8a, we ﬁrst note that a hypersolver trained with K unrolling steps is able to generalize (i.e., converge well) beyond K , although both solvers eventually plateau at some of ε residual value.
Second, with the parameterized ﬁxed-point solving, we see that the neural solver consistently improves over the canonical solver’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 2 × speedup we observe across the board for DEQ model inference efﬁciency in Sec. 5.
D U SE H YPERSOLVER FOR DEQ T RAINING
It 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 f θ also changes, which means the ﬁxed-point manifold of this implicit model constantly gets updated.
Therefore, a neural deep equilibrium solver { h φ ,s ξ } trained on a parameter state f θ 0 may gradually deviate from accurate and stable convergence as f θ also deviates from f θ 0 as a result of more and more parameter updates.
Second, precisely computing the groundtruth ﬁxed points z ⋆ with a classical solver at every training step would add heavy burden to the training overhead.
Third, the usage of hypersolver to solve forward pass ﬁxed point does not affect the backward pass implicit differentiation (notably, decoupled forward and backward trajectories is also a major feature of implicit models), implying the training efﬁciency improvement will be much smaller than that at the inference time (about 2 × , as mentioned in Sec. 5).
Despite these challenges, exploiting the fact that model parameters θ drift only gradually with training iterations (i.e., at training step t and t + 1 , we assume ∥ θ t − θ t + 1 ∥ is small), we propose to train the neural solver { h φ ,s ξ } (which are tiny, lightweight, and have little representational capacity) and the implicit layer f θ (which are huge, representationally rich, and requires many training steps) in an alternating way.
Speciﬁcally, we adopt the following procedure:
Warmup and train the DEQ model for some small number of steps.
Take a snapshot of the current implicit layer f θ and train from scratch a hypersolver { h φ ,s ξ } customized on top of it.
Use the current hypersolver { h φ ,s ξ } to solve for the ﬁxed point and train the DEQ model for the next M steps. 4.
Take a snapshot of the current implicit layer f θ , and ﬁne-tune the existing hypersolver { h φ ,s ξ } on top of it by minimizing L( ω,K ) for some T hypersolver training steps.
Repeat step 3 & 4 until we reach max training steps for the DEQ model.
Empirically, we ﬁnd that the proposed neural solver generalizes surprisingly well for values of M that are not too large (i.e., even though f θ gets updated, the hypersolver can still reliably converge to the correct ﬁxed points during the next M DEQ training steps).
Moreover, as previously mentioned, training neural solver is extremely lightweight, and even more so for ﬁne-tuning it, suggesting that T is very small in pratice.
For example, on WikiText-103 language modeling with a transformer-based f θ (which has 300K DEQ training steps), we set M = 5000 and T = 200 ; in other words, we ﬁne-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 ﬁne-tuning hypersolver is counterbalanced by the huge gain in using hypersolver over a canonical generic root solver.
Fig.
8b shows the training efﬁciency improvement as a result of adopting hypersolver in DEQ training.
With the hypersolver ﬁne-tuned intermittently (green curve), we are able to reduce the total training time over that of the classical-solver-based 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-based methods, like Jacobian regularizations (cf.
purple and pink curves).
E A DDITIONAL E XPERIMENT ON T INY E QUILIBRIUM M ODELS
While Fig. 4 provides strong evidence on how the hypersolver could improve the speed/accuracy [Figure] pareto efﬁciency of these large-scale DEQ models (e.g., DEQ-Transformer contains almost 100M parameters), we also perform additional experiments to validate its efﬁcacy on much smaller networks.
Speciﬁcally, we train a hypersolver on a tiny MDEQ model that is pretrained on CIFARclassiﬁcation task (the same setting as in Bai et al.
(2020), but with Jacobian regularization) which contains only 170K parameters.
The hypersolver contains about 8K parameters (including the initializer h φ ), and is trained for only 1000 training steps.
The results are shown in Fig.
The improvement in pareto frontier again suggests that the neural solvers provide a lightweight and yet effective solution to improving these implicit models.
F C ODE
We also provide code for this submission; see the attached ﬁles in the supplementary material.
Detailed instructions are provided in README.md .
[Figure]