Joint inference and input optimization in equilibrium networks
Authors: Swaminathan Gurumurthy
Shaojie Bai
Zachary Manchester
J Zico Kolter
Venue: NeurIPS 2021 Poster (2021)
Keywords: Deep Equilibrium Models
Bi-level Optimization
Generative Models
Inverse Problems
Adversarial Training

Abstract: Many tasks in deep learning involve optimizing over the inputs to a network to minimize or maximize some objective; examples include optimization over latent spaces in a generative model to match a target image, or adversarially perturbing an input to worsen classifier performance.  Performing such optimization, however, is traditionally quite costly, as it involves a complete forward and backward pass through the network for each gradient step.  In a separate line of work, a recent thread of research has developed the deep equilibrium (DEQ) model, a class of models that foregoes traditional network depth and instead computes the output of a network by finding the fixed point of a single nonlinear layer. In this paper, we show that there is a natural synergy between these two settings. Although, naively using DEQs for these optimization problems is expensive (owing to the time needed to compute a fixed point for each gradient step), we can leverage the fact that gradient-based optimization can itself be cast as a fixed point iteration to substantially improve the overall speed. That is, we simultaneously both solve for the DEQ fixed point and optimize over network inputs, all within a single "augmented" DEQ model that jointly encodes both the original network and the optimization process.  Indeed, the procedure is fast enough that it allows us to efficiently train DEQ models for tasks traditionally relying on an "inner" optimization loop.  We demonstrate this strategy on various tasks such as training generative models while optimizing over latent codes, training models for inverse problems like denoising and inpainting, adversarial training and gradient based meta-learning.

Joint inference and input optimization in equilibrium networks
Swaminathan Gurumurthy' Cameeii Mellon University
Shaojie Bai Carnegie Mellon University
J. Zico Kolter
Zachary Manchester Carnegie Mellon University
Carnegie Mellon University Bosch Center for Al
Abstract
Many tasks in deep learning involve optimizing over the imputs to a network to minimize or maximize some objective; examples include optimization over latent spaces in generative model to match a target image, adversarially perturbing an input worsen classifier performance. Performing such optimization, however, is traditionally quite costly, it involves complete forward and backward pass through the network for each gradient step. In a separate line of work, a recent thread of research has developed the deep equilibrium (DEQ) model, a class of models that foregoes traditional network depth and instead computes the output a network by finding the fixed point of a single nonlinear layer. In this paper, show that there is natural synergy between these two settings, Although, naively using DEQs for these optimization problems is expensive (owing to the time needed compute a fixed point for each gradient step), we can leverage the fact that gradient- based optimization improve the overall speed. That is, we simultaneously both solve for the DEQ fixed point and optimize over network inputs, all within single "augmented" DEQ model that jointly encodes both the original network and the optimization process. Indeed, the procedure is fast enough that allows
itself be cast
fixed point iteration to substantially
efficiently train DEQ models for tasks traditionally relying on an "inner" optimization loop. We demonstrate this strategy on various tasks such as training generative models while optimizing over latent codes, training models for inverse problems like denoising and inpainting, adversarial training and gradient based meta-learning.
Introduction
Many settings in deep learning involve optimization over the inputs network minimize some desired loss. For example, for "generator" network G Z + X' that maps from latent space Z to an observed space X, may be desirable to find latent vector z E that most closely produces some target output æ E X' by solving the optimization problem (e.g. [10 13)
minimize 2E2
||
Go
As another example, constructing adversarial examples for classifiers 28 optimizating over a perturbation to a given input; i.c., given a classifier network X + y, task loss
typically involves
"Correspondence to: Swaminathan Gurumurthy wwwnnmmreedddeew.mmoooo
Code at
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

e: + RR, and a sample E X, we want to solve
maximize [olls
l(g(z + +)).
More generally, wide range of inverse problems 10 and other auxiliary tasks 22 in deep
learning can also be formulated in such manner.
Orthogonal to this line of work, a recent trend has focused on the use of an implicit layer within deep networks to avoid traditional depth. For instance, Bai et al. [3] introduced deep equilibrium models (DEQs) which instead treat the network repeated applications of single layer and compute the output of the network as a solution to an equilibrium-findin sequence of non-linear layer operations. Bai et al. [5] and subsequent work [6] have shown that DEQs can achieve results competitive with traditional deep networks for many realistic tasks.
problem instead of simply specifying
In this work, we highlight the benefit of using these implicit models in the context of input optimization routines. Specifically, because optimization over inputs itself is typically done via iterative method (e.g., gradient descent), we can combine this optimization fixed-point iteration with the forward DEQ fixed point iteration all within single "augmented"" DEQ model that simalaaneussy performs forward model inference as well as optimization over the inputs. This enables the models to more quickly perform both the inference and optimization procedures, and the resulting speedups further allow us to train networks that use such "bi-level" fixed point passes. In addition, we also show close connection between our proposed approach and the primal-dual methods for constrained optimization.
We illustrate our methods on 4 tasks that span across different domains and problems: 1) training DEQ- generative models while optimizing over latent codes; 2) training models for inverse problems such as denoising and inpainting: 3) adversarial training of implicit models; and 4) gradient- based meta-learning. We show that in all cases, performing this simultaneous optimization and forward inference accelerates the process over more naive inner/outer optimization approach. For instance, using the combined approach leads to a 3.5-9x speedup for generative DEQ networks, 3x speedup in adverarial training of DEQ networks and 2.5-3x speedup for gradient based meta-learning. In total, we believe this work points to a variety of new potential applications for optimization with implicit models.
Related Work
Implicit layers. Layers with implicitly defined depth have gained tremendous popularity in recent
y4ars/46) 19 (29]. Rather than static computation graph, these layers define a condition on the output that the model must satisfy, which can represent "infinite" depth, be directly differentiated through via the implicit function theorem 47), and are memory-efficieee to train. Some recent examples of
implicit layers include optimization layers [16] I1 deep equilibrium models[5] 6 68 40 52). neural
ordinary differential equations (ODEs) 14 18 61), logical structure learning [67] J. and continuous
generative models [30].
In particular, deep equilibrium models (DEQs) [5] define the output of the model the fixed point of repeated applications of a layer. They compute this using black-box root-finding methods (5) or accelerated fixed-point iterations 36 (e.g., Broyden's 's method (TT). ). In this work, we propose an efficient approach to perform input optimization with the DEQ by simultaneously optimizing over the inputs and solving the forward fixed point of an equilibrium model a joint, augmented system. As related work, Jeon et al. [36] introduce fixed point iteration networks that generalize DEQs to repeated applications of gradient descent over variables. However, they don't address the specific formulation presented in this paper, which has a number of practical use cases (e.g., adversarial training). Lu et al. [52] proposes an implicit version of normalizing flows by formulating a joint root-finde findinniig problem that defines an invertible function between the input x and output 2". Perhaps the most relevant approach to our work is Gilton et al. (26.. which specifically formulates inverse imaging problems as DEQ model. In contrast, our approach focuses on solving input optimization problems where the network of interest already DEQ, and thus the combined optimization and forward inference task leads to substantially different set of update equations and tradcoffs.

