Dual Training of Energy-Based Models with Overparametrized Shallow Neural Networks

Energy-based models ( EBMs ) are generative models that are usually trained via maximum likelihood estimation . This approach becomes challenging in generic situations where the trained energy is nonconvex , due to the need to sample the Gibbs distribution associated with this energy . Using general Fenchel duality results , we derive variational principles dual to maximum likelihood EBMs with shallow overparametrized neural network energies , both in the active ( aka featurelearning ) and lazy regimes . In the active regime , this dual formulation leads to a training algorithm in which one updates concurrently the particles in the sample space and the neurons in the parameter space of the energy at a faster rate . We also consider a variant of this algorithm in which the particles are sometimes restarted at random samples drawn from the data set , and show that performing these restarts at every iteration step corresponds to score matching training . Using intermediate parameter setups in our dual algorithm thereby gives a way to interpolate between maximum likelihood and score matching training . These results are illustrated in simple numerical experiments . 1 INTRODUCTION . Energy-based models ( EBMs ) are explicit generative models which consider Gibbs measures defined through an energy function f , with a probability density proportional to exp ( −βf ( x ) ) , where β is the inverse temperature . Such models originate in statistical physics ( Gibbs , 2010 ; Ruelle , 1969 ) , and have become a fundamental modeling tool in statistics and machine learning ( Wainwright & Jordan , 2008 ; Ranzato et al. , 2007 ; LeCun et al. , 2006 ; Du & Mordatch , 2019 ; Song & Kingma , 2021 ) . Given data samples from a target distribution , the learning algorithms for EBMs attempt to estimate an energy function f to model the samples density . The resulting learned model can then be used to obtain new samples , typically through Markov Chain Monte Carlo ( MCMC ) techniques . The standard method to train EBMs is maximum likelihood estimation , i.e . the learned energy is the one maximizing the likelihood of the target samples , within a certain function class . One generic approach for this is to use gradient descent , where gradients may be approximated using MCMC samples from the trained model . However , this is computationally difficult for highly non-convex trained energies , due to ‘ metastability ’ , ie the presence of large basins in the energy landscape that trap trajectories for potentially exponential time . This has motivated a myriad of alternative losses to learn EBM energies , such as the popular score matching ; see ( Song & Kingma , 2021 ) for a review . All in all , such weaker losses come at the expense of a loss of statistical power , which motivates exploring computationally efficient methods for EBM maximum-likelihood estimation . EBMs also have structural connections with maximum entropy ( maxent ) models , which have been studied for decades through Fenchel duality . Dai et al . ( 2019b ) was the first work to leverage similar duality arguments for maximum likelihood EBM training . However , their analysis is restricted to energies lying in RKHS balls ( i.e . non-parametric linear models ) . Despite the appealing optimization properties of RKHS , these spaces of functions typically only contain very smooth functions when the dimension is large ( Berlinet & Thomas-Agnan , 2004 ) . A recent line of work—originating in supervised learning—has considered an alternative based on shallow neural networks ( Bach , 2017 ) , which admit a linear representation in terms of a measure over its parameters and are able to adapt to hidden low-dimensional structures in the data . The statistical benefits of the obtained F1 or Barron spaces have recently been studied in the context of shallow EBMs by Domingo-Enrich et al . ( 2021 ) , who show that they may outperform the RKHS models . In this work , we focus instead on the computational aspects of training such shallow EBMs . Relying on infinite-dimensional Fenchel duality results for the KL-regularized L2 and L∞ regression problems over probability measures ( App . B ) , we recast the maximum likelihood training of F1-EBMs into a min-max problem over measures , and derive a gradient descent-ascent algorithm ( Alg . 1 ) in the associated metric spaces , based on the Wasserstein distance . Crucially , these schemes , defined over idealised parametrisations requiring infinite number of neurons , admit a finite-particle approximation , as in regression or classification . When viewed in terms of particle systems , the dynamics evolve two interacting populations simultaneously : one over the neuron parameters , and the other over the data space ( Sec . 3 ) . Moreover , our proposed algorithm naturally interpolates between the primal and dual formualtions , thanks to the relative time-scale between the minimization and maximization steps , and is even able to interpolate between MLE and Score Matching . This dual algorithm is evaluated experimentally in Sec . 5 in a well-calibrated high-dimensional teacherstudent environment , which allows us to assess our models against the ground-truth , and test the effect of input dimension . Our experiments confirm that the dual algorithm converges significatively faster than the primal one , suggesting that dual updates might bypass metastability despite high-dimensional and non-convex energy landscapes . Related work . Our work is based on general Fenchel duality results ( App . B ) that may be useful in applications beyond the main focus of this paper ( see App . D ) . These theorems are a generalization of results stated in the compact case in Domingo-Enrich et al . ( 2021 ) in their Appendix D. Similar duality results have been studied extensively in the area of maximum entropy ( maxent ) models ( reviewed in Ch . 12 of Mohri et al . ( 2012 ) ) . The first maxent duality principle was due to Jaynes ( 1957 ) . Maxent models have been applied since the 1990s in natural language processing and in species habitat modeling among others , and studied theoretically especially since the 2000s ( Altun & Smola , 2006 ; Dudı́k et al. , 2007 ) . Recently Dai et al . ( 2019a ) leveraged duality arguments in the context of maximum likelihood EBMs , although in a form different from ours . Their duality result works in the more restrictive setting of “ lazy ” energies lying in RKHS balls and probability measures with L2 densities , and they derive it directly from a general theorem that works for reflexive Banach spaces ( Ekeland & Temam ( 1999 ) , Ch . 6 , Thm . 2.1 ) . Our Fenchel duality results , which work for Borel probability measures and feature-learning ( F1 ) energies , are more general because we must rely on measure spaces , which are non-reflexive Banach spaces . Their algorithm is also different : they do not evolve generated samples , but rather use a transport parametrization of the energy . Dai et al . ( 2019b ) expand the work ( Dai et al. , 2019a ) combining it with Hamiltonian Monte Carlo . A precursor of modern machine learning EBMs were restricted Boltzmann machines ( RBMs ) , first trained via contrastive divergence or CD ( Hinton , 2002 ) - which estimates the gradient of the loglikelihood via approximate MCMC samples of the trained model . It later led to maximum likelihood training of EBMs ( see e.g . Xie et al . ( 2016 ; 2017 ) ; Du & Mordatch ( 2019 ) among many others ) . A popular variant of CD is persistent contrastive divergence or PCD ( Tieleman , 2008 ; Tieleman & Hinton , 2009 ) , in which the MCMC samples are evolved and reused over gradient computations to be progressively equilibrated . Drawing a comparison with our work , our dual F1-EBM training algorithm resembles PCD in that both evolve a set of samples over training iterations . A vast array of EBM losses alternative to maximum likelihood have been developed in recent years ( Song & Kingma , 2021 ) with the goal of avoiding the MCMC procedure , which may be costly for non-convex densities . Some successful ones are score matching ( Hyvärinen , 2005 ) and related methods such as denoising score matching ( Vincent , 2011 ) . Building on these , recent works have achieved state of the art image generation ( Song & Ermon , 2019 ; 2020 ; Ho et al. , 2020 ; Song et al. , 2021 ) . We derive the score matching algorithm for F1 energies and show a continuum of algorithms interpolating between dual maximum likelihood and score matching training . Finally , our work ( in particular App . C ) also has links with maximum mean discrepancy ( MMD ) flows . MMDs are probability metrics that were first introduced in ( Gretton et al. , 2007 ; 2012 ) for kernel two-sample tests , and that have enjoyed ample success with the advent of deep-learning-based generative modeling as discriminating metrics ( Li et al. , 2015 ; Dziugaite et al. , 2015 ; Li et al. , 2017 ) . Among the MMD literature , the closest work to ours is ( Arbel et al. , 2019 ) , which study theoretically the convergence of unregularized MMD gradient flow ( our equation ( 26 ) with β̃−1 = 0 ) . In their experiments , they observe that noisy updates ( β̃−1 > 0 ) are needed for good generalization . Our work shows that their algorithm is exactly training maximum likelihood EBMs energies in an RKHS ball of radius that depends on the noise level ( see App . C ) . 2 BACKGROUND AND SETUP . In this section , we provide preliminary background on the functional spaces associated to overparametrized two-layer networks , and on EBMs and their training losses . Notation . If V is a normed vector space , we use BV ( β ) to denote the closed ball of V of radius β , and BV : = BV ( 1 ) for the unit ball . If K denotes a subset of the Euclidean space , P ( K ) is the set of Borel probability measures , M ( K ) is the space of Radon ( i.e . signed and finite ) measures , and M+ ( K ) is the set of non-negative Radon measures . If γ is a Radon measure over K , then ‖γ‖TV = ∫ K d|γ| is the total variation ( TV ) norm of γ , which turns M ( K ) into a Banach space . Throughout the paper , and unless otherwise specified , σ : R → R denotes a generic non-linear activation function . We use ( · ) + : R → R to denote the ReLU activation , defined as ( z ) + = max { z , 0 } . We use τ to denote a fixed base probability measure , possibly used with a subindex to specify the space it is defined over . We use Sd ⊆ Rd+1 for the d-dimensional hypersphere and log for the natural logarithm . We denote the Lebesgue measure by λ . Given two probability measures ν , ν′ ∈ P ( K ) , DKL ( ν‖ν′ ) = ∫ K log dνdν′ dν denotes the KL divergence from ν ′ to ν and H ( ν , ν′ ) = − ∫ K log ( dν ′ dτ ) dν is the cross-entropy . 2.1 OVERPARAMETRIZED TWO-LAYER NEURAL NETWORK SPACES . In this work , we will focus on dense function approximation classes generated by overparametrized shallow neural networks . One can distinguish two canonical models , depending on the asymptotic scaling regime . For further background on these regimes , we refer the reader to Chizat et al . ( 2019 ) . Kernel regime . Let X ⊆ Rd1 , Θ ⊆ Rd2 , ϕ : X ×Θ→ R , and τΘ be a fixed base probability measure over Θ . We define F2 as the reproducing kernel Hilbert space ( RKHS ) of functions f : X → R such that for some h ∈ L2 ( Θ , τΘ ) , we have that , for all x ∈ X , f ( x ) = ∫ Θ ϕ ( x , θ ) h ( θ ) dτΘ ( θ ) . The RKHS norm of F2 is defined as ‖f‖F2 = inf { ‖h‖L2 ( Θ ) | f ( · ) = ∫ Θ ϕ ( · , θ ) h ( θ ) dτΘ ( θ ) } where ‖h‖2L2 ( Θ ) : = ∫ Θ |h ( θ ) |2 dτΘ ( θ ) ( c.f . Bach ( 2017 ) ) . As an RKHS , the kernel of F2 is k ( x , y ) = ∫ Θ ϕ ( x , θ ) ϕ ( x , θ ) dτΘ ( θ ) . ( 1 ) Feature-learning regime . Set X , Θ and ϕ as in the previous paragraph , and define F1 as the Banach space of functions f : X → R such that , for some Radon measure γ ∈ M ( Θ ) , for all x ∈ X we have f ( x ) = ∫ Θ ϕ ( x , θ ) dγ ( θ ) . We define the norm of F1 as ‖f‖F1 = inf { ‖γ‖TV |f ( · ) = ∫ Θ ϕ ( · , θ ) dγ ( θ ) } . This construction was introduced by Bach ( 2017 ) , who first used the notation F1 and focused in particular on the case X ⊆ Rd , Θ = Sd and ϕ ( x , θ ) = ReLuk ( 〈 ( x , 1 ) , θ〉 ) for some k ∈ Z+ . This space is also known by the name of Barron space ( E et al. , 2019 ; E & Wojtowytsch , 2020 ) in reference to the classic work ( Barron , 1993 ) . Remark that since ‖h‖L1 ( Θ ) = ∫ Θ |h ( θ ) | dτΘ ( θ ) ≤ ( ∫ Θ |h ( θ ) |2 dτΘ ( θ ) ) 1/2 = ‖h‖L2 ( Θ ) by the Cauchy-Schwarz inequality , we have F2 ⊂ F1 : in particular finite-width neural networks belong to F1 but not to F2 ( Bach , 2017 ) . The TV norm in F1 acts as a sparsity-promoting penalty , which encourages the selection of few well-chosen neurons and may lead to favorable adaptivity properties when the target has a low-dimensional structure .

Summary: This paper derives a Fenchel duality formulation of the maximum likelihood loss of F1-EBMs, which turns the optimization into a min-max problem on probability measures over the sample space. A dual algorithm with mean-field dynamics is proposed to solve this problem. Using the dual aogrithm they also draw a connection between maximum-likelihood training and score matching. Simple numerical experiments on two-layer ReLU network show that the algorithm converges much faster than the original MLE.