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 studies the training of Energy Based Models (EBM). The “standard” technique is maximum likelihood maximization of the observed dataset. While theoretically sound, there are practical difficulties in simulating the MCMC dynamics due to large basins in the energy landscape. For this reason a class of alternative methods, like Score matching, have been explored in the literature. Score matching, however, is deemed to suffer in statistical power.  This work explores how, leveraging known results about Fenchel dualities, it is possible to connect the two training modalities (Maximum Likelihood and Score based) on a continuum, and proposes a practical algorithm based on such considerations. The authors consider for their theoretical analysis of the expressiveness of the models,  shallow neural networks in the lazy and kernel regime (spaces F_1 and F_2 respectively).  Section 2.1 introduces the two considered spaces, referring to the works of Chizat and Bach.   An energy based model is defined at the beginning of section 2.2 by means of the Radon Nikodym derivative. Training an EBM model via maximum likelihood corresponds to maximizing eq. (2), that, when considering networks in F_1, can be rewritten as the expression (3). At the end of Section 2.2. score matching is rigorously defined and the relationship between theoretical and empirical implementations (Hyvarinen) is clarified.  Section 3 is the core of the paper. If Assumption 1 holds (growth conditions), then Thm 1 states that eq. (3) is the Fenchel dual of (5). As (5) is a static functional, the authors propose to use results from Chizat and Santambrogio to rewrite it as the dynamical system (eq. (6))  together with the definitions in (7). Importantly, this formulation contains \alpha, a free parameter >0, that determines relative time-scales over which functions \gamma and \nu are updated. System (6) is still “algorithmically” unsolvable. The authors then propose to use the two following approximations: first, the mean field technique known as propagation of chaos, is used to switch from (6)  to the system (8), where gamma is substituted by neural network  parameters and \nu by and SDE. Then the continuous time system (8) is discretized using Euler-Maruyama. Algorithm 1 is the result of the process.  Section 4 explores how, in the infinite alpha regime, the dynamical system (6) is equivalent to score matching. By means of eq. (11) and proposition (3) the technical analysis is performed. Intuitively, parameters are updated infinitely faster than the samples.  Section 5 explores the implementation of Algorithm 1 on a synthetic dataset. The authors consider two cases (the two teacher neurons aligned to 164 and 78 degrees respectively). Comparison of the primal and dual (alpha<<1, alpha>>1) shows that the dual performs much better than the primal.  The conclusions stress the theoretical contribution based on Fenchel duality and propose as future work the exploration of convergence rates of (6) and (10) and numerical investigation of the restarting probability p_R. In addition, the authors suggest that “porting” the presented work to deeper network architectures (e.g. to work on realistic dataset, such as images) is straightforward from the practical point of view, and unpractical from the theoretical point of view.