Taming GANs with Lookahead-Minmax

Generative Adversarial Networks are notoriously challenging to train . The underlying minmax optimization is highly susceptible to the variance of the stochastic gradient and the rotational component of the associated game vector field . To tackle these challenges , we propose the Lookahead algorithm for minmax optimization , originally developed for single objective minimization only . The backtracking step of our Lookahead–minmax naturally handles the rotational game dynamics , a property which was identified to be key for enabling gradient ascent descent methods to converge on challenging examples often analyzed in the literature . Moreover , it implicitly handles high variance without using large mini-batches , known to be essential for reaching state of the art performance . Experimental results on MNIST , SVHN , CIFAR-10 , and ImageNet demonstrate a clear advantage of combining Lookahead–minmax with Adam or extragradient , in terms of performance and improved stability , for negligible memory and computational cost . Using 30-fold fewer parameters and 16-fold smaller minibatches we outperform the reported performance of the class-dependent BigGAN on CIFAR-10 by obtaining FID of 12.19 without using the class labels , bringing state-of-the-art GAN training within reach of common computational resources . Our source code is available : https : //github.com/Chavdarova/LAGAN-Lookahead_Minimax . 1 INTRODUCTION . Gradient-based methods are the workhorse of machine learning . These methods optimize the parameters of a model with respect to a single objective f : X → R . However , an increasing interest for multi-objective optimization arises in various domains—such as mathematics , economics , multiagent reinforcement learning ( Omidshafiei et al. , 2017 ) —where several agents aim at optimizing their own cost function fi : X1 × · · · × XN → R simultaneously . A particularly successful class of algorithms of this kind are the Generative Adversarial Networks ( Goodfellow et al. , 2014 , ( GANs ) ) , which consist of two players referred to as a generator and a discriminator . GANs were originally formulated as minmax optimization f : X ×Y → R ( Von Neumann & Morgenstern , 1944 ) , where the generator and the discriminator aim at minimizing and maximizing the same value function , see § 2 . A natural generalization of gradient descent for minmax problems is the gradient descent ascent algorithm ( GDA ) , which alternates between a gradient descent step for the min-player and a gradient ascent step for the max-player . This minmax training aims at finding a Nash equilibrium where no player has the incentive of changing its parameters . Despite the impressive quality of the samples generated by the GANs—relative to classical maximum likelihood-based generative models—these models remain notoriously difficult to train . In particular , poor performance ( sometimes manifesting as “ mode collapse ” ) , brittle dependency on hyperparameters , or divergence are often reported . Consequently , obtaining state-of-the-art performance was shown to require large computational resources ( Brock et al. , 2019 ) , making well-performing models unavailable for common computational budgets . ∗Equal contributions . Correspondence to firstname.lastname @ epfl.ch . It was empirically shown that : ( i ) GANs often converge to a locally stable stationary point that is not a differential Nash equilibrium ( Berard et al. , 2020 ) ; ( ii ) increased batch size improves GAN performances ( Brock et al. , 2019 ) in contrast to minimization ( Defazio & Bottou , 2019 ; Shallue et al. , 2018 ) . A principal reason is attributed to the rotations arising due to the adversarial component of the associated vector field of the gradient of the two player ’ s parameters ( Mescheder et al. , 2018 ; Balduzzi et al. , 2018 ) , which are atypical for minimization . More precisely , the Jacobian of the associated vector field ( see def . in § 2 ) can be decomposed into a symmetric and antisymmetric component ( Balduzzi et al. , 2018 ) , which behave as a potential ( Monderer & Shapley , 1996 ) and a Hamiltonian game , resp . Games are often combination of the two , making this general case harder to solve . In the context of single objective minimization , Zhang et al . ( 2019 ) recently proposed the Lookahead algorithm , which intuitively uses an update direction by “ looking ahead ” at the sequence of parameters that change with higher variance due to the “ stochasticity ” of the gradient estimates–called fast weights–generated by an inner optimizer . Lookahead was shown to improve the stability during training and to reduce the variance of the so called slow weights . Contributions . Our contributions can be summarized as follows : • We propose Lookahead–minmax for optimizing minmax problems , that applies extrapolation in the joint parameter space ( see Alg 1 ) , so as to account for the rotational component of the associated game vector field ( defined in § 2 ) . • In the context of : ( i ) single objective minimization : by building on insights of Wang et al . ( 2020 ) , who argue that Lookahead can be interpreted as an instance of local SGD , we derive improved convergence guarantees for the Lookahead algorithm ; ( ii ) two-player games : we elaborate why Lookahead–minmax suppresses the rotational part in a simple bilinear game , and prove its convergence for a given converging base-optimizer ; in § 3 and 4 , resp . • We motivate the use of Lookahead–minmax for games by considering the extensively studied toy bilinear example ( Goodfellow , 2016 ) and show that : ( i ) the use of lookahead allows for convergence of the otherwise diverging GDA on the classical bilinear game in full-batch setting ( see § 4.2.1 ) , as well as ( ii ) it yields good performance on challenging stochastic variants of this game , despite the high variance ( see § 4.2.2 ) . • We empirically benchmark Lookahead–minmax on GANs on four standard datasets— MNIST , CIFAR-10 , SVHN and ImageNet—on two different models ( DCGAN & ResNet ) , with standard optimization methods for GANs , GDA and Extragradient , called LA–AltGAN and LA–ExtraGradient , resp . We consistently observe both stability and performance improvements at a negligible additional cost that does not require additional forward and backward passes , see § 5 . 2 BACKGROUND . GAN formulation . Given the data distribution pd , the generator is a mapping G : z 7→ x , where z is sampled from a known distribution z ∼ pz and ideally x ∼ pd . The discriminator D : x 7→ D ( x ) ∈ [ 0 , 1 ] is a binary classifier whose output represents a conditional probability estimate that an x sampled from a balanced mixture of real data from pd and G-generated data is actually real . The optimization of a GAN is formulated as a differentiable two-player game where the generator G with parameters θ , and the discriminator D with parameters ϕ , aim at minimizing their own cost function Lθ and Lϕ , respectively , as follows : θ ? ∈ arg min θ∈Θ Lθ ( θ , ϕ ? ) and ϕ ? ∈ arg min ϕ∈Φ Lϕ ( θ ? , ϕ ) . ( 2P-G ) When Lϕ = −Lθ the game is called a zero-sum and equation 2P-G is a minmax problem . Minmax optimization methods . As GDA does not converge for some simple convex-concave game , Korpelevich ( 1976 ) proposed the extragradient method , where a “ prediction ” step is performed to obtain an extrapolated point ( θt+ 12 , ϕt+ 12 ) using GDA , and the gradients at the extrapolated point are then applied to the current iterate ( θt , ϕt ) as follows : Extrapolation : { θt+ 12 =θt−η∇θL θ ( θt , ϕt ) ϕt+ 12 =ϕt−η∇ϕL ϕ ( θt , ϕt ) Update : { θt+1=θt−η∇θLθ ( θt+ 12 , ϕt+ 12 ) ϕt+1=ϕt−η∇ϕLϕ ( θt+ 12 , ϕt+ 12 ) ( EG ) where η denotes the step size . In the context of zero-sum games , the extragradient method converges for any convex-concave function L and any closed convex sets Θ and Φ ( Facchinei & Pang , 2003 ) . The joint vector field . Mescheder et al . ( 2017 ) and Balduzzi et al . ( 2018 ) argue that the vector field obtained by concatenating the gradients of the two players gives more insights of the dynamics than studying the loss surface . The joint vector field ( JVF ) and the Jacobian of JVF are defined as : v ( θ , ϕ ) = ( ∇θLθ ( θ , ϕ ) ∇ϕLϕ ( θ , ϕ ) ) , and v′ ( θ , ϕ ) = ( ∇2θLθ ( θ , ϕ ) ∇ϕ∇θLθ ( θ , ϕ ) ∇θ∇ϕLϕ ( θ , ϕ ) ∇2ϕLϕ ( θ , ϕ ) ) , resp . ( JVF ) Rotational component of the game vector field . Berard et al . ( 2020 ) show empirically that GANs converge to a locally stable stationary point ( Verhulst , 1990 , LSSP ) that is not a differential Nash equilibrium–defined as a point where the norm of the Jacobian is zero and where the Hessian of both the players are definite positive , see § C. LSSP is defined as a point ( θ ? , ϕ ? ) where : v ( θ ? , ϕ ? ) = 0 , and R ( λ ) > 0 , ∀λ ∈ Sp ( v′ ( θ ? , ϕ ? ) ) , ( LSSP ) where Sp ( · ) denotes the spectrum of v′ ( · ) andR ( · ) the real part . In summary , ( i ) if all the eigenvalues of v′ ( θt , ϕt ) have positive real part the point ( θt , ϕt ) is LSSP , and ( ii ) if the eigenvalues of v′ ( θt , ϕt ) have imaginary part , the dynamics of the game exhibit rotations . Impact of noise due to the stochastic gradient estimates on games . Chavdarova et al . ( 2019 ) point out that relative to minimization , noise impedes more the game optimization , and show that there exists a class of zero-sum games for which the stochastic extragradient method diverges . Intuitively , bounded noise of the stochastic gradient hurts the convergence as with higher probability the noisy gradient points in a direction that makes the algorithm to diverge from the equilibrium , due to the properties of v′ ( · ) ( see Fig.1 , Chavdarova et al. , 2019 ) . 3 LOOKAHEAD FOR SINGLE OBJECTIVE . In the context of single objective minimization , Zhang et al . ( 2019 ) recently proposed the Lookahead algorithm where at every step t : ( i ) a copy of the current iterate ω̃t is made : ω̃t ← ωt , ( ii ) ω̃t is then updated for k ≥ 1 times yielding ω̃t+k , and finally ( iii ) the actual update ωt+1 is obtained as a point that lies on a line between the two iterates : the current ωt and the predicted one ω̃t+k : ωt+1 ← ωt + α ( ω̃t+k − ωt ) , where α ∈ [ 0 , 1 ] . ( LA ) Algorithm 1 General Lookahead–Minmax pseudocode . 1 : Input : Stopping time T , learning rates ηθ , ηϕ , ini- tial weights θ0 , ϕ0 , lookahead k and α , losses Lθ , Lϕ , dϕ ( · ) , d θ ( · ) base optimizer updates defined in § 4 , pd and pz real and noise–data distributions , resp . 2 : θ̃0,0 , ϕ̃0,0 ← θ0 , ϕ0 3 : for t ∈ 0 , . . . , T−1 do 4 : for i ∈ 0 , . . . , k−1 do 5 : Sample x , z ∼ pd , pz 6 : ϕ̃t , i+1 = ϕ̃t , i − ηϕdϕt , i ( θ̃t , i , ϕ̃t , i , x , z ) 7 : Sample z ∼ pz 8 : θ̃t , i+1 = θ̃t , i+1 − ηθdθt , i ( θ̃t , i , ϕ̃t , i , z ) 9 : end for 10 : ϕt+1 = ϕt + αϕ ( ϕ̃t , k −ϕt ) 11 : θt+1 = θt + αθ ( θ̃t , k − θt ) 12 : θ̃t+1,0 , ϕ̃t+1,0 ← θt+1 , ϕt+1 13 : end for 14 : Output : θT , ϕT Lookahead uses two additional hyperparameters : ( i ) k–the number of steps used to obtain the prediction ω̃t+k , as well as ( ii ) α– that controls how large a step we make towards the predicted iterate ω̃ : the larger the closest , and when α = 1 equation LA is equivalent to regular optimization ( has no impact ) . Besides the extra hyperparameters , LA was shown to help the used optimizer to be more resilient to the choice of its hyperparameters , achieve faster convergence across different tasks , as well as to reduce the variance of the gradient estimates ( Zhang et al. , 2019 ) . Theoretical Analysis . Zhang et al . ( 2019 ) study LA on quadratic functions and Wang et al . ( 2020 ) recently provided an analysis for general smooth non-convex functions . One of their main observations is that LA can be viewed as an instance of local SGD ( or parallel SGD , Stich , 2019 ; Koloskova et al. , 2020 ; Woodworth et al. , 2020b ) which allows us to further tighten prior results . Theorem 1 . Let f : Rd → R be L-smooth ( possibly non-convex ) and assume access to unbiased stochastic gradients σ2-bounded variance . Then the LA optimizer with hyperparemeters ( k , α ) converges to a stationary point E‖∇f ( ωout ) ‖2 ≤ ε ( for the proof refer to Appendix A ) , after at most O ( σ2 ε2 + 1 ε + 1− α α ( σ √ k − 1 ε3/2 + k ε ) ) iterations . Here ωout denotes uniformly at random chosen iterate of LA . Remark 1 . When in addition f is also quadratic , the complexity estimate improves to O ( σ2 ε2 + 1 ε ) . The asymptotically most significant term , O ( σ2 ε2 ) , matches with the corresponding term in the SGD convergence rate for all choices of α ∈ ( 0 , 1 ] , and when α→ 1 , the same convergence guarantees as for SGD can be attained . When σ2 = 0 the rate improves to O ( 1 ε ) , in contrast to O ( 1 ε2 ) in ( Wang et al. , 2020 ) .1 For small values of α , the worst-case complexity estimates of LA can in general be k times worse than for SGD ( except for quadratic functions , where the rates match ) . Deriving tighter analyses for LA that corroborate the observed practical advantages is still an open problem .

Summary: This work extends the recently proposed lookahead optimizer (which was designed for single-objective optimization) to minimax optimization, particularly GAN training. The authors claim that the backtracking step in lookahead optimizer alleviates the notorious rotational behavior in GAN dynamics. Moreover, the authors argue that the lookahead optimizer implicitly handles the high variance in the small-batch setting. Both arguments are backed up by toy experiments on stochastic bilinear games.  Finally, on standard image datasets, the lookahead minimax algorithm outperforms some popular algorithms and achieves state-of-the-art performance on CIFAR-10.