On Training Implicit Models
Year: 2022

Abstract:
This paper focuses on training implicit models of infinite layers. Specifically, previous works employ implicit differentiation and solve the exact gradient for the backward propagation. However, is it necessary to compute such an exact but expensive gradient for training? In this work, we propose a novel gradient estimate for implicit models, named phantom gradient, that 1) forgoes the costly computation of the exact gradient; and 2) provides an update direction empirically preferable to the implicit model training. We theoretically analyze the condition under which an ascent direction of the loss landscape could be found, and provide two specific instantiations of the phantom gradient based on the damped unrolling and Neumann series. Experiments on large-scale tasks demonstrate that these lightweight phantom gradients significantly accelerate the backward passes in training implicit models by roughly 1.7×, and even boost the performance over approaches based on the exact gradient on ImageNet.

heading: 1 Introduction
text: Conventional neural networks are typically constructed by explicitly stacking multiple linear and non-linear operators in a feed-forward manner. Recently, the implicitly-defined models [1, 2, 3, 4, 5] have attracted increasing attentions and are able to match the state-of-the-art results by explicit models on several vision [3, 6], language [2] and graph [4] tasks. These works treat the evolution of the intermediate hidden states as a certain form of dynamics, such as fixed-point equations [2, 3] or ordinary differential equations (ODEs) [1, 7], which represents infinite latent states. The forward passes of implicit models are therefore formulated as solving the underlying dynamics, by either black-box ODE solvers [1, 7] or root-finding algorithms [2, 3]. As for the backward passes, however, directly differentiating through the forward pass trajectories could induce a heavy memory overhead [8, 9]. To this end, researchers have developed memory-efficient backpropagation via implicit differentiation, such as solving a Jacobian-based linear fixed-point equation for the backward pass of deep equilibrium models (DEQs) [2], which eventually makes the backpropagation trajectories independent of the forward passes. This technique allows one to train implicit models with essentially constant memory consumption, as we only need to store the final output and the layer itself without saving any intermediate states. However, in order to estimate the exact gradient by implicit differentiation, implicit models have to rely on expensive black-box solvers for backward passes, e.g., ODE solvers or root-solving algorithms. These black-box solver usually makes the gradient computation very costly in practice, even taking weeks to train state-of-the-art implicit models on ImageNet [10] with 8 GPUs.
This work investigates fast approximate gradients for training implicit models. We found that a firstorder oracle that produces good gradient estimates is enough to efficiently and effectively train implicit
*Equal contribution †Corresponding author, zlin@pku.edu.cn
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
models, circumventing laboriously computing the exact gradient as in prior arts [2, 3, 4, 11, 12]. We develop a framework in which a balanced trade-off is made between the precision and conditioning of the gradient estimate. Specifically, we provide the general condition under which the phantom gradient can provide an ascent direction of the loss landscape. We further propose two instantiations of phantom gradients in the context of DEQ models, which are based on the the damped fixed-point unrolling and the Neumann series, respectively. Importantly, we show that our proposed instantiations satisfy the theoretical condition, and that the stochastic gradient descent (SGD) algorithm based on the phantom gradient enjoys a sound convergence property as long as the relevant hyperparameters, e.g., the damping factor, are wisely selected. Note that our method only affects, and thus accelerates, the backward formulation of the implicit models, leaving the forward pass formulation (i.e., the root-solving process) and the inference behavior unchanged so that our method is applicable to a wide range of implicit models, forward solvers, and inference strategies. We conduct an extensive set of synthetic, ablation, and large-scale experiments to both analyze the theoretical properties of the phantom gradient and validate its speedup and performances on various tasks, such as ImageNet [10] classification and Wikitext-103 [13] language modeling.
Overall, our results suggest that: 1) the phantom gradient estimates an ascent direction; 2) it is applicable to large-scale tasks and is capable of achieving a strong performance which is comparable with or even better than that of the exact gradient; and 3) it significantly shortens the total training time needed for implicit models roughly by a factor of 1.4 ∼ 1.7×, and even accelerates the backward passes by astonishingly 12× on ImageNet. We believe that our results provide strong evidence for effectively training implicit models with the lightweight phantom gradient.

heading: 2 Method

