On Training Implicit Models
Authors: Zhengyang Geng
Xin-Yu Zhang
Shaojie Bai
Yisen Wang
Zhouchen Lin
Venue: NeurIPS 2021 Poster (2021)
Keywords: implicit models
deep learning
optimization
training

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 $\times$ and even boost the performance over approaches based on the exact gradient on ImageNet.

On Training Implicit Models
Zhengyang Gengl.
Xun-Yu Zhang**
Shaojie Bai
Yisen Wang
Zhouchen 23n
'Zeejiang Lab, China
2 Key Lab. of Machine Perception, School of AI, Peking University IInstitute for Artificial Intelligence, Peking University
Carnegie Mellon University PPazhou Lab, China
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 il 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 x, and even boost the performance over approaches based on the exact gradient on ImageNet.
Introduction
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, 4, 5] have attracted increasing attentions and are able to match the state-of-the-ard results by explicit models 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, 3J. As for the backward passes, however, directly differentiating through the forward pass trajectories could induce heavy memory overhead [8, 9]. To this end, researchers have developed memory-efficeee via implicit differentiation. such as solving a Jacobian- based linear exed-poinn 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 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 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 GPUs.
backpropagation
This work investigates fast approximate gradients for training implicit models. We found that first- order 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 (NcurIPS 2021).

---