---

Input optimization in deep learning, Many problems in deep learning can be framed as optimizing over the inputs to minimize some objective. Some canonical examples of this include finding
learning generative models [ 72), meta-learning [58 22 74 etc. For most of these examples, input optimization is typically
adversarial examples 533 45 solving inverse problems I0,
13 56
done using gradient descent on the input, i.e., we feed the input through the network and compute some loss, which we minimize by optimizing over the input with gradient descent. While some of these problems might not require differentiating through the entire optimization process, many do (introduced below), and can further slow down training and impose massive memory requirements.
Input optimization has recently been applied to train generative models. Zadeh et al. 72]. Bojanowski et al. [9] proposed train generator networks by jointly optimizing the parameters and the latent variables corresponding to each example. Similarly, optimizing a latent variable to make the corresponding output match a target image is common in decoder-only models like GANs to get correspondences [10, (39] and has been found useful to stabilize GAN training 771. However, in all of these cases, the input is optimized for just few (mostly 1) iterations. In this work, we present generative model, where we optimize and find the optimal latent code for each image at each training step. Additionally, Bora et al. [10] Chang et al. [3] showed that we can take a pretrained generative model and use it as a prior to solve for the likely solutions to inverse problems by optimizing on the input space of the generative model (i.e., unsupervised inverse problem solving). Furthermore. Diamond et al. [I5]. Gilton et al. [25]. Gregor and LeCun [I have shown that networks can also be trained to solve specific inverse problems by effectively unrolling the optimization procedure and iteratively updating the input. We demonstrate our approach in the unsupervised setting as in Bora et al. [T], Chang et al. [13), but also show flexible extension of our framework to train implicit models for supervised inverse problem solving.
This manifests as optimizing an objective that incentivices an incorrect prediction by the classifier, while constraining the input be within a bounded region of the original input. Many attempts have been made on the defense side [37] 65 69). The most successful strategy thus far has been adversarial training with projected gradient descent (PGD) adversary (53) which involves training the network on the adversarial examples computed using PGD online during training. We show that our joint optimization approach can be easily applied to this setting, allowing us train implicit models to perform competitively with PGD in guaranteeing adversarial robustness, but at much faster speeds.
Another crucial application of input optimization is to find adversarial examples 64
28
While the examples above were illustrated with non non-convex networks, attempts have also been made to design networks whose output is a convex function of the input [2]. This allows one use more sophisticated optimization algorithms, but usually at heavy cost of model capacity. They have been demonstrated to work in a variety of problems including multi-label prediction, image completion 2 learning stable dynamical systems 44 and optimal transport mappings [54] MPC [12, etc.
3 Joint inference and input optimization in DEQs
Here we present our main methodological contribution, which sets up an augmented DEQ that jointly performs inference and input optimization over existing DEQ model. We first define the base DEQ model, and then illustrate a joint approach that simultaneously finds it's forward fixed point and optimizes over its inputs. We discuss several methodological details and extensions.
3.1 Preliminaries: DEQ-based models
To begin with, we recall the deep equilibrium model setting from Bai et al. [5), but with the notation slightly adapted to better align with its usage in this paper. Specifically, we consider an input- spput-injected layer fa: 2 x X' - Z where Z denotes the hidden state of the network, X' denotes the input space, and 0 denotes the parameters of the layer. Given an input & E X, computing the forward pass in DEQ model involves finding a fixed point = (x) E Z, such that
z (x) =
(3)
which (under proper stability conditions) corresponds to the "infinite depth" limit of repeatedly applying the fe function. We emphasize that under this setting, we can effectively think of 2 itself
3

---

as the implicitly defined network (which thus is also parameterized by f), and one can differentiate through this "network" via the implicit function theorem g 47
The fixed point of a DEQ could be computed via the simple forward iteration
z² : f0(2,2)
starting at some artibrary initial value of & (typically 0). However, in practice DEQ models will typically compute this fixed point not simply by iterating the function for but by using more accelerated root-finding or fixed-point approach such as Broyden's 's method [11] or Anderson acceleration [6). Further, although little be said about e.g., the existence uniqueness of these fixed points in general (though there do exist restrictive settings where this is possible [68] [5] (3]). in practice wide suite of techniques have been used to ensure that such fixed points exist, can be found using relatively few function evaluations, and are able to competitively model large-scale tasks
3.2 Joint inference and input optimization
Now we consider the setting of performing input optimization for such a DEQ model. Specifically, consider the task of attempting optimize the input E X to minimize some loss f: Z x R+.
minimize (x)(
(5)
ZEX
where y E y represents the data point. To solve this, we typically perform such an optimization via e.g., gradient descent, which repeats the update
= x (
(0l(2(()).)) dx
(6)
until convergence, where we term 2" alone to denote the fixed output of the network =ă (i.c., just as a fixed output rather than a function). Using the chain rule and the implicit function theorem, we can further expand update (6) using the following analytical expression of the gradient:
Ozg (x) (1(=)()). af(=*,9) Ōr
Ol(=".9) 0z
(==Of(("")))) -T Ofo(2",
Oz*
dx
2
Thinking about =j an implicit function of permits combine the fixed-point equation in Eq. (on 2) with this input optimization update (on x)), thus performing joint forward update:
f0 x))
@fa(2.)) dxe
af.(2.2) a8(=.v) 85
(8)
x a
ds
should be apparent that, if both iterates converge, then they have converged to a simultaneous fixed point =* and optimal x* value for the optimization problem (5). However, simply performing this update can still be inefficient, because computing the inverse Jacobian Eq. 7 expensive and typically computed via an iterative update namely, we would first compute the variable
(I 8fe(s.2)
8f(2.w)
via the following iteration (i.e., a Richardson iteration 660
02
02
(0f0l2.0) Oz
T 14 +
(l((z,)) 9)
1r
Oz
Therefore, to efficiently solve the joint inference and input optimization problem, we propose combining all three iterative procedures into the update
+ a
f(2,2)
(So((..))
38(x.))
+
(10)
x @felz.z)
1
Like Eq. (8), if this joint process converges to a fixed point, then it corresponds to simultaneous
optimum of both the inference and optimization processes. Such formulation especially appealing, as the iteration (10) is äself just an augmented DEQ network (y) (i.e., with input injection y) whose forward pass optimizes on a joint inference-opiimizatinn space = (x, μ,, z). Moreover, we standard techniques to differentiate through this process, though there also optimizations

---

Target
x
2
Augmented
Forward
Backward
2041 (0+
F(z
x
GD update
x
[0]
[++1]
Concelanated Joint pointvariable
Targelacting Inputinjectin
Inputinjection
Conventional Input Optimization in DEQs
Joint updates using "Augmented" DEQ
(Vanilla Gradient Descent)
Figure I: Left: Performing each gradient update on DEQ inputs requires fixed point computation in the forward and backward pass. Right: Solving the fixed points simultaneously as an "augmented" DEQ where the targets y act input and the function F represents the joint fixed point updates in Eq.
can apply in several settings that discuss below. This is in contrast prior works where fe is an explicit deep neural network, where the model forward-pas pasarddsssssssaa and optimization processes are disentangled and have to dealt with separately. We illustrated this in Figure I where the figure on the left shows the input optimization naively performed using gradient descent in DEQs v/s the figure on the right which shows the joint updates performed using the augmented DEQ network.
As final note, we mention that, practice, just as the gradient-descent update has a step size o, often beneficial to add similar "damping" step sizes the other updates as well. This leads to the full iteration over the augmented DEQ
(1 (1-ag)2+ azf(2,x)
(1 an) + au (( (Dfolem))
#
=
1
82
(II)
2
(..(2.))) dx
x Ox
Finally, in order to speed up convergence, as common in DEQ models, we apply more involved fixed point solver, such as Anderson acceleration, top of this naive iteration. We analyze the effect of these different root-finding approaches in the Appendix.
Notes on Convergence Our treatment of the above system augmented DEQ allows to borrow results from 66 to ensure convergence of the fixed point iteration. Specifically, we assume the joint Jacobian of the fixed point iterations we describe are strongly monotone with smoothness parameter m and Lipschitz constant L, then by standard arguments (see c.g., Section 5.1 of (62]), the fixed point iteration with step size m/L² will converge. Note that these substantially weaker rates and constants than required for typical gradient descent or the minimization of locally convex function because the coupling between the three fixed point iterations introduce cross-terms in the joint Jacobian.
Due to these cross-terms, going from the strong monotonicity assumption on the joint fixed point iterations to specific assumptions on fe and f is less straightforward. However, empirically, observed that as long as the step sizes a' were kept reasonably small and the functions fe and & were designed appropriately (e.g & respecting notions of local convexity and fe with Jacobian eigenvalues less than 1, etc.) the fixed point iterations converged reliably.
3.3 Iterpretation as primal-dual optimization
While the problem above was introduced as an input-optimization problem, its formulation as joint optimization problem in the augmented DEQ system (10) can also be viewed as constrained

---

optimization problem where the DEQ fixed-point conditions are treated as constraints,
minimize ((z,y) subject =f (2,
(12)
2,
which yield the Lagrangian
minimize maximize L(x, 2,) = ((2.8) + 1 (Jo(z,0)
(13)
P
and the corresponding KKT conditions
fe fa(z, 2) 2
f0 x)) x-)))2 -2
06(2,2,4)
De(2.v) 82 + @folz.0)
82 I)
E IR 2n
(14)
88
0C(2.2.14
μ (f.(2.0)
82
where μl are the dual variables corresponding to the equality constraints. Rearranging the terms in the KKT conditions of the above problem, introducing the step size parameters a's and treating it as fixed point iteration gives us the updates in Eq. [II. Indeed, performing such iterations is a variation of the classical primal-dual gradient method for solving equality optimization problems
20 34 17
3.4
Outer Optimization (Backward Pass)
notable advantage of formulating the entire joint inference and input optimization problem as an augmented DEQ (v) is that it allows us to abstract away the detailed function of and simply train parameters 0 of this joint process as an outer optimization problem:
minimize gouter (vg (v)
(15)
where gouter : 2 2
R
Given the solutions x², 22, the inner problem, computing updates in the outer optimization problem is equivalent to the backward pass of the augmented DEQ and correspondingly is optimized using standard stochastic gradient optimizers like Adam 42). Thus, as with any other DEQ model, we assume the inner problem was solved to fixed point, and apply the implicit function theorem to
compute the gradients w.r.t the augmented system (T].. This gives us a constant memory backward pass which is invariant to the underlying optimizer used to solve the inner problem.
Ocouter (vā. v)
Ocouter (y, v²) Oug
@eouter (y,
(oKa(v")
@Ko(v*)
(16)
00
Ou*
08
Ou*
Ou"
00
where [x., 2, 14] and Ko(v) = 0 represents the KKT conditions from (1.. As with the original
DEQ, instead of computing the above expression explicitly, we first solve the following linear system
to compute and then substitute it back in the equation above obtain the full gradient,
u Hgouter (y. v²)
(oKo(w) Ov
Ogouter (y, 2*) +
(OKo(e) ==0 (17)
Ou+
Ov
Ov
Although we can train any joint DEQ in this manner, doing so in practice (c.g., via automatic differentiation), will require double backpropagation, because the definition of Ko(u) above already includes vector-Jacobian products, and this expression will require differentiating again. However, in the case that pouter is the same substantial simplification of the outer optimization gradient. These cases are indeed quite common, as we may want e.g., train the parameters of generative model minimize the same reconstruction error that we attempt to optimize via the latent variable; or in the case of adversarial examples, the inner adversarial optimization is precisely the negation of the outer objective.
(or in fact where it is the negation of (), then there exists
In these cases, we have that
gouter (v, v% vx)) Ely. 2 (v))
(18)
so we have that dcouter (y, vw)) Of(y, x*) (1 000(3", Oz*
000(2", (4+*) 05o(2 20
xx*)
(19)
00 In other words, we can compute the exact needed derivatives with respect to 0 by simply re-using the converged solution v*, without the need double backpropagate through the KKT system. The same considerations, but just negative, apply to the case where gouter -l.
Oz

---

4 Experiments
As our approach provides a generic framework for joint modeling of an implicit network's forward dynamics and the "inner" optimization over the input space, we demonstrate its effectiveness and generalizability on 4 different types of problems that are popular areas of research in machine learning:
generative modeling [43 27 inverse problems [T] 25 31]. adversarial training 669 [33 and gradient
58 (results for the latter are in the appendix). In all cases, we show that our joint inference and input optimization (JIIO) provides significant speedups over projected gradient descent applied to DEQ models and that the models trained using JIIO achieve results competitive with standard baselines. In all of our experiments, the design of our model layer fa follows from the prior work on multiscale deep equilibrium (MDEQ) models [6] that have been applied on large-scar scale computer vision tasks, and where we replace all occurrences of batch normalization [35] with group normalization [70] in order to ensure the inner optimization can be done independently for each instance in the batch. We elaborate on the details of the choice of other hyperparameters and design decisions of our model (such the damping parameters in the update step [1 ), as well as that of the datasets for each task in the Appendix.
based meta-learning 22
We introduce below each problem instantiation, how they fit into our methodology described in Sec. S.c. and the result of applying the JIIO framework compared to the alternative methods trained in similar settings. Overall, our results provide strong evidence of benefits of performing joint optimizations on implicit models, thus opening new opportunities for future research in this direction.
4.1 Generative Modeling
We study the application of JIIO to learning decoder-only generative models that compute the latent representations by dircetly minimizing the reconstruction loss 72 9]: i.c., given decoder network D, the latent representation x of a sample y (e.g., an image) is x = minzex (xr) Bl 2 Moreover, instead of placing explicit regularizations on the latent space X' (as in VAEs), we follow (24 to directly train the decoder for reconstruction (and then after training, we fit the resulting latents using simple density model, post-hoe, for sampling). Formally, given sample data y1..... . V/. (e.g., images), the generative model we study takes the following form:
minimize Σ llue ho(2t))1
(20)
subject to 22,2
11ya - ho()))1 i= i..,..,, 2,0:a=f((2.a)
argmin
where he is final output layer that transform the activations 2* to the target dimensionality. We train the MDEQ-based fo with the JIIO framework on standard 64 cropped images from CelebA dataset, which consists of 202,599 images. We use the standard train-val-test split as used in Liu et al. [51 and train the model for 50k training steps. We use Fréchet inception distance (FID) 3 to measure the quality of the sampling and test-time reconstruction of the implicit model trained with JIIO and compare with the other standard baselines such as VAEs [43]. The results are shown in Table I JIO-MDEQ refers to the MDEQ model trained using our setup with 40 JIIO iterations in the inner loop during training (and tested with 100 iterations). MDEQ-VAL VAE refers to an equivalent MDEQ model but with encoder and decoder trained as a VAE. We observe that our model's s generation quality is competitive with, or better than, each of these encoder-decoder based approaches. Moreover, with the joint optimization proposed, JIIO-MDEQ achieves the best reconstruction quality.
We additionally apply JIIO on pre-trained MDEQ-VAIE (i.e., train an MDEQ-based VAE usual on optimizing ELBO 43). and take the decoder out) for test-time image reconstruction. The result (shown in Table T suggests that the reconstructions obtained result better even than the original MDEQ-VAL In other words, JIIO can be used with general implicit-mode-basee decoders at test time even if the decoder wasn't trained with JIIO.
"Thi notation differs from the "standard"' notation of latent variable models (where the latent variable is typically denoted by 2). However, because 2 all have standard meanings in setting above, we change from the common notation here to be more consistent with the remainder of this paper.
7

---

Model
Generation
Reconstruction
VAE
43
48.12
39.12
RAE 24
40.96
36.01
MDEQ-VAL
57.15
45.81
MDEQ- VAE (w/ JIIO)
42.36
JIIO-MDEQ
46.82
32.52
Table I: Comparison of FID scores attained by standard generative models with our method, which performs joint optimization. We use 40 solver iterations (for the augmented DEQ) to train the JIIO model reported in this table.
Figure 2: Samples generated with JIIO on a small MDEQ network.
One of the key advantages presented by JIIO is the relative speed of optimization over simply running gradient descent (or adaptive variants like Adam 4T). Table 2 shows our timing results for one optimization run on a single example for various models (averaged over 200 examples). We observe that performing 40 iterations of projected Adam takes more than 9x the time taken by 40 iterations of JIIO, which we used during training and more than 3.5x the time taken by 100 iterations of JIIO which we use for reconstructions at test time (e-g., to produce the results in Table [ though both of them lead to similar levels of reconstruction loss). Fig shows the reconstruction loss evolves across single optimization run for an MDEQ model trained with JIIO. This again clearly shows that JIIO converges vastly faster (in terms of wall-clock time) than if we handle the inner optimization separately as in prior works, demonstrating the advantage of joint optimization. However, it's s interesting to note that JIIO optimization seems somewhat unstable (see Fig. 3 and fluctuates more as well. This seems to be an artifact of the specific accelcration scheme we use (sce more details in Appendix A.2)
4.2 Inverse Problems
We also extend the setup mentioned in section 4.1 directly to inverse problems. These problems, specifically, can be approached as either an unsupervised or a supervised learning problem, which we discuss separately in this section. To demonstrate how JIIO can be applied, we will be using example inverse problem tasks, which was extensively
image inpainting and image denoising studied in prior works like Chang et al. [33, Gilton et al. [26]. For the inpainting task, we randomly mask 20x20 window from the image and train the model to adequately fill the missing pixels based on the surrounding image context. For the image denoising tasks, we add random gaussian noise ~ (O, with 0.2 and 0.4, respectively, to all pixels in the image, and train the model to recover the original image. We use the same datasets and train-test setups as in the generative modeling experiments in Sec. 4.1
0.60
|II (40 [ters)
Model
time taken (ms)
Adam (4C iters)
0.40
[IO] (300 nees)
PGD 20 iters
4360
Adam (100
J1IO 80 iters
1401
Table 2: Time taken to compute adversarial
0.22
0.00
example of MDEQ model on MNIST
0.01
Model
time taken (ms)
0.52
Adam: 40 iters
7659
J11O 40 iters
862
0.01
30000
JIIO 100 iters
2156
(ms)
Figure 3: Cost changing with time for Adam v/s JIIO optimization. Tested models trained with 40 and 100 JIIO iterations respectively
Table 3: Time taken to perform JIIO optimization v/s Adam in the generative modeling/inverse problem experiments

---

Task
Model
Inpainting
Denoising
0.2)
Denoising (o 0.4)
AE
17.9
18.72
18.32
Supervised
MDEQ-AE
17.06
18.58
18.49
BI1O-MDE-110
16.90
18.22
17.89
VAE (Adam) [10]
15.34
15.31
15.24
Unsupervised
MDEQ-VAE (Adam)
16.62
16.96
16.87
11O-MDEQ-40
15.88
17.08
16.03
JIO--DEE--100
15.87
17.86
17.55
Table 4:
Comparison of Median PSNR values for supervised and unsupervised inverse problem solving approaches. The top 3 rows show models that are trained for the specific inverse problem and the latter show pre-trained generative models re-purposed for solving inverse problems
4.2.1 Unsupervised inverse problem solving
Bora et al. [0). Chang et al. [3 have showed that we can solve most inverse problems by taking pre-trained generative model and using that prior solve for the likely solutions the inverse problems by optimizing on the input space of the generative model. Specifically, given "generator" network X + V, mapping from the latent space X' an observed space y, that models the data generating distribution, they show that one can solve any inverse problem by optimizing the following objective:
minimize ZEX
|| AG(x)
(21)
where y = Ay E Y represents the corrupted data point, y E y is the uncorrupted data and A y y denotes the measurement matrix that defines the specific type of inverse problem that try to solve (e.g., fg., for image inpainting, it would be mask with the missing regions filled in with zeros. For deblurring, it would be convolution with gaussian blur operator etc.). They call unsupervised inverse problem solving. Likewise, we can use the pre- prettrinee generator from section 4.1 to solve most inverse problems by simply solving slightly modified version of the inner problem in (2
minimize a,2:a= fe(.aa)
|| (21
(22)
In table 4 the unsupervised results for VAE and MDEQ-VAL VAE generators are obtained by optimizing (2T) using Adam for 40 iterations, while for the JIIO trained models, we optimize (22) with 100 JIIO iterations. JIIO-MDEQ
and JIO-MDEC 100 refer JIOO-MDEQ models trained with 40 and 100 inner-loop iterations respectively. The results in Table 4 show that on all problems, JIIO trained generators produce results comparable the VAE and MDEQ VAE generators. Moreover, as shown in section 4.1 JIIO also converges much faster than Adam applied to a MDEQ VAE generator.
4.2.2 Supervised Inverse problem solving
While the unsupervised inverse problem solving works reasonably well, we can also learn models to solve specific inverse problems to obtain better performance. Specifically, given uncorrupted data y1..... Yo and the measurement matrix A, we can train a network Gg y mapping from the corrupted sample D, Ay to the uncorrupted sample y4 by minimizing:
minimize Σ 11 (Ay))))
(23)
Now, instead of modeling Gg as an explicit network, we could also model it as solution to the inverse problem in (22) and the resulting parameters can be trained follows:
minimize Σ llus h(2)))1
(24)
subject to 4/,2 23 argmin
|| - Ah(=)) i= 1,..,,..7
2.2:2=fo(2.2)
As shown in Table this yields models competitive with their autoencoder based counterparts, while being better than all the unsupervised approaches. Each of the baseline models in the supervised section of Table 4 are trained by simply optimizing (23) with the corresponding model replacing Go- However, note that given the models here are trained on specific inverse problems, one would have to train new model for each new problem opposed to the unsupervised approach.

---

Datasets
Train (1) Test (-)
Clean
PGD
JI1
Clean
99.45 0.03
80.1 = 1.87
0.88:4.77
MNIST
PGD
99.18 0.03
06.53
95.74
JI1O
99.32 = 0.09
95.74 0.22
96.63 0.58
Clean
78.47 0.94
2.38 = 0.41
0.41
3.71 = 4.01
CIFAR
PGD
54.91 1.01
37.4 = 0.26
36.17 = 0.55
J110
= 0.82
37.31 = 0.67
37.77 = 0.84
Table Comparison of adversarial training approaches on 1.2 norm perturbations with e 1. The rows represent the training procedure and the columns represent the testing procedure
4.3 Adversarial Training
Although the previous two tasks were based on image generation, we note that our approach can be used more generally for input optimization in DEQs and illustrate that by applying it to & adversarial training on DEQ-based classification models. Specifically, given inputs Ti Vs! i
1,..., n, adversarial training seeks to optimize the objective
minimize
max ((ho(1; + 6,), 3/) óllas
(25)
We apply our setting to a DEQ-based classifier ha(x) = he (2(2((z)) where 2" = fo( xx) In this setting, embed the iterative optimization over 5 (with projected gradient descent) into the augmented DEQ. and write the problem as
minimize Σ E((h((z))))) 1=1
(26)
subject to 2
argmin
-E(ho(2), i i= 1,...,7
2.8. fo
Specifically train MDEQ models CIFARIO (48) and MNIST (50) using adversarial training against L2 attacks with e 1 for 20 epochs and 10 epochs respectively using the standard train-val- test splits. Table 5 shows the robust and clean accuracy of models trained using PGD adversarial training, JIIO adversarial training and vanilla training. We find that models trained using JIIO have comparable robust and clean accuracy both datasets. Furthermore, when tested models trained without adversarial training, we observe that JIIO serves as a comparable attack method to PGD. Table 4.2 shows the time taken find the adversarial example for single image of MNIST using 20 iterations of PGD and 80 iterations of JIIO. We again observe more than 3x speedups when using JIIO over using PGD while obtaining competitive robust accuracy. However, note that, unlike previous experiments in generative modeling/inverse problems, performing adversarial training with truncated JIIO iterations would lead to significant reduction in robust performance due the adversarial setting.
5 Concluding remarks
We present novel optimization procedure for jointly optimizing over the input and the forward fixed point in DEQ models and show that, for the same class of models, it is 3 9x faster than performing vanilla gradient descent Adam the inputs. We also apply this approach to 3 different settings to show it's effectiveness: training generative models, solving inverse problems, adversarial training of DEQs and gradient based meta-learning. In the process, also introduce entirely decoder only generative model that performs competitively with it's autoencoder based counterparts.
type of
Despite these features, we note that there is substantial room for future work in these directions. Notably, despite the fact that the augmented joint inference and input optimization DEQ can embed both processes in "single" DEQ model, in practice these joint models take substantially more iterations to converge as well (often in the range of 50-100) than traditional DEQs (often 10-20 iterations). and correspondingly often use larger memory caches within methods like Anderson acceleration. Thus, while we make the statement that these augmented DEQs are "just" another DEQ, this relative difficulty in finding fixed point likely adds challenges to training the underlying models. Thus, despite ongoing work improving the inference time of typical DEQ models, there substantial room for improvement here in making these joint models truly efficient.
10

---

Acknowledgements
Swaminathan Gurumurthy and Shaojie Bai are supported by grant from the Bosch Center for Artificial Intelligence. We further thank Zhichun Huang and Taylor Howell for their help with some implementations and for various brainstorming sessions and feedback throughout the course of this project. We would also like to thank Vladlen Koltun for brainstorming ideas especially in the initial stages of the project.
References
[1] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization layer in neural networks. In International Conference on Machine Learning (ICML), pages 136-145, 2017.
[2] B. Amos, L. Xu, and J. Kolter. Input convex neural networks. In D. Precup and Y. W. Tch, editors, Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 146-155, PMLR, 06-11 Aug 2017. URL http://proceedin ngs m.l...r.......m prwss/ v 70/ amos 1 7b.7b.4 17b.html
[3] B. Amos, I. Jimenez, J. Sacks, B. Boots, and J. 7. Kolter. Differentiable mpc for end-to-end planning and control. S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL httpss: http:///aeaaaiasoo/o/o/oo/////// file/ba6d84 3eb4 www.eeeeeeeeeeeeeeeeeeeee 2al.aaaaee..11112ceeceecce211333 Pap099-Papr.... p
[4] D. G. Anderson. Iterative procedures for nonlinear integral equations. Journal of the ACM (JACM), 12(4):547-560, 1965.
[5] S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32:690-7011 1001 2019.
[6] S. Bai, V. Koltun, and Z. Kolter. Multiscale deep equilibrium models. In Advances in Neural Information Processing Systems, volume 33, 2020. URL httphs://gtthub...... lccus lab/mdeq
[7] S. Bai, V. Koltun, and Z. Kolter. Stabilizing equilibrium models by jacobian regularization, In
M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 554-565 PMLR, html
18 24 Jul 2021. Utthtttss://prooeeee
ngs ml wttpsmmr..rr.sssssseeseeem
[8] K. Binmore and J. Davies. Calculus: concepts and methods. Cambridge University Press, 2001.
[9] P. Bojanowski, A. Joulin, D. Lopez-Paz, and A. Szlam. Optimizing the latent space of generative networks. arXiv preprint arXiv: 707.0007000007 7707.05776, 2017.
[10] A. Bora, A. Jalal, E. Price, and A. G. Dimakis. Compressed sensing using generative models.
In International Conference on Machine Learning (ICML), pages 537-546, 2017.
[11] C. G. Broyden. A class of methods for solving nonlinear simultaneous equations. Mathematics of Computation, 1965.
[12] F. Bünning, A. Schalbetter, A. Aboudonia, M. H. de Badyn, P. Heer, and J. Lygeros. Input convex neural networks for building mpc. ArXiv, abs/201 1.1227, 2020.
[13] R. Chang, C.-L. Li, Póczos, B. Vijaya Kumar, and A. C. Sankaranarayanan. Sannaaaaaaaaaann One network to solve them all - solving linear inverse problems using deep projection models. In 2017 IEEE International Conference on Computer Vision (ICCV), pages 5889-5898, 2017. doi: 10.1109//CCV22117...77
[14] R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential Bengio, H. Wallach, Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neurul Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https: /proceed igg. neur wwtpseneww.rppppeeeeeeeeeee wtt.cccppppe//11111
equations.
Paper. pdf

---

[15] S. Diamond, V. Sitzmann, F. Heide. and G. Wetzstein. Unrolled optimization with deep priors. arXiv preprint arXiv: 1705000005. 2017.
[16] J. Djolonga and A. Krause. Differentiable learning of submodular models. In I. Guyon, U. V. Luxburg. Bengio. Н. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http::::/potoeeeee///ssss/ 11le/192fc04 4 74df fea 441
neur i iss. www.scccppppeeee... cc/pappe//017 pdf
[17] S. Du and Hu. Linear convergence the primal-dual gradient method for convex-coneave saddle point problems without strong convexity. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 196-205. PMLR, 2019.
(18) E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural odes. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpps://proceedngg wwwpiiinnnnnnn.inssooeeemm newrips wctpapee//199//EiI. Paper. pddf
[19] L. El Ghaoui, F. Gu, B. Travacca, A. Askari, and A. Y. Tsai. Implicit deep learning. arXiv preprint 1181v:1908.06635, 2, 2019.
[20] H. C. Elman and Golub. Inexact and preconditioned uzawa algorithms for saddle point problems. SIAM Journal on Numerical Analysis, 3166):645-16611 1994.
[21] H.-r. Fang and Y. Saad. Two classes of multisecant methods for nonlinear acceleration. Numerical Linear Algebra with Applications, 16(3):197-221, 2009.
[22] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1135. PMLR, 06-11 Aug 2017. URL |
mlr. mlr.sesss/
v70/finn html
[23] S. W. Fung, H. Heaton, Q. Li, D. McKenzie, S. Osher, and W. Yin. Fixed point networks: Implicit depth models with Jacobian-fres backprop. urXiv:2103. 12803, 2021.
[24] P. Ghosh, M. S. Sajjadi, A. Vergari, M. Black, and B. Schölkopf. From variational to deterministic autoencoders. In International Conference on Learning Representations, 2020.
[25] D. Gilton, G. Ongie, and R. Willett. Neumann networks for linear inverse problems in imaging. IEEE Transactions on Computational Imaging, 6:328-343, 2020. doi: 10.1109/TC1.219. 2948732.
[26] D. Gilton, G. Ongie, and R. Willett. Deep equilibrium architectures for inverse problems in imaging, arXiv preprint arXiv:2102.078 2021.
[27] I. Goodfellow, Pouget-Abadee Abadie, M. Mirza, Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL
igss. cettoriss.ccppaeee Paa.saapp.pppp..
[28] J. Goodfellow, J. Shlens, and C. Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6572, 2014.
[29] S. Gould, R. Hartley and D. Campbell. Deep declarative networks: A new hope. arXiv: 1909.0486. 2019.
[30] W. Grathwohl, R. T. Q. Chen, J. Bettencourt, and D. Duvenaud. Scalable reversible generative models with free-form continuous dynamics. International Conference on Learning Representations, 2019. URL
net
12

---

[31] K. Gregor and Y. LeCun. Learning fast approximations of sparse coding. In Proceedings of the 27th international conference on international conference on machine learning, pages 399-406, 2010.
[32] S. Gurumurthy, S. Kumar, and K. Sycara. Mame: Conference on Robot Learning, pages 910. PMLR, 2020.
Model-agnostic meta-exploration. In
[33] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time-scale update rule converge a local nash equilibrium. arXiv preprint 2017.
08500000
[34] M. Hong, M. Razaviyayn, and J. Lee. Gradient primal-dual algorithm converges to second-order stationary solution for nonconvex distributed optimization over networks. In International Conference Machine Learning, pages 2009-2018. PMLR, 2018.
[35] loffe and Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning (ICML), 2015.
[36] Y. Jeon, M. Lee, and J. Y. Choi. Differentiable forward and backward fixed-point iteration layers. IEEE Access, 9: 13883-11839,, 2021.
[37] H. Kannan, A. Kurakin, and I. Goodfellow. Adversarial logit pairing, ArXiv, abs/ 1803.06373, 2018.
[38] T. Karras, S. Laine, and T. Aila. A style-based generator architecture for generative adversarial networks, 2019 iece. In CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 4396-4405, 2018.
[39] T. Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8110-8119, 2020.
[40] K. Kawaguchi. On the theory of implicit deep learning: Global convergence with implicit layers. In International Conference on Learning Representations (ICLR), 2021.
[41] D. Kingma and J. Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
[42] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In Y. Bengio and Y. LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA May 7-9, 2015, Conference Track Proceedings, 2015. URL http: farxiv.orgy abs 1412.64266 6980
[43] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint 1r1iv: 2013.
[44] J. Z. Kolter and G. Manek. Learning stable deep dynamics models. In NeurlPS, 2019.
[45] Kolter and Wong. Provable defenses against adversarial examples via the convex outer adversarial polytope. In ICML, 2018.
[46] J. Z. Kolter, D. Duvenaud, and M. Johnson. Deep implicit layers neural odes, deep equilibirum models, and beyond. 2020. URL Mttttttt/mmppiiee http///mmiiiittyaaeeetttttteeeee
[47] S. G. Krantz and H. R. Parks. The implicit function theorem: History, theory, and applications. Springer, 2012.
[48] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[49] B. M. Lake, R. Salakhutdinov, and J. B. 'enenhaum. Human-level concept learning through probabilistic program induction. Science, 350(6266); 35066266):1332-11338 1332-1338, 2015.
[50] Y. LeCun, I.. Bottou, Y. Bengio, and Haffner. Gradient-based learning applied document recognition. Proceedings of the IEEE, 86 (6611)):2782224, 1998.
13

---

[51] Z. Liu, P. Luo. X. Wang, and X. Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[52] C. Lu, J. Chen, C. Li, Q. Wang, and J. Zhu. Implicit normalizing flows. In International Conference on Learning Representations (ICLR), 2021.
[53] A. Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu. Towards deep learning models resistant to adversarial attacks. ArXiv, abs/ fass///66.0083. 2018.
[54] A. Makkuva, A. Taghvaci, S. Oh, and J. Lee. Optimal transport mapping via input convex neural networks. In ICML, 2020.
[55] A. Nichol and J. Schulman. Reptile: scalable metalearning algorithm. arXiv: Learning, 2018.
[56] G. Ongie, A. Jalal, C. A. Metzler, R. G. Baraniuk, A. G. Dimakis, and R. Willett. Deep learning techniques for inverse problems in imaging, IEEE Journal on Selected Areas in Information Theory, 1(1):39-56, 2020.
[57] N. Papernot, P. D. McDaniel, X. Wu, S. Jha, and A. Swami. Distillation as a defense to adversarial perturbations against deep neural networks, 2016. In 2016 IEEE Symposium on Security and Privacy (SP), 2015.
[58] A. Rajeswaran, C. Finn, S. M. Kakade, and S. Levine. Meta-learning with implicit gradients. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019. URL htt.tt/e//ccccccco.......o..oopp
Paper.
[59] M. Revay, R. Wang, and R. Manchester. Lipschitz bounded equilibrium networks. arXiv:2010.01332, 2020.
[60] I.. F. Richardson. Ix. the approximate arithmetical solution by finite differences of physical problems involving differential equations, with an application to the stresses in a masonry dam. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of Mathematical or Physical Character, 110er,,1044555 2106459-4700)300-337,, 1911.
[61] Y. Rubanova. R. T. Chen, and D. Duvenaud. Latent ODEs for treeuuarly-sampll time series. 1rXiv:190..03907, 2019.
[62] E. K. Ryu and S. Boyd. Primer on monotone operator methods. Appl. Comput. Math, 15(1): 3-43, 2016.
[63] T. Salimans, A. Karpathy, X. Chen, and D. Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:1001.05517, 2017.
[64] C. Szegedy, W. Zaremba, I. Sutskever, J. Bruna, D. Erhan, I. Goodfellow, and R. Fergus. Intriguing properties of neural networks. 2014.
[65] G. Tao, S. Ma, Y. Liu, and X. Zhang. Attacks meet interpretability: Attribute-steered detection of adversarial samples. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc 2018 URL httpttp:::ppoceeeeceeetetooooeee Eile/b99463 21114444705777711017/11111111110
[66] H. F. Walker and P. Ni. Anderson acceleration for fixed-point iterations. SIAM Journal on Numerical Analysis, 49( 2011.
[67] P.-W. Wang, P. Donti, B. Wilder, and Z. Kolter. SATNet: Bridging deep learning and logical reasoning using a differentiable satisfibillity solver. K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1545-6554 PMLR, 09-15 Jun 2019. URI http://proceeding: m.lr. wwt.:/swww.......eetemmmmmme.eee /97 v97 wang 19e.
14

---

[68] E. Winston and J. Z. Kolter. Monotone operator equilibrium networks. In Neural Information Processing Systems, 2020.
[69] E. Wong, L. Rice, and J. Z. Kolter. Fast is better than free: Revisiting adversarial training. In International Conference Learning Representations, 2020. URL /
net Seet/Eoruum2.d=BB0000EEEE
[70] Y. Wu and K. He. Group normalization. In Proceedings the European conference on computer vision (ECCV), pages 1-1 2018.
[71] Y. Wu, Donahue, D. Balduzzi, K. Simonyan, and T. Lillicrap. [LOGAN]: Latent optimisation for generative adversarial networks, 2020. URL http::://ppennevieww.ettt///
id-rJeU 1SFvr
[72] A. Zadeh, Y.-.. Lim, Liang, and L..P. Morency. Variational auto-decoder. arXiv preprint arXiv: a903v: 190.00000.00 00840, 2019.
[73] J. Zhang, B. O' O'Donoguue, and S. Boyd. Globally convergent type-i anderson acceleration for nonsmooth fixed-point iterations. SIAM Journal on Optimization, 30(4):3170-3197, 2020.
[74] L. Zintgraf, K. Shiarli, Kurin, K. Hofmann, and S. Whiteson. Fast context adaptation via meta-learning. In International Conference Machine Learning, pages 7693-7702. PMLR, 2019.
15


Reviews:
[{"review_id": "taH4cZ3pkPw", "reviewer": "Reviewer_EACo", "summary": "The paper describes a method for optimizing the inputs of a deep equilibrium model (DEQ) at inference or training time by combining the input optimization procedure with the equilibrium finding procedure, transforrming the problem in an augmented DEQ.\n\nThe method is evaluated on several image-based tasks: generative model with optimized latents, inpainting/deblurring and generative adversarial training, yielding competitive results with baselines and a significant speedup compared to the typical multi-level optimization procedure.", "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 combines two active areas or research: implicit models, specifically DEQs, and forward-pass optimization on latent or input variables which is used on a variety of approaches.\n\nIt has been previously noted that  forward-pass optimization can be interpreted as computing the fixed point of an implicit model, here the authors show that when this is applied to a model which is already a DEQ in its base (unoptimized) forward pass, the computation of the DEQ internal fixed point and the external optimization fixed point can be combined as the computation of the fixed point of an augmented DEQ. This is more efficient than backpropagating through the DEQ and solving the optimization problem by gradient descent.\n\nThe research question is interesting, the paper is well written and easy to follow. The experimental results show comparable results with the baselines and substantial speedups. The tasks are reasonably varied. The only drawback is that no code release is planned.\n\nOverall I think this is a good contribution to the field.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors discuss the main limitation which is the substantial memory penalty of their method.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "sl5APsDTFF1", "reviewer": "Reviewer_DfXs", "summary": "This submission considers the problem of joint inference and input optimization for DEQ model. The main contribution is that it can perform inference and input optimization simultaneously. \n", "questions": "", "limitations": "", "rating": 5, "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": "Pros:\nThe most important contribution of this submission is that it shows that there is no need to exactly finds the fixed point before performing the input optimization. In some special cases, the input p[timizaion won't introduce additional calculation since we can re-use the converged solution for the inner loop for parameters updating. \n\nCons:\n1. Writing is too casual, and the motivation is very unclear. It is very unfriendly to readers who don't know much about DEQ and its related underlying problem.\n\n2. The convergence guarantee for the problems, such as Eq.(5), relies on the assumptions on the  $\\ell(\\cdot)$ and $f_{\\theta}$, which are all missing in this submission. Things should be more rigorous, especially for the mathematical process.  \n\n3. For experimental parts, there are too few methods to compare. The experiment mission is too simple, and the scale is too small, which cannot fully show the advantages of the proposed method.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No need to consider the Societal Impact", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4h"}, "extra_text": {}}, {"review_id": "7zP8raRUMoY", "reviewer": "Reviewer_yNXE", "summary": "The paper argues that input optimization problems are pervasive in modern machine learning methods (adversarial example optimization, latent optimization in generative models, image completion/denoising). The authors then make a connection to deep equilibrium models (DEQ), which generate an output as a fixed-point of a single non-linear function. The connection is that both input optimization and DEQ process are fixed point iterations and can be solved jointly as an equality-constrained optimization problem. This results in 3-9x speeds on adversarial example generation and generative image modeling tasks.", "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": "Originality:\nThe idea of solving for input and equilibrium point optimization as a joint problem is to the best of my knowledge novel and is well-motivated - indeed input optimization with DEQs is wasteful if done naively as nested optimization and taking into account the structure of the problem allows for more efficient methods.\n\nQuality:\nThe mathematical formulation is sound to the best of my knowledge.\n\nClarity:\nThe paper is clearly written and all technical components are well-motivated and given sufficient intuition. I especially appreciated an alternative interpretation of the formulation in section 3.3 as that more closely connected to my background and made it easier to reason about properties of this method.\n\nThe paper and appendix contain enough details to recreate the experiments.\n\nAs authors admit, the variable naming convention sometimes clashes with current generative modeling literature. While I understand the reasoning behind this, it would be nice to have an algorithm figure with all variables annotated so that the reader can quickly refer to for refresher on the meaning of variable names.\n\nSignificance:\nDEQs are an intriguing class of models that deserve continuing investigation. I see the impact of this paper as two-fold: 1) it extends the capabilities of DEQ models by making input optimization tractable in these models 2) it gives an inspirational example of how DEQ models can be extended and combined with other mathematical tools.\n\nThe empirical results in generative image modeling on CelebA and FFHQ qualitatively don’t look very strong and quantitatively the VAE baseline is rather weak (HDCGAN reported FID of 8.44 on 64x64 CelebA generation in 2017, if I’m comparing the right benchmark). As such, I don’t see the community focused on best generative modeling performance adopting this work at the moment (of course, this is not the point of the paper, but it does limit significance).\n\nThe other two concerns I have about general adoption and impact of this work is stability of training and time and space complexity. For the latter, authors provide wallclock time comparisons and briefly mention memory footprint challenges. However, it would be very helpful to people considering building on this work to know time and space complexity of your algorithm, especially since memory could be a bottleneck in scaling this approach up.\n\nFor the concerns of stability, first I really appreciate authors for sharing their observations of stability issues rather than avoiding disclosing them - this is very helpful for readers to know and for presents opportunities for people to build on the work! Do the authors have any hypotheses for stability issues they could share in the paper? Is it potentially from fixed-point iteration not running to convergence? General instability of non-convex Lagrangian in eq 13? In this case, have you tried augmented Lagrangian known to improve stability? Any additional discussion here would be very helpful.\n\nTowards applicability, I want to mention that another area where input optimization is prevalent recently is energy-based models. These models generate a datapoint as the minimum of the energy function (which could be a DEQ model) [A Tutorial on Energy-Based Learning by LeCun et al] and often use this optimization process in contrastive training of the models [i.e. Implicit Generation and Generalization in Energy-Based Models by Du et al, Divergence Triangle by Han et al, and many others].\n\nBut even on the first two merits alone, combined with the paper being sound I believe this work is worthy of sharing with our community and will inspire follow-up research.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I discussed the limitation in Significance section of my main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3-4 hours"}, "extra_text": {}}, {"review_id": "3kaJLU7pmNn", "reviewer": "Reviewer_XjMd", "summary": "The paper proposes an extension to the DEQ framework that allows to simultaneously perform input and output optimization of the implicit DEQ layer. The extension called JIIO can itself be cast as another DEQ and backpropagated through. This appears to be useful in a number of applications such as decoder-only autoencoders where the encoder is defined implicitly as optimization over the latent representation. ", "questions": "", "limitations": "", "rating": 6, "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": "Novelty:\n\nThe proposed method is novel to my knowledge.\n\nClarity:\n\nThe paper is well-written and easy to follow. I have some comments that are related to presentation of the method as well and I list them in the Quality section.\n\nQuality:\n\nI followed derivations of the optimization scheme and I find it sound.\n\nMy main point of concern is instability of the method which is important to properly investigate (if not fix completely) to make JIIO really useful to community.\n\n1) I would like to have a better explanation for this rather than the trivial intuition that the optimization problem is more difficult. Especially given that some of the optimization methods simply failed to converge.\n2) Similarly, it is unclear how to choose the output in JIIO given the optimization instability. Is the lowest-loss solution returned? Is it done both during training and during inference?\n3) As an idea on the potential source of instability, did authors measure error in estimating the inverse Jacobian / its positive-definiteness? \n4) Is it an artifact arising from using non-optimal step sizes?\n\nSignificance:\n\nThe list of considered applications is somewhat niche in the sense that DEQs could be not necessarily the best framework to target those, but I still think it's a valid contribution and JIIO will find more use in the community.\n\nI'm keen to revise my rating based on authors reply.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No issues.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers found the paper to be interesting, novel and well-written. This paper will likely be of interest to the community and inspire follow-on work. In review, the authors agreed to add a discussion on assumptions on assumptions on l() and f_theta required to ensure (local) convergence; we expect the authors to follow through on this.
Decision: Accept (Poster)