heading: 2.1 Inspection of Implicit Differentiation
text: In this work, we primarily focus on the formulation of implicit models based on root-solving, represented by the DEQ models [2]. The table of notations is arranged in Appendix A. Specifically, given an equilibrium module F , the output of the implicit model is characterized by the solution h∗ to the following fixed-point equation: h∗ = F(h∗, z), (1) where z ∈ Rdu+dθ is the union of the module’s input u ∈ Rdu and parameters θ ∈ Rdθ , i.e., z> = [u>,θ>]. Here, u is usually a projection of the original data point x ∈ Rdx , e.g., u =M(x). In this section, we assume F is a contraction mapping w.r.t. h so that its Lipschitz constant Lh w.r.t. h is less than one, i.e., Lh < 1, a setting that has been analyzed in recent works [14, 15]1.
To differentiate through the fixed point by Eq. (1), we need to calculate the gradient of h∗ w.r.t. the input z. By Implicit Function Theorem (IFT), we have
∂h∗ ∂z = ∂F ∂z ∣∣∣∣ h∗ ( I − ∂F ∂h ∣∣∣∣ h∗ )−1 . (2)
Here, (∂a/∂b)ij = ∂aj/∂bi. The equilibrium point h ∗ of Eq. (1) is then passed to a post-processing function G to obtain a prediction ŷ = G(h∗). In the generic learning scenario, the training objective is the following expected loss:
R(θ) = E(x,y)∼P [L(ŷ(x;θ),y)] , (3) where y is the groundtruth corresponding to the training example x, and P is the data distribution. Here, we omit the parameters of G, because given the output h∗ of the implicit module F , training the post-processing part G is the same as training explicit neural networks. The most crucial component is the gradient of the loss function L w.r.t. the input vector z> = [u>,θ>], which is used to train both the implicit module F and the input projection moduleM. Using Eq. (2) with the condition h = h∗, we have
∂L ∂u = ∂F ∂u
( I − ∂F
∂h )−1 ∂L ∂h , ∂L ∂θ = ∂F ∂θ ( I − ∂F ∂h )−1 ∂L ∂h . (4)
The gradients in Eq. (4) are in the same form w.r.t. u and θ. Without loss of generality, we only discuss the gradient w.r.t. θ in the following sections.
1Note that the contraction condition could be largely relaxed to the one that the spectral radius of ∂F/∂h∗ on the given data is less than 1, i.e., ρ(∂F/∂h∗) < 1, as indicated by the well-posedness condition in [5].

heading: 2.2 Motivation
text: The most intriguing part lies in the Jacobian-inverse term, i.e., (I − ∂F/∂h)−1. Computing the inverse term by brute force is intractable due to the O(n3) complexity. Previous implicit models [2] approach this by solving a linear system involving a Jacobian-vector product iteratively via a gradient solver, introducing over 30 Broyden [16] iterations in the backward pass. However, the scale of the Jacobian matrix can exceed 106 × 106 in the real scenarios, leading to a prohibitive cost in computing the exact gradient. For example, training a small-scale state-of-the-art implicit model on ImageNet can consume weeks using 8 GPUs while training explicit models usually takes days, demonstrating that pursuing the exact gradient severely slows down the training process of implicit models compared with explicit models.
Secondly, because of the inversion operation, we cast doubt on the conditioning of the gradient and the stability of training process from the numerical aspect. The Jacobian-inverse can be numerically unstable when encountering the ill-conditioning issue. The conditioning problem might further undermine the training stability, as studied in the recent work [17].
Plus, the inexact gradient [9, 18, 19, 20, 21, 22] is widely applied in the previous learning protocol, like linear propagation [23] and synthetic gradient [24]. Here, the Jacobian-inverse is used to calculate the exact gradient which is not always optimal for model training. Moreover, previous research has used a moderate gradient noise as a regularization approach [25], which has been shown to play a central role in escaping poor local minima and improving generalization ability [26, 27, 28].
The concerns and observations motivate us to rethink the possibility of replacing the Jacobian-inverse term in the standard implicit differentiation with a cheaper and more stable counterpart. We believe that an exact gradient estimate is not always required, especially for a black-box layer like those in the implicit models. Hence this work designs an inexact, theoretically sound, and practically efficient gradient for training implicit models under various settings. We name the proposed gradient estimate as the phantom gradient.
Suppose the Jacobian ∂h∗/∂θ is replaced with a matrixA, and the corresponding phantom gradient is defined as
∂̂L ∂θ := A ∂L ∂h . (5)
Next, we give the general condition on A so that the phantom gradient can be guaranteed valid for optimization (Sec. 2.3), and provide two concrete instantiations of A based on either damped fixed-point unrolling or the Neumann series (Sec. 2.4). The proofs of all our theoretical results are presented in Appendix C.