models, circumventing laboriously computing the exact gradient as in prior arts [2, 3, 4, 11, 12]. We develop a framework 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-poing 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 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 12x on ImageNet. We believe that our results provide strong evidence for effectively training implicit models with the lightweight phantom gradient.
2 Method
2.1 Inspection of Implicit Differentiation
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- h"" 2),
(1)
where 2 E Rdu+do is the union of the module's input u E TR and parameters 0 E RRdo i.e., 2 (u=" J... Here, a is usually projection of the original data point x E RR e.g., u = M(x). In this section, assume F is contraction mapping w.r.f. h so that its Lipschitz constant La w.r.f. h is less than one, i.e., Ln < 1, a setting that has been analyzed in recent works [14, 15)'.
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
dh
0F
OF
dz
dz h- Ōh
Here, (Ba/Ob).; 0b. The equilibrium point he of Bq. (1) is then passed to post-processing function 9 to obtain prediction ý G(h"). In the generic learning scenario, the training objective is the following expected loss:
R = p [[(((;; ( 0),
(3)
where y is the groundtruth corresponding to the training example æ. 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 training explicit neural networks. The most crucial component is the gradient of the loss function w.n.t. the input vector x = [u" 1 which is used to train both the implicit module F and the input projection module M. Using Eq. (2) with the condition h = h², we have
OC
0F
0F
-1
Oc
OC
OF
0F
-1 Oc
r-
(4)
ou
du
0h
0h
1
00
00
Öh ah
The gradients in Eq. (4) are in the same form w.r.t. u and 0. Without loss of generality, we only
discuss the gradient w.r.t in the following sections.
Note that the contraction condition could be largely relaxed to the one that the spectral radius of 0F/0h' on the given data is less than 1, iee. p(0F/Oh") ) < 1, as indicated by the well posedness condition in [5].

---

2.2 Motivation
The most intriguing part lies in the grrobian-inverss term, i.e. (I /F/ah) Computing the inverse term by brute force is intractable due to the O(n²) complexity. Previous implicit models [2] approach this by solving linear system involving Jacobian
product iteratively via gradient solver, introducing over 30 Broyden [16] iterations in the backward pass. However, the scale of the Jacobian matrix can exceed 10 x 10 the real scenarios, leading to a prohibitive cost in computing the exact gradient. For example, training small-scale state-of-the- 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.
implicit model
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 cerobian-inverey 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-invers, 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 regularization approach [25], which has been shown play central role in escaping poor local minima and improving generalization ability [26, 27, 28).
The concerns and observations motivate us rethink the possibility of replacing the nicobian-inverse inverse term in the standard implicit differentiation with a cheaper and more stable counterpart. We believe that 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 Oh") /00 is replaced with matrix A, and the corresponding phantom gradient defined as
ac ac A
(5)
dh 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-poine -oint unrolling or the Neumann series (Sec. 2.4). The proofs of all our theoretical results are presented in Appendix C.
00
2.3 General Condition the Phantom Gradient
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 0F /00. If
A
0F
OF
A
(6)
oh
00
G
then the phantom gradient provides an ascent direction of the function C. i.e..
ØL
ac
0.
(7)
00
ae
Remark 1. Suppose only the (I 1J-07 (0//0)) term is replaced with matrix D, namely, A
(DF/00) D. Then, the condition in (6) can be reduced into
|
|
0F
1
(8)
0h
R2
where k is the condition number of 0F/00. (See Appendix C.I for the derivation.)
3

---

2.4 Instantiations of the Phantom Gradient
In this section, we present two practical instantiations of the phantom gradient. We also verify that the general condition in Theorem can be satisfied if the hyperparameters in our instantiations are wisely selected.
Suppose we hope to differentiate through implicit dynamics, egg. either a root-solving process or 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, we solve the root of Eq. (I) via the fixed-point iteration:
h4+1 Flhs, 2),
0, T 1,
then by differentiating through the unrolled steps of Eq. (9), we have
öhr
T-1
OF
T-1
OF
II
(10)
00
00
dh
t=
be s=t+1
hh,
Besides, the Neumann series of the Jacobian-inverse (I - (0)
OF
0F
OF
I +
+
(11)
0h
0h
Oh
Notably, computing the Jacobian 0h" 0h"/00 using the Neumann series in (11) is equivalent to differen- tiating 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), consider damped variant of the fixed-point iteration:
h++1 = (h, x) = XF x) + (1 - A)hr,
t t= 0, . 1.
(12)
Differentiating through the unrolled steps of Eq. (12), Eq. (10) is adapted
ohT
T-1
OF
T-1
II
OF
+(1-
(13)
00
00
0h
1=0
a=t+1
The Neumann series of (I 1--0 OF (O)) -1 is correspondingly adapted as
OF Jh.
( B B B +..., where B + (1 X)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 (14) converges to the Jacobian- inverse (I (F/ as well.
Theorem 2. Suppose the Jacobian 0F / is contraction mapping. Then,
(i) the Neumann series (14) converges the iecobiaa-invverse (F/ and
(ii) the function F is continuously differentiable w.r.t. both and 0. the sequence in Eq. (13) converges to the exact Jacobian Oh /00 as T + 0o, i.e.,
Ohr
0F
0F
lim T-000
I
(15)
00
00 h
Öh
However, discussed in Sec. 2.2, it is unnecessary to compute the exact gradient with infinite terms. the following context, we introduce two instantiations of the phantom gradient based on the finite-term truncation of Eq. (13) or (14).
Unrolling-band based Phantom Gradient (UPG).
In the unrolling form, the matrix A is defined as
OF
0F
pnr
II
+(1-1)1)
(16)
80
Th hh
bi s=t-1

---

Neumann-series-banee Phantom Gradient (NPG).
In the Neumann form, the matrix A is defined
as
0F 00| h-
OF (I+ B B B) where B X Sh h
A nCu
+ X)I. (17)
Note that both the initial point of the fixed-poine iteration (i.e., ho in Eq. (16) and the point at which the Neumann series evaluated (i.e. h" in (17)) are the solution of the root-finding solver. (See Appendix B for implementation of the phantom gradient.)
According Theorem 2, the matrix A defined by either Eq. (16) (17) converges the exact Jacobian 0h' /00 as l - oo for any À E (0, 1]. Therefore, by Theorem 2, the condition in (6) can be satisfied if a sufficiently large step & is selected, since
0F
OF
OF
0F
A
I
S (1+ LNd
A
(18)
0h
00
00
Ōh
Next, characterize the impact of the two hyperparameters, i.e., k and X, on the precision and conditioning of A. Take the NPG (Eq. (17)) as an example.
(i) On the precision of the phantom gradient,
a large 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 decreases.
iii
On the conditioning of the phantom gradient,
a large impairs the conditioning of A since the condition number of B grows exponentially as & increases, while a small À helps maintain a small condition number of A because the singular values of 0F/ /F/Oh are "smoothed" by the identity matrix.
In a word, a large k is preferable for more accurate A, while small contributes to the well- conditioning of A. Practically, these hyperparameters should be selected in consideration of balanced trade-off between the precision and conditioning of A. See Sec. 3 for experimental results.
2.5 Convergence Theory
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 eттоr of the phantom gradient is sufficiently small, the SGD algorithm converges to an e-approximate stationary point in the expectation sense. We will discuss the feasibility of our assumptions in Appendix C.3.
Theorem 3. Suppose the loss function R in Eq. (3) is t-smooth, lower-bounded, and has bounded gradient almost surely in the training process. Besides, assume the gradient in Eq. (4) is unbiased estimator of (R(0) with bounded covariance. If the phantom gradient in Eq. (5) is an e-approximation to the gradient in Eq. (4), i.e.
0C
ac
Sc,
almost surely,
(19)
00
00
then using Eq. (5) as a stochastic first-order oracle with a step size on = O(1/ Vn) to update with gradient descent, it follows after N iterations that
BT. VR((0)
I 50 c+
log
E
(20)
IN
E Tn
VN
Remark 2. Consider the condition in (19):
Øc
ac
aF
aF
0c
<
A
I
(21)
00
00
30
0h
0h
Suppose the gradient OC/ah h is almost surely bounded. By Theorem 2, the condition in (19) can be guaranteed as long sufficiently large & is selected.
5

---

À 0.5
A 1.0
À 0.5
1.0
10
3.9
0.52
6:34
5.50
0.98
1.20
0.90
0:92
0.94
0.35
0.58
1.0
0.50
0.92
2.94
0.56
3.00
0.92
0:94
0.96
0.9
1.00
Cosire Sis
Coslvo Seniarty
Codire Simibeity
(a) Neumann series based phaniom grudient.
(b) Unrolling baved phantoen gradient.
Figure 1: Cosine similarity between the phantom gradient and the exact gradient in the synthetic setting.
3 Experiments
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 Neumann- series-based phantom gradients? (3) How is the phantom gradient influenced by the hyperparameters k and X? (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 (1), (2), and (3) in Sec. 2.4. Now answer (1) and (2) and demonstrate the performance curves under different hyperparameters and on CIFAR-10 [31]. Besides, we also study other factors that have potential influences the training process of the ofate-of-the-art implicit models [2, 3] like pretraining. For (4) and (5), we conduct experiments on large-scale datasets 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 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, the CIFAR-10 10 dataset. parameters) as the backbone model, denoted as the ablation setting. Additional implementation details and experimental results are presented in Appendix D
use the MDEQ model [3] (170K
Precision of the Phantom Gradient. The precision of the phantom gradient 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 under the static and randomly generated weights, while the latter provides characterization of the training dynamics.
In the synthetic setting, the function F restricted to be a contraction mapping. Specifically, we directly set the Lipschitz constant of F as Ln. 0.9, and use 100 fixed-point iterations to solve the root h* of Eq. (1) until the relative error satisfies | - F(h, /)]l < 10-5. Here, the exact gradient is estimated by backpropagation through the fixed-point iterations, and cross-validatee implicit differentiation solved with 20 iterations of the Broyden's '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 grows and that small À tends to slow down the convergence of the phantom gradient, allowing it to explore in wider range regarding the angle against the exact gradient.
by
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 schedule". shows that the phantom gradient still provides ascent direction in the real training process, 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 adaptive gradient solver for implicit models.
2 All training sources this work are available at httpp:://giiie com http:/hhhw.....//mmmmeeeeeeaaaaa
Code available at https:/ https:///gtthtn com/locuslab/mden

---

(a) Neumann-scries- phantom gradient.
(b) Unrolling-base phantom gradient.
Figure 2: Cosine similarity between the phantom gradient and the exact gradient in the real scenario. The horizontal axis corresponds to the cosine similarity, and the vertical axis to the training step.
To Pretrain, or not to Pretrain? To better understand the components of the implicit models' training schedule, we first illustrate detailed ablation study of the baseline model in the ablation setting. The average accuracy with standard deviation reported in Tab. 1.
The MDEQ model employs a pretraining stage in which the model F unrolled as a recurrent network. We study the impact of the pretrain- ing 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 severe performance drop and apparent training instability among different tri- als 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 strong baseline for our method to compare with.
Table Ablation settings on CIFAR-10.
Method
Acc. (%)
Implicit Differentiation
85.0 + 0.2
w/o Pretraining
82.3 = 1.3
w/o Dropout
83.7 + 0.1
Adam - SGD
84.5 + 0.3
SGD w/o Pretraining
82.9 1.5
UPG (A5,0.5. w/o Dropout)
85.8 NPG (A5.0.5, w/o Dropout)
85.6 + 0.5
UPG (A9.0.5, w/ Dropout)
86. +.5
However, pretraining is not always indispensable for training implicit models. It introduces extra hyperparameter, i.e., how many steps should involved in the pretraining stage. Next, we discuss how the UPG could circumvent this issue.
Trade-off's between Unrolling and Neumann. For an exact fixed point h", i.e,, h- = F(h", 2), there is no difference between UPG and NPG. However, when the numerical error exists in solving h" ie., ||h - F(h =)Il O, 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 pretrain- ing and training stages for implicit models. When the model is not sufficiently trained and the solver converges poorly (see [3]), the UPG defines for- ward computation graph that is essentially equiva- lent shallow weight-tied network to refine the coarse equilibrium states. In this stage, the phan- tom gradient serves as a backpropagation through time (BPTT) algorithm and hence behaves in the pretraining stage. Then, as training progresses, the solver becomes more stable and converges to the fixed point ht better. This makes the UPG behave more like the NPG. Therefore, the unrolled pretraining is gradually transited into the regular training phase based implicit differentiation, and the hyperparameter tuning of pretraining steps be waived. We argue that such an ability to adaptively switch training stages is benign to the implicit models' training protocol, which also supported by the performance gain in Tab. 1.
Table 2:
Complexity comparison. Mem means the memory cost, and K and k denote the solver's 's steps and the unrolling/Neumman steps, respectively. Here, K >>> 1.
Method Implicit UPG NPG
Time
Mem
Peak Mem
O(K)
O(1)
O(k)
O(k)
O(k)
O(k)
O(k)
O(1)
0(1)
Although the UPG requires higher memory overhead than implicit differentiation or the NPG, does not surpass the peak memory usage in the entire training protocol by implicit differentiation
7

---

05
12E
461
Impact and
NPG UPG
Figure 3: Ablation studies on (a) the hyperparameters and k and (b) two forms of phantom gradient.
due to the pretraining stage. In the ablation setting. the MDEQ model employs 10-layer unrolling for pretraining, which actually consumes double memory compared with 5-step unrolling scheme, e,g,, A5.0.5 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 "exaet" gradient of 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 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 state-free counterpart, We empirically note that can be seen as certain form of implicit Jacobian regularization for implicit models supplement to the explicit counterpart [17], indicated by the 1.
stable estimated Jacobian spectral radius during training, i.e., p(OFx
The experimental results also cohere with our insight. The performance curves in Fig. 3 demonstrate the influence of and and further validate that the UPG is more robust wide range of steps than the NPG. In fact, when the Jacobian spectral radius p(aFF /(h) increases freely without proper regularization, the exploding high-order terms in the NPG could exert negative impact on the overall direction of the phantom gradient, leading to performance degradation when large & is selected (see Fig. 3()). also observe surging explosion issue in the NPG experiments. On the contrary, the UPG can circumvent these problems thanks to its implicit Jacobian regularization.
estimated Jacobian spectral radius well as the gradient
Table 3: Experiments using DEQ [2] and MDEQ [3] vision and language tasks. Metrics stand for accuracy(%) ( for image classification on CIFAR-10 10 and ImageNet, and perplexity for language modeling Wikitext- 103. JR stands for Jacobian Regularization [17]. t indicates additional steps in the forward equilibrium solver.
Datasets
Model
Method
Params
Metries
Speed
CIFAR-10
MDEQ
Implicit
10M
93.8 + 0.17
1.0x
CIFAR-10
MDEQ
UPG As.0.5
10M
95.0 + 0.16
1.4x
ImageNet
MDEQ
Implicit
18M
75.3
1.0x
ImageNet
MDEQ
UPG Ar.0.6
18M
75.7
1.7x
Wikitext-103
DEQ (PostLN)
Implicit
98M
24.0
1.0x
Wikitext-103
DEQ (PostLN)
UPG As.o.
98M
25.7
1.7x
Wikitext-103
DEQ (PreLN)
JR + Implicit
98M
24.5
1.7x
Wikitext-101
DEQ (PreLN)
JR + UPG As,a.s
98M
24.4
2.2x
Wikitext-103
DEQ (PreLN)
JR + UPG As,0.8
98M
24.0t
1.7x

---

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 otate-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 12x on ImageNet classification.
Table 4: Experiments using IGNN [4] graph tasks. Metrics stand for accuracy( ( for graph classification on COX2 and PROTEINS, Micro-FI($6) for node classification on PPI.
Datasets
Model
Method
Params
Metrics (%)
COX2
IGNN
Implicit
38K
84.1 + 2.9
COX2
IGNN
UPG A5,0.5
38K
83.9 + 3.0
COX2
IGNN
UPG A5,0.8
38K
83.9 2.7
COX2
IGNN
UPG A5,10
38K
83.0 1.0
PROTEINS
IGNN
Implicit
34K
78.6 4.1
PROTEINS
IGNN
UPG A5,0.5
34K
78.4 4.2
PROTEINS
IGNN
UPG A5,0.8
34K
78.6
PROTEINS
IGNN
UPG 15,10
34K
78.8 18.8=0 4.2
PPI
IGNN
Implicit
4.7M
97.6
PPI
IGNN
UPG A5,0.5
4.7M
98.2
PPI
IGNN
UPG A5,0.8
4.7M
97.4
PPI
IGNN
UPG A5,10
4.7M
96.2
4 Related Work
Implicit Models. Implicit models generalize the recursive forward/backward rules of neural net- works 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-baet [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-bet layers of neural networks by solving fixed-poit equations. The equations are solved by either the black-box root-finding solver [2, 3] or the fixed-point iteration [4, 22]. The optimization-bai 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 problem- defining equations of the root-solving- based models [2, 3] or the KKT conditions of the optimization- based models [35]. As such, the gradient can be obtained from solving the backward linear system.
implicit models [2. 5, 3, 4, 6, 11, 22] characterize
In this work, we focus on the voot-solving-basee implicit models. Theoretical works towards noot-solving-baseee implicit models include the well- well-posedie poleedossssess [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.
Optimization in Deep Learning. Non-end- -ened-t-eend 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- pre-assigeed 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 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], meta- learning [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- -nssesnninverre),, 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 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 small number of steps may induce 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 I 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-base approaches [61, 55] have been studied in this field as well. Since the can be viewed as solution of a linear system, algorithms like the conjugate gradient method can be used to solve it.
methods, optimization-ban vector product
5 Limitation and Future Work
The main limitation of this work lies in the hyperparameter tuning of the phantom gradient, especially for the damping factor X, 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.
Conclusion
In this work, we explore the possibility of training implicit models via the efficient approximate phantom gradient. We systematically analyze the general condition of 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 1.4 ~ 1.7x acceleration with comparable or better performances on large-scale benchmarks. Overall, this paper provides practical perspective on training implicit models with theoretical guarantees.
10

---

Acknowledgments
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 grant from the Bosch Center for Artificial Intelligence.
References
[1] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural Ordinary Differential Equations. In Neural Information Processing Systems (NeurlPS), 2018. 1,9
[2] Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep Equilibrium Models. In Neural Information Processing Systems (NeurIPS), 2019. 1, 2, 6, 8,
[3] Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale Deep Equilibrium Models. In Neural Information Processing Systems (NeurIPS), pages 5238-5250, 2020. 1, 2, 6, 7, 8, 9
[4] Fangda Gu, Heng Chang, Wenwu Zhu, Somayeh Sojoudi, and Laurent El Ghaoui. Implicit Graph Neural Networks. In Neural Information Processing Systems (NeurlPS), pages 1 11984- 11995, 2020. 1, 2, 9
[5] Laurent EI Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai. Implicit Deep Learning. SIAM Journal on Mathematics of Data Science, 3(3):930-955) 15809958,.00 2021. 1,
[6] Tiancai Wang, Xiangyu Zhang, and Jian Sun. Implicit Feature Pyramid Network for Object Detection. arXiv preprint arXiv:2012.135633 2020. 1,9
[7] Emilien Dupont, Arnaud Doucet, and Yee Whye Tch. Augmented Neural ODEs. In Neural Information Processing Systems (NeurIPS), 2019. 1,9
[8] Luca Franceschi, Michele Donini, Paolo Frasconi, and Massimiliano Pontil. Forward and Reverse Grasient-Based Hyperparameter Optimization. In International Conference on Machine Learning (ICML), pages 1165-1173, 2017. 1, 10
[9] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing Millions of Hyperparameters
by Implicit Differentiation. In International Conference on Artificial Intelligence and Statisties (AISTATS), pages 1540-1552, 2020. 1, 1,3,4 4, 10
[10] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet: Large Scale Visual Recognition Challenge. International Journal on Computer Vision (IJCV), 11(())211122))
2015. 1, 2, 6
[11] Ezra Winston and J. Zico Kolter. Monotone operator equilibrium networks. In Neural Informa- tion Processing Systems (NeurIPS), pages 10718-10728, 2020. 2,
[12] Cheng Lu, Jianfei Chen, Chongxuan Li, Qiuhao Wang, and Jun Zhu. Implicit Normalizing Flows. In International Conference on Learning Representations (ICLR), 2021. 2
[13] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. In International Conference Learning Representations (ICLR), 2017. 2,6 6
[14] Chirag Pabbaraju, Ezra Winston, and J. Zico Kolter. Estimating Lipschitz constants of monotone deep equilibrium models. International Conference on Learning Representations (ICLR), 2021. 2
[15] Max Revay, Ruigang Wang, and Ian R Manchester. Lipschitz Bounded Equilibrium Networks. arXiv:2010.01732. 2020. 2, 9
[16) Charles G Broyden. A Class of Methods for Solving Nonlinear Simultaneous Equations. Mathematics of computation, 19(9)):571-593, 1965. 6,
[17] Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Stabilizing Equilibrium Models by Jacobian Regularization. In International Conference Machine Learning (ICML), 2021. 3, 8
[18] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks. In International Conference on Machine Learning (ICML), pages 1126-11323 2017. 3, 10
11

---

[19] Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated Back- propagation for Bilevel Optimization. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 1723-1732, 2019. 3, 4, 10
[20] Jens Behrmann, David Kristjanson Duvenaud, and Jörn-Henrik Jacobsen. Invertible Residual Networks. In International Conference on Machine Learning (ICML), 2019. 3
[21] Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li, Ke Wei, and Zhouchen Lin. Is Attention Better Than Matrix Decomposition? In International Conference Learning Representations (ICLR), 2021. 3, 8, 9, 10
[22] Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley J. Osher, and Wotao Yin. Fixed Point Networks: Implicit Depth Models with Jacobian-Free Backprop. arXiv preprint arXiv:2103. 12803, 2021. 3, 8, 9, 10
[23] Mehrdad Yazdani. Linear Backprop in non-linear networks. In Neural Information Processing
Systems Workshops, 2018. 3
[24] Max Jaderberg, Wojciech Marian Czarnecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled Neural Interfaces using Synthetic Gradients. In International Conference on Machine Learning (ICML), pages 1627-1635, 2017. 3, 9, 10
[25] Xavier Gastaldi. Shake-Shake Regularization. arXiv preprint arXiv: / 705.07485, 2017. 3
[26] Guozhong An. The Effects of Adding Noise During Backpropagation Training on a Generaliza- tion Performance. Neural Computation, 8(3):643-674, 1996. 3
[27]
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects. In International Conference on Machine Learning (ICML), pages 7654-7663, 2019.
[28] Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the Noisy Gradient Descent that Generalizes as SGD. In International Conference on Machine Learning (ICML). pages 10367-10376. 2020. 3
[29] Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Joemn-Henrik Jacobsen. Residual Flows for Invertible Generative Modeling. Neural Information Processing Systems (NeurIPS),
pages 9916-9926, 2019. 3
[30] Xin- Yu Zhang, Taihong Xiao, Haolin Jia, Ming-Ming Cheng, and Ming-Hsuan Yang. Semi- Supervised Learning with Meta- reta-Gradient. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 73 11, 2021. 3, 8, 10
[31] Alex Krizhevsky. Learning Multiple Layers of Features from Tiny Images. Technical report, Citeseer, 2009.
[32] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral Normaliza- tion for Generative Adversarial Networks. In International Conference on Learning Represen- lations (ICLR), 2018. 6
[33] Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhut- dinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv:1207.0580, 2012. 7
[34] Bryan Wilder, Eric Ewing, Bistra Dilkina, and Milind Tambe. End to end learning and optimization on graphs. Neural Information Processing Systems (NeurIPS), pages 4672- 4683, 2019. 8, 9, 10
[35] Brandon Amos and J. Zico Kolter. OptNet: Networks. In International Conference on Machine Learning (ICML), pages 136-145, 2017. 9
Differentiable Optimization as a Layer in Neural
[36] Josip Djolonga and Andreas Krause. Differentiable Learning of Submodular Models. Neural Information Processing Systems (NeurIPS), pages 1013-1023, 2017. 9
[37] Po- Wei Wang, Priya Donti, Bryan Wilder, and Zico Kolter. SATNet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver. In International Conference on Machine Learning (ICML). pages 6545-6554, 2019.
[38] VG Boltyanskiy, Revaz. V Gamkrelidze, YEF Mishchenko, and L.S Pontryagin. Mathematical theory of optimal processes. 1962. 9
12

---

[39] Marin Vlastelica, Anselm Paulus, Vít Musil, Georg Martius, and Michal Rolinek. Differentiation of Blackbox Combinatorial Solvers. In International Conference on Learning Representations (ICI.R), 2020. 9
[40] Stephen Gould, Richard Hartley, and Dylan Campbell. Deep Declarative Networks. IEEE Transactions on Pattern Recognition and Machine Intelligence (PAMI), 2021. 9
[41] Xingyu Xie, Qiuhao Wang, Zenan Ling, Xia Li, Yisen Wang, Guangean Liu, and Zhouchen Lin. Optimization Induced Equilibrium Networks. arXiv preprint arXiv:2105.13228 2021. 9
[42] Kenji Kawaguchi. On the Theory of Implicit Deep Learning: Global Convergence with Implicit Layers. In International Conference Learning Reprexentations (ICI.R), 2020. 9
[43] Miguel Carreira- Carreira-Perpinte Perpinan and Weiran Wang. Distributed Optimization of Deeply Nested Systems. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 10-19, 2014. 9
[44] Gavin Taylor, Ryan Burmeister, Zheng Xu, Bharat Singh, Ankit Patel, and Tom Goldstein. Training Neural Networks Without Gradients: A Scalable ADMM Approach. In International Conference Machine Learning (ICML), pages 2722-2731, 2016. 9
[45] Ziming Zhang, Yuting Chen, and Venkatesh Saligrama. Efficient Training of Very Deep Neural Networks for Supervised Hashing. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 1487-1495, 2016. 9
[46] Ziming Zhang and Matthew Brand. Convergent Block Coordinate Descent for Training Tikhonov Regularized Deep Neural Networks. In Neural Information Processing Systems (NeurlPS), 2017. 9
[47] Jinshan Zeng, Shikang Ouyang, Tim Tiz-Kit Lau, Shaobo Lin, and Yuan Yao. Global Conver- gence in Deep Learning with Variable Splitting via the preprint arXiv: 00225, 2018. 9
Property. arXiv
[48] Jia Li, Cong Fang, and Zhouchen Lin. Lifted Proximal Operator Machines. In Association for the Advancement of Artificial Intelligence (AAAI). pages 4181-4188, 2019. 9
[49] Fangda Gu, Armin Askari, and Laurent El Ghaoui. Fenchel Lifted Networks: A Lagrange Relaxation of Neural Network Training. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 3362-3371, 2020. 9
[50] Yoshua Bengio. How Auto-Encodes Could Provide Credit Assignment in Deep Networks via Target Propagation. arXiv preprint /007.:140.79900 2014. 9
[51] Dong-Hyun Lee, Saizheng Zhang, Asja Fischer, and Yoshua Bengio. Difference Target Propa- gation. In International Conference on Learning Representations (ICLR), pages 498 008--15,55 2015.
[52] Alexander Meulemans, Francesco Carzaniga, Johan Suykens, João Sacramento, and Benjamin F. Grewe. A Theoretical Framework for Target Propagation. In Neural Information Processing Systems (NeurIPS), pages 20024-20036, 2020. 9
[53] Wojeiech Marian Czarnecki, Grzegorz Šwirszcz, Max Jaderberg, Simon Osindero, Oriol Vinyals, and Koray Kavukcuoglu. Understanding Synthetic Gradients and Decoupled Neural Interfaces. In International Conference on Machine Learning (ICML), pages 904-912 2017. 9, 10
[54] Benjamin James Lansdell, Prashanth Ravi Prakash, and Konrad Paul Kording. Learning to solve the credit assignment problem. In International Conference on Learning Representations (ICLR), 2020. 9, 10
[55] Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine Meta-Learning with Implicit Gradients. In Neural Information Processing Systems (NeurIPS), 2019. 10
[56] J Frédéric Bonnans and Alexander Shapiro. Perturbation analysis of optimization problems. Springer Science & Business Media, 2013. 10
[57] Justin Domke. Generic Methods for Optimization-Ban -asieaattti---Bedddddeed Modeling. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 318-328- 2012. 10
[58] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based Hyperparameter Optimization through Reversible Learning. In International Conference on Machine Learning (ICML), pages 2113-2122, 2015. 10
13

---

[59] Jelena Luketina, Mathias Berglund, Klaus Greff, and Tapani Raiko. Scalable Gradient-Bate Tuning of Continuous Regularization Hyperparameters. In International Conference on Machine Learning (ICML), pages 2952-2960, 2016. 10
[60] Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable Architecture Search. International Conference Learning Representations (ICLR), 2018. 10
[61] Fabian Pedregosa. Hyperparameter Optimization with Approximate Gradient. In International Conference on Machine Learning (ICML), pages 737-746, 2016. 10
14


Reviews:
[{"review_id": "vU9H4j33LU", "reviewer": "Reviewer_FWcz", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is nicely written and is well-organized. The paper explains well how their method differs from prior works. Although the work is incremental, the approximate gradient due to its reduced cost computation is of interest. Here are my detailed comments.\n\n- The term (1 - lambda) h_t in (13) is introducing a momentum. Please discuss the results from the perspective of momentum.\n\n- It would be nice to emphasize that when \\lambda = 1, the proposed gradient is reduced to backprop through the unrolled iterations (e.g., Figure 1b right). All reported results are based on lambda = 0.5, and k = 5. a) the choice reason should be explained. b) the experiments should cover a range for lambda and k, specially the case lambda = 1 (backprop but not exact gradient). Indeed, their approach should be compared against the case lambda = 1 to highlight the advantages of the phantom gradient.\n\n\n- At line 12 and 54, please be more specific that the speed up is compared to which methods (backprop, which implicit diff method)?\n\n- It would be nice to move some of the questions asked on paragraph line 174 to the introduction or contribution.\n\n- Although included in the supp, explain for the tables that A_{5,0.5} is for k = 5, and lambda = 0.5.\n\n- Place the tables and figures close to their descriptions.\n\n- It's nice that the code has been made available. I suggest to include a readme.\n\nMinor comments:\n\n- typo on supp line 20. It should be Remark.\n\n\n\n----------------------------------------------------------------------------------------------------------------\nPost response review\n----------------------------------------------------------------------------------------------------------------\n----------------------------------------------------------------------------------------------------------------\n\nI have found the authors' response to the reviews satisfactory. I strongly recommend to add some of the discussions made here to either the paper or supplementary, specifically elaboration on NTK, hyperparameter tuning, result discussion.\n\nOverall, given the modifications suggested by the reviewers, I have increased my score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations is not included. A need for tuning lambda can be one limitation of the gradient.\n\nI suggest to include some societal impact related to the computational cost of training implicit models and their robustness.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "oCYXA2_OsoE", "reviewer": "Reviewer_kurx", "summary": "This paper proposes a new method to estimate gradient for implicit models, which empirically accelerates  the training of implicit models.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall this is an interesting paper. The major contributions are:\n\n1. This paper proposes a new method to estimate gradient in training implicit models, where the idea is quite straight-forward and easy to follow; \n2. Some analyses are provided on how well the approximation could be, including some theoretical results and their intuitive explanations;\n3. Experimental results on CIFAR-10 and ImageNet are also provided, in which the proposed method shows its effectiveness in terms of accuracy and efficiency.\n\nBelow are some other comments and questions:\n\n1. There are three major theorems in the paper. The first one relies on the condition  based on singular values and the last one relies on the condition of good approximation. As mentioned in the paper, the condition for the first theorem is related to  NTK. Consequently, when the model is 'sufficiently wide', the condition can be easily satisfied. However, Theorem 3.1 & 3.2 in [Arora et al. 2019] suggest that this 'sufficiently wide' requirement could be rather impractical, as the required width  is much wider than those commonly used models in practice. Thus, the conditions for the first and third theorems are somewhat strong. Maybe the authors can consider adding an experiment to study the impact of different model designs. \n2. It would strengthen the paper if it compares more related works in experiments and reports all the results with mean (standard deviation);\n3. In the experiments, the proposed method leads to worse results sometime. It would be interesting to know whether this is due to the approximation quality or something else. Hope the authors can provide some detailed discussions on this.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concern is identified.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours."}, "extra_text": {}}, {"review_id": "bqNAlxzf4H", "reviewer": "Reviewer_hVtm", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Strength] The paper has the following strength.\n- The paper studies the very interesting and important problem of accelerating implicit model training. Which is also a difficult problem. The paper gives a novel insight that memorizing the first few steps of the forward fixed point iterations can help more efficient and stable gradient evaluation.\n- The proposed method gives theoretical analysis of the method on the validity of phantom gradients and convergence analysis of the method. \n- The method offers significant speedup in backward steps over the current implicit gradient method in model training and offer a moderate speedup overall.\n- Thorough empirical analysis on the property of the model is well conducted in the paper, showing cleverly chosen parameters are key to the speedup in implicit training. It is also demonstrated that the method scales well to large experiments in computer vision.\n\n[Weekness] The paper is well written but has a few glitches. I think the paper would be a good addition to NeurIPS if the authors can carefully address them.\n- The notations of the paper are not clearly given. (e.g. The norm and inner product of Theorem 1). I sugguest they are operator norm and matrix inner product. Though careful inspection could infer the notations, a clearer definition can help make the paper more readable for general public.\n- The paper should more adequately discuss related works such as (El Ghaoui, 2019; Revay, 2020; Gu, 2020) on theoretical contributions like well-posedness (refereed as well-conditioned in this paper) for implicit models and (Wang, 2020; Gu, 2020) on recent applications of implicit models on object detection and graph neural networks.\n- I have seen a varient of Theorem 2 in Theorem 6.1 of (El Ghaoui, 2019). I think the condition can be further relaxed to something like $\\partial F / \\partial h*$ has spectual radius less than 1. Note that this is not the same as being contractive. I think the relation should not undermine the contribution of this paper though.\n- The parameter choices of k and lambda are not well-connected with the theoretical analysis as in Theorem 1. It would help the consistency if the empirical choices of the parameters in accordance to the theory can be discussed in section 2.3. Also in section 2.3 (i) and (ii), the benefits of the better conditioning of A matrix is not well discussed. It only start to make sense to me after reading the experiment section for a while. I sugguest it helps the accuracy of the model although resulting in less accurate gradient.. Right?\n- Figure 3 (b) could use other colors for UPG and NPG to avoid visual confusion.\n- The paper could also include speed comparison on just the backward steps as I suggest it would show a much more visible improvement. Because the method does not touch the forward pass which takes roughly the same amount of time to do as the backward, the overall improvement cannot exceed approximately 2x, making it appear less impressive.\n- The proposed method should be general to all implicit models. The authors currently fail to show this by only offering experiments regarding computer vision tasks. I think the impact from generalization capability of the method would be better demonstrated if the author can include experiments on tasks in the final version, for example on natural language processing (Bai, 2019) and/or graph neural network (Gu, 2020) experiments whose code is available online.\n- I hope the authors be more generous in sharing the limitations of the method as suggested by NeurIPS 2021 board (see checklist 1b).\n\nRef:\n\nBai, S., Kolter, J. Z., & Koltun, V. (2019). Deep equilibrium models. Advances in Neural Information Processing Systems 32. https://papers.nips.cc/paper/2019/hash/01386bd6d8e091c2ab4c7c7de644d37b-Abstract.html\n\nEl Ghaoui, L., Gu, F., Travacca, B., Askari, A., & Tsai, A. Y. (2019). Implicit deep learning. arXiv preprint arXiv:1908.06315, 2. https://arxiv.org/abs/1908.06315\n\nRevay, M., Wang, R., & Manchester, I. R. (2020). Lipschitz Bounded Equilibrium Networks. arXiv preprint arXiv:2010.01732. https://arxiv.org/abs/2010.01732\n\nGu, F., Chang, H., Zhu, W., Sojoudi, S., El Ghaoui, L. (2020). Implicit graph neural networks. Advances in Neural Information Processing Systems 33, 11984--11995. https://papers.nips.cc/paper/2020/hash/8b5c8441a8ff8e151b191c53c1842a38-Abstract.html\n\nWang, T., Zhang, X., & Sun, J. (2020). Implicit feature pyramid network for object detection. arXiv preprint arXiv:2012.13563. https://arxiv.org/abs/2012.13563\n\n\n---------------------------------------\nPost Rebuttal:\nI appreciate the authors' detailed responses to my questions. I am glad that the authors find my comments useful and will include them in the final draft. I have no further concerns about the paper being accepted by NeurIPS. Thus, I will update my score to 7.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "YdnfGlJUEUQ", "reviewer": "Reviewer_f2Rd", "summary": "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$. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n1. The paper is well-written and organized very well. Every step is motivated properly and the remarks provide good intuition for the theory.\n2. The empirical experimentation is reasonably thorough, in particular lines (174-179) provide a good roadmap for section 3. \n3. The speed improvement is significant. The results on ImageNet show that this algorithm scales well with minimal performance degradation (75.3 -> 75.0). \n4. The convergence theory provided in section 2.4 along with the cosine alignment plots in Figure 1 are very convincing.\n5. The relevant hyperparameters of the method are ablated appropriately (Figure 3). \n6. I like the takeaway that precise gradient estimates are not always required, especially in cases where a layer is a black box. I think this will be of interest to the community.\n\n### Weaknesses\n\n1. Remark 2 seems incorrect. This holds in the linear model case, however for generic multilayer networks with ReLU or sigmoid non-linearity, the NTK is not a scalar network. See [2] for details. Perhaps there's some other effect keeping the Jacobian well-conditioned with large width however.\n2. Does the Neumann Series for computing the phantom gradient always converge faster than Broyden's fixed point in the standard DEQ setup? It seems like this algorithm can be seen as trading one iterative solver for another.\n3. There are no error bars for the small scale experiment. Is training with this gradient estimate stable in terms of performance?\n\n\nOverall I think this method is quite interesting and will be of benefit to the larger NeurIPS community. My rating can be improved if the authors answer some of the questions I have posed, or clarify some misconceptions with their method.\n\n\n### Questions to the Author\n\n1. Is the Jacobian-inverse likely to be ill-conditioned in this setting, or is this mostly a motivation?\n2. I would like to see a plot of the condition in Eq. 6, perhaps for a small dataset. Is this a reasonable assumption?\n3. Figure 2 shows that cosine similarity with the exact gradient decreases with the training iterations. How does this affect training later on? Could the algorithm be improved by keeping $k$ low early in training and increasing it later?\n4. Are there cases where Broyden converges faster than the Neumann series?\n5. Do you find that these hyperparameters transfer well between tasks? Do you need to re-tune?\n\n\n[1] Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale Deep Equilibrium Models. In Neural Information Processing Systems (NeurIPS), pages 5238–5250, 2020.\n[2] Youngmin Cho, Lawrence Saul. Kernel Methods for Deep Learning. In Advances in Neural Information Processing Systems 22 (NIPS 2009)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "They do not discuss the limitations or social impact of their work. Given that this paper is rather abstract, dealing with a learning method rather than an application, I believe this to be justified. One suggestion for limitations, however, would be to try to study cases where their gradient estimator is unstable.", "review_text": "", "extra_scores": {"ethical_concerns": "This paper does not have direct ethical issues.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Many current architectures use "implicit" layers where a solver (e.g. root finder or ODE solver) gives the output. Exactly differentiating these models requires implicit differentiation in the backprop step, which is expensive. This paper introduces a novel "phantom" gradient which is easy to compute, and shows that the inner-product of this gradient with the true gradient is >0, meaning it is a descent direction. An empirical verification shows that optimization using this gradient succeeds.

Reviewers generally felt the paper was novel, well-organized, convincing, clear, and had experiments that show the method truly works. There were a few weaknesses that were identified. First, there were some technical issues in the theoretical results. From the author response, these appear to be relatively minor and specific corrections have been identified. Second, there were some concerns about the rigor of the experiments (e.g. the lack of error bars in the small-scale experiment). Third, there were some concerns about notation and related works. For these latter two concerns the authors have responded at length in the feedback, and this has improved reviewer confidence about these issues. We trust the authors will integrate this content into the main paper, which will strength it considerably.
Decision: Accept (Poster)