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 paper proposes a lookahead-minmax algorithm for optimizing minmax problems such as GANs, which updates the parameters (of both the generator and the discriminator) with the extrapolation. With a bilinear example, the authors show that the use of lookahead-minimax allows for convergence in cases where other methods does not, and yields good performance under high variance. Experiments of generative performance on several well-known public datasets demonstrates the effectiveness of the proposed method.