Joint inference and input optimization in equilibrium networks

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 gradientbased 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 . 1 Introduction . Many settings in deep learning involve optimization over the inputs to a network to minimize some desired loss . For example , for a “ generator ” network G : Z → X that maps from latent space Z to an observed space X , it may be desirable to find a latent vector z ∈ Z that most closely produces some target output x ∈ X by solving the optimization problem ( e.g . [ 10 , 13 ] ) minimize z∈Z ‖x−Gθ ( z ) ‖22 . ( 1 ) As another example , constructing adversarial examples for classifiers [ 28 , 53 ] typically involves optimizating over a perturbation to a given input ; i.e. , given a classifier network g : X → Y , task loss ∗Correspondence to : Swaminathan Gurumurthy < sgurumur @ andrew.cmu.edu > Code available at https : //github.com/locuslab/JIIO-DEQ 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ` : Y → R+ , and a sample x ∈ X , we want to solve maximize ‖δ‖≤ ` ( g ( x+ δ ) ) . ( 2 ) More generally , a wide range of inverse problems [ 10 ] and other auxiliary tasks [ 22 , 3 ] in deep learning can also be formulated in such a 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 . [ 5 ] introduced deep equilibrium models ( DEQs ) which instead treat the network as repeated applications of a single layer and compute the output of the network as a solution to an equilibrium-finding problem instead of simply specifying a 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 . 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 an iterative method ( e.g. , gradient descent ) , we can combine this optimization fixed-point iteration with the forward DEQ fixed point iteration all within a single “ augmented ” DEQ model that simultaneously 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 a 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-based 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 ) gradientbased meta-learning . We show that in all cases , performing this simultaneous optimization and forward inference accelerates the process over a more naive inner/outer optimization approach . For instance , using the combined approach leads to a 3.5-9x speedup for generative DEQ networks , a 3x speedup in adverarial training of DEQ networks and a 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 . 2 Related Work . Implicit layers . Layers with implicitly defined depth have gained tremendous popularity in recent years [ 46 , 19 , 29 ] . Rather than a 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-efficient to train . Some recent examples of implicit layers include optimization layers [ 16 , 1 ] , deep equilibrium models [ 5 , 6 , 68 , 40 , 52 ] , neural ordinary differential equations ( ODEs ) [ 14 , 18 , 61 ] , logical structure learning [ 67 ] , and continuous generative models [ 30 ] . In particular , deep equilibrium models ( DEQs ) [ 5 ] define the output of the model as 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 method [ 11 ] ) . 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 as 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-finding problem that defines an invertible function between the input x and output z ? . Perhaps the most relevant approach to our work is Gilton et al . [ 26 ] , which specifically formulates inverse imaging problems as a DEQ model . In contrast , our approach focuses on solving input optimization problems where the network of interest is already a DEQ , and thus the combined optimization and forward inference task leads to a substantially different set of update equations and tradeoffs . 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 adversarial examples [ 53 , 45 ] , solving inverse problems [ 10 , 13 , 56 ] , learning generative models [ 9 , 72 ] , meta-learning [ 58 , 22 , 74 , 32 ] etc . For most of these examples , input optimization is typically 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 to 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 [ 71 ] . However , in all of these cases , the input is optimized for just a few ( mostly 1 ) iterations . In this work , we present a 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 . [ 13 ] 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 . [ 15 ] , Gilton et al . [ 25 ] , Gregor and LeCun [ 31 ] 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 . [ 10 ] , Chang et al . [ 13 ] , but also show flexible extension of our framework to train implicit models for supervised inverse problem solving . Another crucial application of input optimization is to find adversarial examples [ 64 , 28 ] . This manifests as optimizing an objective that incentivices an incorrect prediction by the classifier , while constraining the input to be within a bounded region of the original input . Many attempts have been made on the defense side [ 57 , 37 , 65 , 69 ] . The most successful strategy thus far has been adversarial training with a 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 to train implicit models to perform competitively with PGD in guaranteeing adversarial robustness , but at much faster speeds . While the examples above were illustrated with non-convex networks , attempts have also been made to design networks whose output is a convex function of the input [ 2 ] . This allows one to use more sophisticated optimization algorithms , but usually at a 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 an 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 .

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.