Bounds all around: training energy-based models with bidirectional bounds

1 Energy-based models . Energy-based models ( EBMs ) are probabilistic models that draw inspiration from physics and have a long history in machine learning ( Hopfield , 1982 ; Hinton and Sejnowski , 1983 ; Smolensky , 1986 ) . An EBM is specified in terms of an energy function Eθ : X → R that is parameterized by θ and defines a probability distribution over X from the Gibbs density : pθ ( x ) = exp ( −Eθ ( x ) ) Zθ , Zθ = ∫ exp ( −Eθ ( x ) ) dx , ( 1 ) where Zθ is the normalization constant or partition function . In principle , any density can be described this way for a suitable choice ofEθ . EBMs are typically learned using maximum likelihood estimation ( MLE ) , where we wish to find a value of θ that minimized the negative data log-likelihood : L ( θ ) : = −Ex∼pdata ( x ) [ log pθ ( x ) ] = Ex∼pdata [ Eθ ( x ) ] + logZθ , ( 2 ) where pdata is the data generating distribution . The fundamental challenge with EBMs for non-standard energy functions is the lack of a closed-form expression for the normalization constant , Zθ , which hinders exact pointwise density evaluation , sampling and learning . Therefore , these tasks are traditionally performed using approximate methods such as Markov chain Monte Carlo ( MCMC , Metropolis et al. , 1953 ; Neal et al. , 2011 ) . Similar to the Boltzmann learning rule ( Hinton and Sejnowski , 1983 ; Osogami , 2017 ) , gradient-based learning for MLE involves evaluating the gradient ∇θL ( θ ) = Ex∼pdata [ ∇θEθ ( x ) ] − Ex∼pθ [ ∇θEθ ( x ) ] , where the expectation in the last term necessitates MCMC approximation . Recently , a series of papers ( Kim and Bengio , 2016 ; Kumar et al. , 2019 ; Abbasnejad et al. , 2020 ; Che et al. , 2020 ) have established links between EBMs and Wasserstein GANs ( WGANs , Arjovsky et al. , 2017 ) , such that EBMs can be approximately learned through a minimax game . Xie et al . ( 2017 , 2018c ) and Wu et al . ( 2018 ) also employed the adversarial game of EBM . This allows for significant improvements in EBM learning over MCMC based methods ( Grathwohl et al. , 2021 ) . ∗equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we remark that variational bounds on the value function of the minimax game can be problematic as this is both maximized and minimized . We propose to both upper and lower bound the negative log-likelihood , and alternate between their respective minimization and maximization to alleviate this concern . Evaluation of the bounds requires evaluating the entropy of the generator , and we provide a new efficient estimator . We link the upper bound to the use of gradient penalties which are known to stabilize training . Experimentally , we demonstrate that our approach matches or surpasses state-of-the-art on diverse tasks at negligible performance increase . 1.1 Background : Variational bounds and minimax games . Following Grathwohl et al . ( 2021 ) , we can lower bound the log normalization constant using a proposal ( variational ) distribution pg ( x ) and Jensen ’ s inequality : logZθ = logEx∼pg [ exp ( −Eθ ( x ) ) pg ( x ) ] ≥ Ex∼pg [ log exp ( −Eθ ( x ) ) pg ( x ) ] = −Ex∼pg [ Eθ ( x ) ] +H [ pg ] ( 3 ) where H [ pg ] = −Epg [ log pg ] is the ( differential ) entropy of the proposal distribution . This means that we also obtain a lower bound on the negative log likelihood function given by L ( θ ) : = Ex∼pdata ( x ) [ Eθ ( x ) ] − Ex∼pg ( x ) [ Eθ ( x ) ] +H [ pg ] ≤ L ( θ ) . ( 4 ) We note that this bound is tight when pg = pθ , since the bound on the log normalization constant ( 3 ) can be equivalently expressed as logZθ ≥ logZθ − KL ( pg‖pθ ) . To tighten the bound , we seek a proposal distribution pg that maximises L ( θ ) , while for MLE we want to find a θ that minimizes L ( θ ) . Accordingly , MLE can be formulated as a minimax game : min Eθ max pg { L ( θ ) } = min Eθ max pg { Ex∼pdata ( x ) [ Eθ ( x ) ] − Ex∼pg ( x ) [ Eθ ( x ) ] +H [ pg ] } . ( 5 ) This gives a tractable and MCMC-free approach to EBM learning , assuming we can evaluate the entropy ( Grathwohl et al. , 2021 ) . Our key issue is that both minimizing and maximizing a lower bound is potentially unstable . By minimizing a lower bound , we run the risk of finding an optimum , where the bound is loosest , rather than where it is informative about the optima of the true objective . In particular , a minima of the lower bound may be −∞ , which is rather unhelpful . Good results have been reported from using a lower bound in the minimax game ( Dai et al. , 2017 ; Kumar et al. , 2019 ; Grathwohl et al. , 2021 ) , but the conceptual issue remains . The WGAN loss function is , as noted by Arjovsky et al . ( 2017 ) , very similar to the lower boundL ( θ ) , LWGAN = Ex∼pdata [ Eθ ( x ) ] − Ex∼pg [ Eθ ( x ) ] , ( 6 ) with the missing entropy term being the only difference . Lessons from the WGAN can , thus , be expected to apply to the EBM setting as well . For example , the success of gradient clipping ( Arjovsky et al. , 2017 ) and gradient penalties ( Gulrajani et al. , 2017 ) in WGAN training , inspired Grathwohl et al . ( 2021 ) to heuristically use gradient penalties in variational inference to stabilize training . Our work will provide justification to such heuristics . 2 Approximate minimax games through bidirectional bounds . Instead of solving a minimax game with a lower bounded value function ( 5 ) , we propose to bound L ( θ ) from both above and below : bL ( θ ) c ≤ L ( θ ) ≤ dL ( θ ) e. ( 7 ) With this , we can now follow an optimization strategy of alternating 1. minimize dL ( θ ) e with respect to Eθ . 2. maximize bL ( θ ) c with respect to pg . This avoids the potential pitfalls of minimizing a lower bound . 2.1 A lower bound for maximizing pg . We already have a lower bound L ( θ ) on L ( θ ) in Equation ( 4 ) . The first two terms of the bound can readily be evaluated by sampling from pdata and pg , respectively . The entropy of pg is less straightforward . Following standard practice in GANs , we define pg through a base variable z ∼ N ( 0 , I ) , which is passed through a ‘ generator ’ network G : Rd → RD , i.e . x = G ( z ) ∼ pg . This constructs a density over a d-dimensional manifold2 in RD , which by standard change-of-variables is log pg ( G ( z ) ) = log p0 ( z ) − 1 2 log det ( JᵀzJz ) , ( 8 ) where Jz ∈ RD×d is the Jacobian ofG at z . This assume that the Jacobian exist , i.e . that the generator G has at least one derivative . The entropy of pg is then H [ pg ] = −Ex∼pg [ log pg ( x ) ] = H [ p0 ] + Ez∼p0 [ 1 2 log det ( JᵀzJz ) ] . ( 9 ) The entropy of p0 trivially evaluates to H [ p0 ] = d/2 ( 1 + log ( 2π ) ) . To avoid an expensive evaluation of the log-determinant-term , we note that it is easily bounded from below : 1 2 log det ( JᵀzJz ) = 1 2 d∑ i=1 log s2i ≥ d log s1 , ( 10 ) where sd ≥ . . . ≥ s1 are the singular values of the Jacobian Jz . Note that in general , this bound is not tight . With this , we have our final lower bound : L ( θ ) ≥ bL ( θ ) c = Ex∼pdata ( x ) [ Eθ ( x ) ] − Ex∼pg ( x ) [ Eθ ( x ) ] +H [ p0 ] + Ez∼p0 [ d log s1 ( z ) ] . ( 11 ) 2.2 An upper bound for minimizing Eθ . There are many ways to arrive at an upper bound for the log-likelihood . Due to the strong ties between the lower bound ( 4 ) and the WGAN objective , we take inspiration from how WGANs are optimized . The original WGAN paper ( Arjovsky et al. , 2017 ) noted optimization instabilities unless gradient clipping was applied . Gulrajani et al . ( 2017 ) noted that this clipping could be avoided by adding a gradient penalty , such that the loss become LWGAN-GP = LWGAN + λEx̂∼g ( x ) [ ( ‖∇x̂Eθ ( x̂ ) ‖2 − 1 ) 2 ] , ( 12 ) where x̂ is a uniform interpolation between real data and generated samples . This is generally acknowledged as a good way to train a WGAN ( Gulrajani et al. , 2017 ) . We observe that this loss is similar to our lower bound , but the gradient penalty adds a positive term , such that one could speculate that we might have an upper bound . With this intuition in hand , we prove the following statement in the supplementary material . Theorem 1 Suppose f : X → R is L-Lipschitz continuous , g ( x ) is a probability density function with finite support , then there exists constants M , m ≥ 0 and p ≥ 1 such that : logEx∼g ( x ) [ f ( x ) ] − Ex∼g ( x ) [ log f ( x ) ] ≤M ( Ex∼g ( x ) [ |∇x log f ( x ) |p ] +m ) 1/p . ( 13 ) Note that if Ex∼g ( x ) [ |∇x log f ( x ) |p ] +m ≥ 1 , the bound in Equation ( 13 ) can be simplified by dropping the power 1/p , log [ Ex∼g ( x ) [ f ( x ) ] ] − Ex∼g ( x ) [ log f ( x ) ] ≤MEx∼g ( x ) [ |∇x log f ( x ) |p ] +Mm . ( 14 ) For our model , we set f ( x ) = exp ( −Eθ ( x ) ) pg ( x ) and g ( x ) = pg ( x ) . Empirically , we observe the bound simplification in Equation ( 14 ) holds through-out most of the training , and use the following upper bound : dL ( θ ) e = L ( θ ) +MEx∼pg ( x ) [ |∇xEθ ( x ) +∇x log pg ( x ) | p ] +Mm . ( 15 ) For a given choice of f ( x ) , we can view m as a constant and the term Mm can be ignored during optimization . This upper bound can directly be interpreted as the lower bound ( 4 ) plus a gradient penalty , albeit one of a different form than the traditional WGAN penalty , which is derived purely from a regularization perspective . Our upper bound can , thus , be seen as a justification of the regularization from a maximum likelihood perspective . 2Note that Eq . 8 is only the density on the spanned manifold , and that the density is zero off the manifold . Bound tightness When pθ = pg , we have that |∇xEθ ( x ) +∇x log pg ( x ) |p = 0 . In Theorem 1 , m is a constant related to the Lipschitz constant of log f ( x ) satisfying |∇x log f ( x ) |p ≤ |∇y log f ( y ) |p+m for all x , y ( see proof for details ) . When pθ = pg we also have |∇x log f ( x ) |p = 0 , such that m = 0 . Our upper bound is then dL ( θ ) e = L ( θ ) = L ( θ ) , and hence it is tight .

Summary: This paper proposes a method to train EBMs with lower and upper variational bounds. Previous approaches used only one side of the bound. This causes significant training instabilities as the training objective involves alternately maximizing and minimizing the same quantity.  The lower bound is estimated by lower-bounding the log determinant of the jacobian multiplied with itself with the smallest singular value of the jacobian.   The upper bound is inspired by tricks used to make WGAN and related EBM training approaches stable by regularizing the gradient.   The minimum singular value, needed for computing the lower bound, is computed using an iterative algorithm run for a small number of steps. Automatic differentiation is used cleverly to ensure full Jacobians don’t need to be instantiated.  The upper bound is loosened, and Hutchinson’s trace estimator is employed to efficiently compute.