heading: 2.3 General Condition on the Phantom Gradient
text: Previous research on theoretical properties for the inexact gradient include several aspects, such as the gradient direction [22], the unbiasedness of the estimator [29], and the convergence theory of the stochastic algorithm [19, 30]. The following theorem formulates a sufficient condition that the phantom gradient gives an ascent direction of the loss landscape. Theorem 1. Suppose the exact gradient and the phantom gradient are given by Eq. (4) and (5), respectively. Let σmax and σmin be the maximal and minimal singular value of ∂F/∂θ. If∥∥∥∥A(I − ∂F∂h ) − ∂F ∂θ
∥∥∥∥ < σ2minσmax , (6) then the phantom gradient provides an ascent direction of the function L, i.e.,〈
∂̂L ∂θ , ∂L ∂θ
〉 > 0. (7)
Remark 1. Suppose only the (I − ∂F/∂h)−1 term is replaced with a matrix D, namely, A = (∂F/∂θ)D. Then, the condition in (6) can be reduced into∥∥∥∥D(I − ∂F∂h ) − I ∥∥∥∥ < 1κ2 , (8) where κ is the condition number of ∂F/∂θ. (See Appendix C.1 for the derivation.)

heading: 2.4 Instantiations of the Phantom Gradient
text: In this section, we present two practical instantiations of the phantom gradient. We also verify that the general condition in Theorem 1 can be satisfied if the hyperparameters in our instantiations are wisely selected.
Suppose we hope to differentiate through implicit dynamics, e.g., either a root-solving process or an optimization problem. In the context of hyperparameter optimization (HO), previous solutions include differentiating through the unrolled steps of the dynamics [19] or employing the Neumann series of the Jacobian-inverse term [9]. In our case, if we solve the root of Eq. (1) via the fixed-point iteration: ht+1 = F(ht, z), t = 0, 1, · · · , T − 1, (9) then by differentiating through the unrolled steps of Eq. (9), we have
∂hT ∂θ = T−1∑ t=0 ∂F ∂θ ∣∣∣∣ ht T−1∏ s=t+1 ∂F ∂h ∣∣∣∣ hs . (10)
Besides, the Neumann series of the Jacobian-inverse (I − ∂F/∂h)−1 is
I + ∂F ∂h + ( ∂F ∂h )2 + ( ∂F ∂h )3 + · · · . (11)
Notably, computing the Jacobian ∂h∗/∂θ using the Neumann series in (11) is equivalent to differentiating through the unrolled steps of Eq. (9) at the exact equilibrium point h∗ and taking the limit of infinite steps [9].
Without altering the root of Eq. (1), we consider a damped variant of the fixed-point iteration:
ht+1 = Fλ(ht, z) = λF(ht, z) + (1− λ)ht, t = 0, 1, · · · , T − 1. (12)
Differentiating through the unrolled steps of Eq. (12), Eq. (10) is adapted as
∂hT ∂θ = λ T−1∑ t=0 ∂F ∂θ ∣∣∣∣ ht T−1∏ s=t+1 ( λ ∂F ∂h ∣∣∣∣ hs + (1− λ) I ) . (13)
The Neumann series of (I − ∂F/∂h)−1 is correspondingly adapted as
λ ( I +B +B2 +B3 + · · · ) , where B = λ
∂F ∂h + (1− λ)I. (14)
The next theorem shows that under mild conditions, the Jacobian from the damped unrolling in Eq. (13) converges to the exact Jacobian and the Neumann series in (14) converges to the Jacobianinverse (I − ∂F/∂h)−1 as well. Theorem 2. Suppose the Jacobian ∂F/∂h is a contraction mapping. Then,
(i) the Neumann series in (14) converges to the Jacobian-inverse (I − ∂F/∂h)−1; and
(ii) if the function F is continuously differentiable w.r.t. both h and θ, the sequence in Eq. (13) converges to the exact Jacobian ∂h∗/∂θ as T →∞, i.e.,
lim T→∞ ∂hT ∂θ = ∂F ∂θ ∣∣∣∣ h∗ ( I − ∂F ∂h ∣∣∣∣ h∗ )−1 . (15)
However, as discussed in Sec. 2.2, it is unnecessary to compute the exact gradient with infinite terms. In the following context, we introduce two instantiations of the phantom gradient based on the finite-term truncation of Eq. (13) or (14).
Unrolling-based Phantom Gradient (UPG). In the unrolling form, the matrixA is defined as
Aunrk,λ = λ k−1∑ t=0 ∂F ∂θ ∣∣∣∣ ht k−1∏ s=t+1 ( λ ∂F ∂h ∣∣∣∣ hs + (1− λ) I ) . (16)
Neumann-series-based Phantom Gradient (NPG). In the Neumann form, the matrixA is defined as
Aneuk,λ = λ ∂F ∂θ ∣∣∣∣ h∗ ( I +B +B2 + · · ·+Bk−1 ) , where B = λ ∂F ∂h ∣∣∣∣ h∗ + (1− λ)I. (17)
Note that both the initial point of the fixed-point iteration (i.e., h0 in Eq. (16)) and the point at which the Neumann series is evaluated (i.e., h∗ in (17)) are the solution of the root-finding solver. (See Appendix B for implementation of the phantom gradient.)
According to Theorem 2, the matrix A defined by either Eq. (16) or (17) converges to the exact Jacobian ∂h∗/∂θ as k →∞ for any λ ∈ (0, 1]. Therefore, by Theorem 2, the condition in (6) can be satisfied if a sufficiently large step k is selected, since∥∥∥∥A(I − ∂F∂h ) − ∂F ∂θ ∥∥∥∥ ≤ (1 + Lh) ∥∥∥∥∥A− ∂F∂θ ( I − ∂F ∂h
)−1∥∥∥∥∥ . (18) Next, we characterize the impact of the two hyperparameters, i.e., k and λ, on the precision and conditioning ofA. Take the NPG (Eq. (17)) as an example.
(i) On the precision of the phantom gradient, • a large k makes the gradient estimate more accurate, as higher-order terms of the
Neumann series are included, while • a small λ slows down the convergence of the Neumann series because the norm ‖B‖
increases as λ decreases. (ii) On the conditioning of the phantom gradient,
• a large k impairs the conditioning of A since the condition number of Bk grows exponentially as k increases, while
• a small λ helps maintain a small condition number ofA because the singular values of ∂F/∂h are “smoothed” by the identity matrix.
In a word, a large k is preferable for a more accurate A, while a small λ contributes to the wellconditioning of A. Practically, these hyperparameters should be selected in consideration of a balanced trade-off between the precision and conditioning ofA. See Sec. 3 for experimental results.

heading: 2.5 Convergence Theory
text: In this section, we provide the convergence guarantee of the SGD algorithm using the phantom gradient. We prove that under mild conditions, if the approximation error of the phantom gradient is sufficiently small, the SGD algorithm converges to an -approximate stationary point in the expectation sense. We will discuss the feasibility of our assumptions in Appendix C.3. Theorem 3. Suppose the loss functionR in Eq. (3) is `-smooth, lower-bounded, and has bounded gradient almost surely in the training process. Besides, assume the gradient in Eq. (4) is an unbiased estimator of∇R(θ) with a bounded covariance. If the phantom gradient in Eq. (5) is an -approximation to the gradient in Eq. (4), i.e.,∥∥∥∥∥ ∂̂L∂θ − ∂L∂θ
∥∥∥∥∥ ≤ , almost surely, (19) then using Eq. (5) as a stochastic first-order oracle with a step size of ηn = O(1/ √ n) to update θ with gradient descent, it follows after N iterations that
E [∑N n=1 ηn‖∇R(θn)‖2∑N
n=1 ηn
] ≤ O ( +
logN√ N
) . (20)
Remark 2. Consider the condition in (19):∥∥∥∥∥ ∂̂L∂θ − ∂L∂θ ∥∥∥∥∥ ≤ ∥∥∥∥∥A− ∂F∂θ ( I − ∂F ∂h )−1∥∥∥∥∥ ∥∥∥∥∂L∂h ∥∥∥∥ . (21) Suppose the gradient ∂L/∂h is almost surely bounded. By Theorem 2, the condition in (19) can be guaranteed as long as a sufficiently large k is selected.

heading: 3 Experiments
text: In this section, we aim to answer the following questions via empirical results: (1) How precise is the phantom gradient? (2) What is the difference between the unrolling-based and the Neumannseries-based phantom gradients? (3) How is the phantom gradient influenced by the hyperparameters k and λ? (4) How about the computational cost of the phantom gradient compared with implicit differentiation? (5) Can the phantom gradient work at large-scale settings for various tasks?
We have provided some theoretical analysis and intuitions to (1), (2), and (3) in Sec. 2.4. Now we answer (1) and (2) and demonstrate the performance curves under different hyperparameters k and λ on CIFAR-10 [31]. Besides, we also study other factors that have potential influences on the training process of the state-of-the-art implicit models [2, 3] like pretraining. For (4) and (5), we conduct experiments on large-scale datasets to highlight the ultra-fast speed and competitive performances, including image classification on ImageNet [10] and language modeling on Wikitext-103 [13].
We start by introducing two experiment settings. We adopt a single-layer neural network with spectral normalization [32] as the function F and fixed-point iterations as the equilibrium solver, which is the synthetic setting. Moreover, on the CIFAR-10 dataset, we use the MDEQ-Tiny model [3] (170K parameters) as the backbone model, denoted as the ablation setting. Additional implementation details and experimental results are presented in Appendix D2.
Precision of the Phantom Gradient. The precision of the phantom gradient is measured by its angle against the exact gradient, indicated by the cosine similarity between the two. We discuss its precision in both the synthetic setting and the ablation setting. The former is under the static and randomly generated weights, while the latter provides characterization of the training dynamics.
In the synthetic setting, the function F is restricted to be a contraction mapping. Specifically, we directly set the Lipschitz constant of F as Lh = 0.9, and use 100 fixed-point iterations to solve the root h∗ of Eq. (1) until the relative error satisfies ‖h−F(h, z)‖/‖h‖ < 10−5. Here, the exact gradient is estimated by backpropagation through the fixed-point iterations, and cross-validated by implicit differentiation solved with 20 iterations of the Broyden’s method [16]. In our experiment, the cosine similarity between these two gradient estimates consistently succeeds 0.9999, indicating the gradient estimate is quite accurate when the relative error of forward solver is minor. The cosine similarity between phantom gradients and exact gradients is shown in Fig. 1. It shows that the cosine similarity tends to increase as k grows and that a small λ tends to slow down the convergence of the phantom gradient, allowing it to explore in a wider range regarding the angle against the exact gradient.
In the ablation setting, the precision of the phantom gradient during the training process is shown in Fig. 2. The model is trained by implicit differentiation under the official schedule3. It shows that the phantom gradient still provides an ascent direction in the real training process, as indicated by the considerable cosine similarity against the exact gradient. Interestingly, the cosine similarity slightly decays as the training progresses, which implies a possibility to construct an adaptive gradient solver for implicit models.
2All training sources of this work are available at https://github.com/Gsunshine/phantom_grad. 3Code available at https://github.com/locuslab/mdeq.
To Pretrain, or not to Pretrain? To better understand the components of the implicit models’ training schedule, we first illustrate a detailed ablation study of the baseline model in the ablation setting. The average accuracy with standard deviation is reported in Tab. 1.
The MDEQ model employs a pretraining stage in which the model F is unrolled as a recurrent network. We study the impact of the pretraining stage, the Dropout [33] operation, and the optimizer separately. It can be seen that the unrolled pretraining stabilizes the training of the MDEQ model. Removing the pretraining stage leads to a severe performance drop and apparent training instability among different trials because the solver cannot obtain an accurate fixed point h∗ when the model is not adequately trained. This ablation study also suggests that the MDEQ model is a strong baseline for our method to compare with.
However, pretraining is not always indispensable for training implicit models. It introduces an extra hyperparameter, i.e., how many steps should be involved in the pretraining stage. Next, we discuss how the UPG could circumvent this issue.
Trade-offs between Unrolling and Neumann. For an exact fixed point h∗, i.e., h∗ = F(h∗, z), there is no difference between UPG and NPG. However, when the numerical error exists in solving h∗, i.e., ‖h∗−F(h∗, z)‖ > 0, these two instatiations of the phantom gradient can behave differently.
We note that a particular benefit of the UPG is its ability to automatically switch between the pretraining and training stages for implicit models. When the model is not sufficiently trained and the solver converges poorly (see [3]), the UPG defines a forward computation graph that is essentially equivalent to a shallow weight-tied network to refine the coarse equilibrium states. In this stage, the phantom gradient serves as a backpropagation through time (BPTT) algorithm and hence behaves as in the pretraining stage. Then, as training progresses, the
solver becomes more stable and converges to the fixed point h∗ better. This makes the UPG behave more like the NPG. Therefore, the unrolled pretraining is gradually transited into the regular training phase based on implicit differentiation, and the hyperparameter tuning of pretraining steps can be waived. We argue that such an ability to adaptively switch training stages is benign to the implicit models’ training protocol, which is also supported by the performance gain in Tab. 1.
Although the UPG requires higher memory overhead than implicit differentiation or the NPG, it does not surpass the peak memory usage in the entire training protocol by implicit differentiation
due to the pretraining stage. In the ablation setting, the MDEQ model employs a 10-layer unrolling for pretraining, which actually consumes double memory compared with a 5-step unrolling scheme, e.g.,A5,0.5 in Tab. 1. In Tab. 2, we also demonstrate the time and memory complexity for implicit differentiation and the two forms of phantom gradient.
In addition, leaving the context of approximate and exact gradient aside, we also develop insights into understanding the subtle differences between UPG and NPG in terms of the state-dependent and state-free gradient. Actually, the UPG is state-dependent, which means it corresponds to the “exact” gradient of a computational sub-graph. Both the NPG and the gradient solved by implicit differentiation, however, do not exactly match the gradient of any forward computation graph unless the numerical error is entirely eliminated in both the forward and the backward passes for implicit differentiation or the one-step gradient [34, 21, 30, 22] is used for the NPG, i.e., k = 1. Interestingly, we observe that models trained on the state-dependent gradient demonstrate an additional training stability regarding the Jacobian spectral radius, compared with those trained on the state-free counterpart. We empirically note that it can be seen as certain form of implicit Jacobian regularization for implicit models as a supplement to the explicit counterpart [17], indicated by the stable estimated Jacobian spectral radius during training, i.e., ρ(∂Fλ/∂h) ≈ 1. The experimental results also cohere with our insight. The performance curves in Fig. 3 demonstrate the influence of λ and k and further validate that the UPG is more robust to a wide range of steps k than the NPG. In fact, when the Jacobian spectral radius ρ(∂F/∂h) increases freely without proper regularization, the exploding high-order terms in the NPG could exert a negative impact on the overall direction of the phantom gradient, leading to performance degradation when a large k is selected (see Fig. 3(b)). We also observe a surging of the estimated Jacobian spectral radius as well as the gradient explosion issue in the NPG experiments. On the contrary, the UPG can circumvent these problems thanks to its implicit Jacobian regularization.
Phantom Gradient at Scale. We conduct large-scale experiments to verify the advantages of the phantom gradient on vision, graph, and language benchmarks. We adopt the UPG in the large-scale experiments. The results are illustrated in Tab. 3 and Tab. 4. Our method matches or surpasses the implicit differentiation training protocol on the state-of-the-art implicit models with a visible reduction on the training time. When only considering the backward pass, the acceleration for MDEQ can be remarkably 12× on ImageNet classification.

heading: 4 Related Work
text: Implicit Models. Implicit models generalize the recursive forward/backward rules of neural networks and characterize their internal mechanism by some pre-specified dynamics. Based on the dynamics, the implicit mechanisms can be broadly categorized into three classes: ODE-based [1, 7], root-solving-based [2, 3, 4, 5, 11], and optimization-based [35, 36, 37] implicit models.
The ODE-based implicit models [1, 7] treat the iterative update rules of residual networks as the Euler discretization of an ODE, which could be solved by any black-box ODE solver. The gradient of the differential equation is calculated using the adjoint method [38], in which the adjoint state is obtained by solving another ODE. The root-solving-based implicit models [2, 5, 3, 4, 6, 11, 22] characterize layers of neural networks by solving fixed-point equations. The equations are solved by either the black-box root-finding solver [2, 3] or the fixed-point iteration [4, 22]. The optimization-based implicit models [35, 36, 37, 39, 34, 21, 40, 41] leverage the optimization programs as layers of neural networks. Previous works have studied differentiable layers of quadratic programming [35], submodular optimization [36], maximum satisfiability (MAXSAT) problems [37], and structured decomposition [21]. As for the backward passes, implicit differentiation is applied to the problemdefining equations of the root-solving-based models [2, 3] or the KKT conditions of the optimizationbased models [35]. As such, the gradient can be obtained from solving the backward linear system.
In this work, we focus on the root-solving-based implicit models. Theoretical works towards root-solving-based implicit models include the well-posedness [5, 4], monotone operators [11], global convergence [42, 41], and Lipschitz analysis [15]. We look into the theoretical aspect of the gradient-based algorithm in training implicit models and the efficient practice guidance. With these considerations, we show that implicit models of the same architecture could enjoy faster training speed and strong generalization in practical applications by using the phantom gradient.
Non-End-to-End Optimization in Deep Learning. Non-end-to-end optimization aims to replace the standard gradient-based training of deep architectures with modular or weakly modular training without the entire forward and backward passes. Currently, there are mainly three research directions in this field, namely, the auxiliary variable methods [43, 44, 45, 46, 47, 48, 49], target propagation [50, 51, 52], and synthetic gradient [24, 53, 54]. The auxiliary variable methods [43, 44, 45, 46, 47, 48, 49] formulate the optimization of neural networks as constrained optimization problems, in which the layer-wise activations are considered as trainable auxiliary variables. Then, the equality constraints are relaxed as penalty terms added to the objectives so that the parameters and auxiliary variables can be divided into blocks and thus optimized in parallel. The target propagation method [50, 51, 52] trains each module by having its activations regress to the pre-assigned targets, which are
propagated backwards from the downstream modules. Specifically, the auto-encoder architecture is used to reconstruct targets at each layer. Finally, the synthetic gradient method [24, 53, 54] estimates the local gradient of neural networks using auxiliary models, and employ the synthetic gradient in place of the exact gradient to perform parameter update. In this way, the forward and backward passes are decoupled and can be executed in an asynchronous manner.
Our work is in line with the non-end-to-end optimization research since we also aims to decouple the forward and backward passes of neural networks. However, we show that finding a reasonable “target” or a precise gradient estimate is not always necessary in training deep architectures. Our paper paves a path that an inexact but well-conditioned gradient estimate can contribute to both fast training and competitive generalization of implicit models.
Differentiation through Implicit Dynamics. Differentiation through certain implicit dynamics is an important aspect in a wide range of research fields, including bilevel optimization [19, 9], metalearning [18, 55, 30], and sensitivity analysis [56]. Since the gradient usually cannot be computed analytically, researchers have to implicitly differentiate the dynamics at the converged point. The formula of the gradient typically contains a term of Jacobian-inverse (or Hessian-inverse), which is computationally prohibitive for large-scale models. (See Eq. (2) in our case.) Herein, several techniques have been developed to approximate the matrix inverse in the previous literature.
An intuitive solution is to differentiate through the unrolled steps of a numerical solver of the dynamics [57, 58, 8]. In particular, if a single step is unrolled, it reduces to the well-known one-step gradient [59, 18, 60, 34, 30, 21, 22], in which the inverse of Jacobian/Hessian is simply approximated by an identity matrix. On the contrary, unrolling a small number of steps may induce a bias [9], while the memory and computational cost grows linearly as the number of unrolled steps increases. Towards this issue, Shaban et al. [19] propose to truncate the long-term dependencies and differentiate through only the last L steps. In fact, if the dynamics have converged to a stationary point, the finite-term truncation in Shaban et al. [19] is exactly the Neumann approximation of the Jacobian-inverse with the first L terms. Based on this, Lorraine et al. [9] directly use the truncated Neumann series as an approximation of the Jacobian-inverse. Besides the unrolling-based methods, optimization-based approaches [61, 55] have been studied in this field as well. Since the Jacobian-inverse-vector product can be viewed as solution of a linear system, algorithms like the conjugate gradient method can be used to solve it.

heading: 5 Limitation and Future Work
text: The main limitation of this work lies in the hyperparameter tuning of the phantom gradient, especially for the damping factor λ, which directly controls the gradient’s precision and conditioning, the implicit Jacobian regularization for UPG, the stability for NPG, and the final generalization behaviors. However, it has not been a hindrance to the application of phantom gradients in training implicit models as one can tune the hyperparameter according to the validation loss in the early training stage.
Regarding future works, we would like to highlight the following aspects: (1) eliminating the bias of the current phantom gradient, (2) constructing an adaptive gradient solver for implicit models, (3) analyzing the damping factor to provide practical guidance, (4) investigating the implicit Jacobian regularization, and (5) understanding how different noises in the gradients can impact the training of implicit models under different loss landscapes.

heading: 6 Conclusion
text: In this work, we explore the possibility of training implicit models via the efficient approximate phantom gradient. We systematically analyze the general condition of a gradient estimate so that the implicit model can be guaranteed to converge to an approximate stationary point of the loss function. Specifically, we give a sufficient condition under which a first-order oracle could always find an ascent direction of the loss landscape in the training process. Moreover, we introduce two instantiations of the proposed phantom gradient, based on either the damped fixed-point unrolling or the Neumann series. The proposed method shows a 1.4 ∼ 1.7× acceleration with comparable or better performances on large-scale benchmarks. Overall, this paper provides a practical perspective on training implicit models with theoretical guarantees.

heading: Acknowledgments
text: Zhouchen Lin was supported by the NSF China (No.s 61625301 and 61731018), NSFC Tianyuan Fund for Mathematics (No. 12026606) and Project 2020BD006 supported by PKU-Baidu Fund. Yisen Wang was partially supported by the National Natural Science Foundation of China under Grant 62006153, and Project 2020BD006 supported by PKU-Baidu Fund. Shaojie Bai was sponsored by a grant from the Bosch Center for Artificial Intelligence.

Review summaries:
The paper presents a method called phantom gradient to increase the efficiency in training implicit deep models by approximating the exact gradient. The authors first exploit Neumann series for the damped inverse matrix in implicit gradient calculation. They then truncate the series to k terms to arrive at finite term approximation of the exact implicit gradient. They finally provide theoretical justification for when phantom gradients are descent directions for implicit models. The authors evaluate their methods by conducting experiments on thoroughly analyzing their descent direction and demonstrating acceleration over a current implicit training method.

The paper proposes an approximate gradient, called phantom gradient, for training implicit models. They highlight that the method reduces the cost of gradient computation and it provides a descent direction suitable to train models. Their gradients are based on two methods, fixed-point unrolling and the Neumann series. The paper proves that under some network conditions on the eigenvalues of the derivative of the implicit module, the phantom gradient has a descent direction positively corelated to the exact gradient. The paper also provide almost surely convergence of the phantom gradient to the exact gradient under some conditions.

This paper focuses on the problem efficiently training Deep Equilibrium Models (DEQs) by estimating the gradient to the "equilibirum layer." Usually, this gradient is computed using a black-box fixed point solver (generally Broyden's), which makes the backward pass expensive. This paper introduces a gradient estimate, which they term the _phantom gradient_. They show that the dot product between this gradient and the true gradient is greater than 0, which means it constitutes a descent direction in a first-order sense. They empirically verify that these estimates are good using CIFAR-10 and ImageNet, where they get performance remarkably close to the standard Multiscale DEQ [1] with a reasonable training speed improvement. They also show that the phantom gradient approaches the true gradient in terms of cosine similarity with increasing series elements $k$.