summaries: 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.

summaries: The authors formulated the Fenchel dual problems to maximum likelihood and score matching training applied to energy-based models defined by shallow neural networks. The authors then proposed practical algorithms based on the Hahn decomposition and compared this to various predecessor algorithms. An interesting part is that the learning dynamics can interpolate between maximum likelihood and score matching. There are also rich results in the Appendix.

summaries: Two approaches exist for learning a Gibbs measure: either by MLE (minimizing the KL divergence) or by score matching (maximizing the Hyvarinen score). With the same goal to minimizing the KL divergence, the author proposes an alternative training approach by augmenting an auxiliary $\gamma$ measure and thereby optimizing a more tractable dual problem. The author claims that this new training method has a quicker convergence rate than the primary problem.

summaries: 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.

abstractText: 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.

authors:
Michael Arbel
Arthur Gretton

title: Kernel conditional exponential family

venue: In Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics,

authors:
Michael Arbel
Anna Korba
Adil Salim
Arthur Gretton

title: Maximum mean discrepancy gradient flow

venue: In Advances in Neural Information Processing Systems,

title: Breaking the curse of dimensionality with convex neural networks

venue: Journal of Machine Learning Research,

title: Universal approximation bounds for superpositions of a sigmoidal function

venue: Information Theory, IEEE Transactions on,

authors:
Alain Berlinet
Christine Thomas-Agnan

title: Reproducing Kernel Hilbert Space in Probability and Statistics

authors:
Jonathan Borwein
Qiji Zhu

title: Techniques of Variational Analysis

venue: CMS Books in Mathematics. Springer-Verlag New York,

authors:
Zhengdao Chen
Grant M. Rotskoff
Joan Bruna
Eric Vanden-Eijnden

title: A dynamical central limit theorem for shallow neural networks, 2020

authors:
Lenaic Chizat
Francis Bach

title: On the global convergence of gradient descent for overparameterized models using optimal transport

venue: In Advances in neural information processing systems,

authors:
Lénaı̈c Chizat
Edouard Oyallon
Francis Bach

title: On lazy training in differentiable programming

authors:
Lénaı̈c Chizat
Gabriel Peyré
Bernhard Schmitzer
François-Xavier Vialard

title: Unbalanced optimal transport: Dynamic and kantorovich formulations

venue: Journal of Functional Analysis,

authors:
Bo Dai
Hanjun Dai
Arthur Gretton
Le Song
Dale Schuurmans
Niao He

title: Kernel exponential family estimation via doubly dual embedding

venue: In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics,

authors:
Bo Dai
Zhen Liu
Hanjun Dai
Niao He
Arthur Gretton
Le Song
Dale Schuurmans

title: Exponential family estimation via adversarial dynamics embedding

authors:
Prafulla Dhariwal
Alex Nichol

title: Diffusion models beat gans on image synthesis

venue: arXiv preprint arXiv:2105.05233,

authors:
Carles Domingo-Enrich
Samy Jelassi
Arthur Mensch
Grant Rotskoff
Joan Bruna

title: A meanfield analysis of two-player zero-sum games

authors:
Carles Domingo-Enrich
Alberto Bietti
Eric Vanden-Eijnden
Joan Bruna

title: On energy-based models with overparametrized shallow neural networks, 2021

authors:
Yilun Du
Igor Mordatch

title: Implicit generation and generalization in energy-based models

venue: In Advances in Neural Information Processing Systems (NeurIPS),

authors:
Miroslav Dudı́k
Steven J. Phillips
Robert E. Schapire

title: Maximum entropy density estimation with generalized regularization and an application to species distribution modeling

venue: J. Mach. Learn. Res.,

authors:
Nelson Dunford
Jacob T Schwartz

title: Linear operators. Part I, General theory

venue: Pure and applied mathematics (Interscience series). Interscience,

authors:
Gintare Karolina Dziugaite
Daniel M Roy
Zoubin Ghahramani

title: Training generative neural networks via maximum mean discrepancy optimization

authors:
Weinan E
Stephan Wojtowytsch

title: On the banach spaces associated with multi-layer relu networks: Function representation, approximation theory and gradient descent dynamics, 2020

authors:
Weinan E
Chao Ma
Lei Wu

title: A priori estimates of the population risk for two-layer neural networks

venue: Communications in Mathematical Sciences, 17:1407–1425,

authors:
Ivar Ekeland
Roger Temam

title: Convex analysis and variational problems

venue: Philadelphia, Pa: Society for Industrial and Applied Mathematics,

title: Elementary Principles in Statistical Mechanics: Developed with Especial Reference to the Rational Foundation of Thermodynamics. Cambridge Library Collection - Mathematics

authors:
Arthur Gretton
Karsten M Borgwardt
Malte Rasch
Bernhard Schölkopf
Alex J Smola

title: A kernel method for the two-sample-problem

authors:
Arthur Gretton
Karsten M. Borgwardt
Malte J. Rasch
Bernhard Schölkopf
Alexander Smola

title: A kernel two-sample test

title: Training products of experts by minimizing contrastive divergence

venue: Neural Comput.,

authors:
Jonathan Ho
Ajay Jain
Pieter Abbeel

title: Denoising diffusion probabilistic models

title: Estimation of non-normalized statistical models by score matching

title: Information theory and statistical mechanics

venue: Phys. Rev.,

authors:
Alexia Jolicoeur-Martineau
Rémi Piché-Taillefer
Rémi Tachet des Combes
Ioannis Mitliagkas

title: Adversarial score matching and improved sampling for image generation

venue: arXiv preprint arXiv:2009.05475,

authors:
Zahra Kadkhodaie
Eero P Simoncelli

title: Solving linear inverse problems using the prior implicit in a denoiser

venue: arXiv preprint arXiv:2007.13640,

authors:
Chun-Liang Li
Wei-Cheng Chang
Yu Cheng
Yiming Yang
Barnabas Poczos

title: Mmd gan: Towards deeper understanding of moment matching network

authors:
Yujia Li
Kevin Swersky
Richard Zemel

title: Generative moment matching networks

venue: In ICML,

title: A class of markov processes associated with nonlinear parabolic equations

venue: Proceedings of the National Academy of Sciences of the United States of America, 56:1907–11,

authors:
Song Mei
Andrea Montanari
Phan-Minh Nguyen

title: A mean field view of the landscape of twolayer neural networks

venue: Proceedings of the National Academy of Sciences,

authors:
Mehryar Mohri
Afshin Rostamizadeh
Ameet Talwalkar

title: Foundations of Machine Learning

title: Zur theorie der gesellschaftsspiele

venue: Mathematische annalen,

title: Random coding strategies for minimum entropy

venue: IEEE Transations on Information Theory,

authors:
Marc Ranzato
Christopher Poultney
Sumit Chopra

title: Efficient learning of sparse representations with an energy-based model

authors:
Grant M Rotskoff
Eric Vanden-Eijnden

title: Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error

venue: arXiv preprint arXiv:1805.00915,

authors:
Grant M Rotskoff
Samy Jelassi
Joan Bruna
Eric Vanden-Eijnden

title: Global convergence of neuron birth-death dynamics

venue: In Proceedings of the 36th International Conference on International Conference on Machine

title: Statistical mechanics: Rigorous results

venue: W.A. Benjamin,

title: On general minimax theorems

venue: Pacific J. Math.,

authors:
Justin Sirignano
Konstantinos Spiliopoulos

title: Mean field analysis of neural networks: A central limit theorem

venue: Stochastic Processes and their Applications,

authors:
Yang Song
Stefano Ermon

title: Generative modeling by estimating gradients of the data distribution

venue: arXiv preprint arXiv:1907.05600,

title: Improved techniques for training score-based generative models

authors:
Yang Song
Diederik P. Kingma

title: How to train your energy-based models, 2021

authors:
Yang Song
Jascha Sohl-Dickstein
Diederik P. Kingma
Abhishek Kumar
Stefano Ermon
Ben Poole

title: Score-based generative modeling through stochastic differential equations

venue: In International Conference on Learning Representations (ICLR

authors:
Danica J. Sutherland
Heiko Strathmann
Michael Arbel
Arthur Gretton

title: Efficient and principled score estimation with nyström kernel exponential families

venue: In Proceedings of the TwentyFirst International Conference on Artificial Intelligence and Statistics,

title: Topics in propagation of chaos

venue: Ecole d’Eté de Probabilités de Saint-Flour XIX —

title: Training restricted boltzmann machines using approximations to the likelihood gradient

venue: In Proceedings of the 25th International Conference on Machine Learning,

authors:
Tijmen Tieleman
Geoffrey Hinton

title: Using fast weights to improve persistent contrastive divergence

venue: In Proceedings of the 26th Annual International Conference on Machine Learning,

title: A connection between score matching and denoising autoencoders

venue: Neural Computation,

authors:
Martin Wainwright
Michael Jordan

title: Graphical models, exponential families, and variational inference

venue: Foundations and Trends in Machine Learning, 1:1–305,

title: High-Dimensional Statistics: A Non-Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics

authors:
Jianwen Xie
Yang Lu
Song-Chun Zhu
Yingnian Wu

title: A theory of generative convnet

venue: In Proceedings of The 33rd International Conference on Machine Learning,

authors:
Jianwen Xie
Song Zhu
Yingnian Wu

title: Under review as a conference paper at ICLR

heading: 1 INTRODUCTION

text:
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).

heading: 2 BACKGROUND AND SETUP

text:
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.

heading: 2.1 OVERPARAMETRIZED TWO-LAYER NEURAL NETWORK SPACES

text:
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.

heading: 2.2 EBMS AND TRAINING LOSSES

text:
Consider a measurable set X ⊆ Rd1 with a fixed base probability measure τX ∈ P(X ). If F is a class of functions (or energies) mapping X to R, for any f ∈ F we can define the probability measure νf as a Gibbs measure with density:
dνβf dτX (x) := Z−1βf e −βf(x) with Zβf := ∫ X e−βf(y)dτX (y) ,
where dνβf/dτX is the Radon-Nikodym derivative of νβf and Zβf is the partition function. The parameter β > 0 is the inverse temperature. We could merge β into F by considering the function class {βf |f ∈ F}, but we have decided to keep them separate to showcase the dependency on β. Gibbs measures are the cornerstone of statistical physics since the seminal works of Boltzmann and Gibbs. Beyond their widespread use across computational sciences, they have also found their application in Machine Learning, by the name of energy-based models (EBMs), where the energy function is parametrised using e.g. a neural network.
We denote by F1-EBMs the energy-based models for which the energy class F is the unit ball BF1(1) of F1. Notice that the class {βf |f ∈ F} is equal to the ball BF1(β). Such models may be regarded as abstractions of more complex deep EBMs, in that they incorporate feature learning, and they were first studied by Domingo-Enrich et al. (2021), which provide statistical guarantees. They are to be contrasted with F2-EBMs, for which F is the unit ball BF2(1). F2-EBMs, which we study in App. C, have fixed features and showed worse statistical performance in experiments (Domingo-Enrich et al., 2021).
Given samples {xi}ni=1 from a target measure νp, training an EBM consists in selecting the best νβf with energy f ∈ F according to a given criterion. Two such criteria are relevant in this paper.
Maximum likelihood. The maximum likelihood estimator (MLE) is defined as f̂ = arg maxf∈F ∏n i=1 dνβf dτX
(xi), or, equivalently, as the minimizer of the cross-entropy H(νn, νβf ) with the empirical measure νn = 1n ∑n i=1 δxi :
f̂ = arg min f∈F − 1 n n∑ i=1 log ( dνβf dτX (xi) ) = arg min f∈F 1 n n∑ i=1 f(xi) + β −1 logZβf . (2)
The estimated distribution is then simply given by dνβf̂ = Z −1 βf̂ e−βf̂dτX . By observing that DKL(ν||ν′) = H(ν, ν′) − H(ν), where H(ν) := H(ν, ν) is the entropy, minimizing the crossentropy is equivalent to minimizing the KL divergence when the latter is finite. However, such equivalence is only well-defined in the population setting where the empirical measure νn is replaced by its population counterpart νp. An appropriate choice of function class F induces a regularization that prevents the learned Gibbs measure to overfit to the empirical data measure, and presumably approximate νp instead. The MLE enjoys strong statistical properties (Wainwright & Jordan, 2008; Wainwright, 2019), as well as a powerful variational principle as soon as one considers convex function classes (see App. A), but its notorious computational challenges (see Related works section) have motivated alternative approximation metrics to be used to learn EBMs.
Since an arbitrary element f of F1 can be expressed as f(x) = ∫ Θ ϕ(x, θ) dγ(θ), with ‖f‖F1 equal to the infimum of ‖γ‖TV for all such γ, the maximum likelihood function for F = BF1(1) the problem (2) can be restated as fMLE = ∫ Ω ϕ(·, θ)dγMLE(θ) where
γMLE = arg min γ∈M(Θ) ‖γ‖TV≤1
1
n n∑ i=1 ∫ Θ ϕ(xi, θ) dγ(θ) + 1 β log (∫ X exp ( −β ∫ Θ ϕ(x, θ) dγ(θ) ) dτX (x) ) .(3)
Score matching. An important instance of such weaker metrics is given by Score Matching (SM). The SM metric between two absolutely continuous probability measures ν, ν′ is defined as SM(ν, ν′) = ∫ X |∇ log dν dτX (x) − ∇ log dν ′ dτX (x)|2 dν(x). The SM metric is known in information theory as the relative Fisher information. Note that this metric cannot be trivially extended to the case ν = νn, because empirical measures do not have a density with respect to τX . To get around this difficulty, note that, if the target measure νp is absolutely continuous with respect to τX and we denote by fp(x) = −β−1 log dνpdτX (x) its energy, learning an EBM with function class F under the population loss corresponds to solving f̂ = arg minf∈F ∫ X |∇f(x) − ∇fp(x)|
2 dνp(x). The insight from Hyvärinen (2005) is that under regularity conditions on fp, via integration by parts we then have ∫ X |∇f −∇fp|
2 dνp = E{xi}ni=1L(f, νn) + C, where E{xi}ni=1 denotes expectation over the data set, C is a constant in f which is therefore irrelevant, and
L(f, νn) = 1
n n∑ i=1 β−1∆f(xi) + 1 2 |∇f(xi)|2.
In practice, we train an EBM via score matching by solving f̂ = arg minf∈F L(f, νn). Score matching is computationally more tractable than maximum likelihood and performs well in practice (see Related works section). Statistically, its main drawback is that the SM metric is weaker than the KL divergence, and may fail to distinguish distributions in some instances.
The following proposition, proved in App. G, provides the expression for the loss L and the resulting score matching problem for F1-EBMs. Proposition 1. Suppose that X ⊆ Rd1 is a manifold without boundaries. Assume that∫ X |∇xϕ(x, θ) · ∇ dνp dτX
(x)| dτX (x) is upper-bounded by some constant K for all θ ∈ Θ. Assume also that supθ∈Θ ‖∇xϕ(x, θ)‖ < η(x) and that ∫ X |η(x)|
2 dνp(x) <∞. The optimization problem to train EBMs under the score matching loss over the ball BF1(1) gives fSM = ∫ Ω ϕ(·, θ)dγSM(θ) where
γSM = arg min γ∈M(Θ) ‖γ‖TV≤1
∫ Θ ∫ X ( 1 2 ∇xϕ(x, θ) · ∇x ∫ Θ ϕ(x, θ′) dγ(θ′)− β−1∆xϕ(x, θ) ) dνn(x)dγ(θ).(4)
3 DUAL F1-EBM TRAINING VIA MAXIMUM LIKELIHOOD
As a corollary of the duality result from Subsec. B.2, we derive an alternative objective forF1-EBMs trained via maximum likelihood, the original objective being (3) and we develop an algorithm to solve this alternative problem. To this end, we make: Assumption 1. Let ϕ : X×Θ→ R be a continuous function such that eitherX is compact or (i) for any fixed θ ∈ Θ, ϕ(x, θ) ≤ ξ(x) for some strictly positive ξ : X → R, and (ii) ξ(x) + log(ξ(x)) =
o ( − log ( dτX dλ (x) ) − (d1 + ) log ‖x‖2 ) as ‖x‖2 → +∞ for some > 0.
In particular, this assumption holds for ReLU network energies when setting X = Rd1 , Θ = Rd1+1, ϕ(x, θ) = σ(〈(x, 1), θ〉)/‖θ‖ and τX Gaussian (and in many other settings). Theorem 1. Under Assumption 1, the problem (3) is the Fenchel dual of
min ν∈P(X ) max γ∈M(Θ), ‖γ‖TV≤1
β−1DKL(ν||τX ) + ∫
Θ ∫ X ϕ(x, θ) d(ν − νn)(x) dγ(θ). (5)
Moreover, the solution ν? of (70) is precisely the Gibbs measure for the optimal γ? in (3), that is, dν? dτX (x) = 1Zβ exp ( −β ∫ Θ ϕ(x, θ) dγ?(x) ) .
If we replace the F1 ball by the F2 ball, the analogous duality result links the maximum likelihood problem with the entropy regularized MMD flow from Arbel et al. (2019) (see App. C).
Training dynamics on nonnegative measures. Let us write the dynamics to solve (5) that can be discretized in terms of parameters and particles (cf Proposition 2 below). To this end we consider the triple (γ+, γ−, ν) where the nonegative measures γ± are defined through the Hahn decomposition of γ = γ+ − γ−. Then we introduce coupled gradient flows for this triple, in which γ+t and γ−t evolve via a Wasserstein-Fisher-Rao gradient flow (Chizat et al., 2018) and νt evolves via a Wasserstein gradient flow (Santambrogio, 2017):
∂tγ σ t = −ασ∇θ · ( γσt ∇θFt(θ) ) + αγσt ( σFt(θ)−Kt ) , σ = ±1, γσt = γ±t
∂tνt = ∇x · ( νt ( ∇xft(x)− β−1∇ log
dτX dλ
)) + β−1∆xνt,
(6)
where α is a tunable parameter and we defined
Ft(θ) = ∫ X ϕ(x, θ) d(νt−νn)(x), ft(x) = ∫ Θ ϕ(x, θ) (dγ+t − dγ−t )(θ), (7)
Kt = 1‖γ+t ‖TV+‖γ − t ‖TV≥1 ∫ Θ Ft(θ)(dγ + t − dγ−t )(θ).
The initialization of (6) is ν0 = νn and γ±0 = 0 (such that the initial energy is null). The term Kt keeps the total variation of γt below one. The parameter α acts as a relative timescale. Notice that different values of α can potentially lead to different behaviors of the dynamics; setting α 1 would correspond to the primal formulation of maximum likelihood with persistent MCMC samples (as in PCD). In contrast if α 1, γ±t evolves faster than νt and if the optimization is well behaved at all times γt = γ+t − γ−t remains close to minimizing the inner maximization problem of (5) with γ = γt. Initializing ν0 = νn is crucial to avoid the kind of metastabilities that curse the behavior of classical (primal) maximum likelihood EBM training.
Proposition 2 below states that the solution (µt, νt) may be approximated using coupled particle systems (see proof in App. F) and is the basis for Alg. 1. The link between particle systems and measure PDEs is through a classical technique known as propagation of chaos (Sznitman, 1991) and it has been used previously for similar coupled systems in the machine learning literature (DomingoEnrich et al., 2020).
Proposition 2. Let {θ(j)0 }mj=1 be initial features sampled uniformly over Θ, let {σj}mj=1 be uniform samples over {±1} and let {w(j)0 = 1}mj=1 be the initial weight values, which are set to 1. Let {X(i)0 }Ni=1 be the initial “generated” samples, which are chosen i.i.d. uniformly from the target sample set {xi}ni=1. Consider the system of ODEs/SDEs:
dθ (j) t
dt = ασj∇F̃t(θ(j)t ),
dw (j) t
dt = αw
(j) t (σjF̃t(θ (j) t )− K̃t)
dX (i) t = ( −∇f̃t(X(i)t ) + β−1∇ log
dτX dλ (X (i) t )
) dt+ √ 2β−1 dW
(i) t
(8)
where
F̃t(θ) = 1
N N∑ i=1 ϕ(X (i) t , θ)− 1 n n∑ i=1 ϕ(xi, θ), f̃t(x) = 1 m m∑ j=1 σjw (j) t ϕ(x, θ (j) t ),
K̃t = 1∑m j=1 w (j) t ≥m 1 m m∑ j=1 σjw (j) t F̃t(θ (j) t ).
(9)
are the empirical counterparts of the functions in (7). Then the system (8) approximates the measure dynamics. Namely, as m,N →∞:
• the empirical measure γ̂t = 1m ∑m j=1 σjw (j) t δθ(j)t
converges weakly to the solution γt = γ+t −γ−t of (6) with uniform initialization for any finite time interval [0, T ], and
• the empirical measure ν̂t = 1N ∑N i=1 δX(i)t
converges weakly to the solution νt of (6) for any finite time interval [0, T ].
Importantly, the system of ODEs/SDEs in (8) may be solved via forward Euler steps on {θj}mj=1 and {wj}mj=1 (or rather, {logwj}mj=1), and Euler-Maruyama updates on {X (i) 0 }Ni=1. Such a discretization yields Algorithm 1. Algorithm 1 makes use of a tunable parameter pR, which stands for the restart probability and will be discussed in Sec. 4 as a natural way to connect maximum likelihood with score matching. To discretize (8) we set pR = 0, i.e., there are no particle restarts.
Algorithm 1 DualF1-EBM training (pR = 0: maximum likelihood, pR = (sα∧1): score matching) Input: n samples {xi}ni=1 of the target distribution, stepsize s, stepsize ratio α. Initialize unif. features (θ(j)0 ) m j=1 over Θ, weights (w (j) 0 ) m j=1 in [0, 1), signs (σj) m j=1 over {±1}.
Initialize generated samples {X(i)0 }Ni=1 uniformly i.i.d. from {xi}ni=1. for t = 0, . . . , T − 1 do
for i = 1, . . . , N do With probability pR, replaceX (i) t by some uniformly chosen sample in {xi}ni=1 (see Sec. 4).
Sample ζ(i)t from the d1-variate standard Gaussian. Perform Euler-Maruyama update: X(i)t+1 = X (i) t − s(∇f̃t(X (i) t ) + β −1∇ log dτXdλ (X (i) t )) +√
2β−1s ζ (i) t , where f̃t is defined in (9).
end for for j = 1, . . . ,m do
Update θ(j)t+1 = θ (j) t + sασj∇F̃t(θ (j) t ), where F̃t is defined in (9). Update w̃(j)t+1 = w (j) t+1 exp(sασjF̃t(θ (j) t )). Normalize if needed w(j)t+1 = w̃ (j) t+1/max ( m−1 ∑m j′=1 w̃ (j′) t+1, 1 ) .
end for end for Output: samples {X(i)T }Ni=1, energy fT (x) := β m ∑m j=1 σjwjϕ(x, θj).

heading: 4 LINKS BETWEEN MAXIMUM LIKELIHOOD AND SCORE MATCHING

text:
F1-EBMS
In this section we uncover how the score matching loss fits seamlessly as a variant of Alg. 1, in the form of particle restarts. Interestingly, we can modify the PDE (6) in a way that allows us to make a connection with score matching. To this end, let us introduce the following coupled measure PDE:
∂tγ σ t = −ασ∇θ · ( γσt ∇θFt(θ) ) + αγσt ( σFt(θ)−Kt ) , σ = ±1, γσt = γ±t ,
∂tνt = ∇x · ( νt ( ∇xft(x)− β−1∇ log
dτX dλ
)) + β−1∆xνt − α (νt − νn) .
(10)
Remark that the only difference between this equation and the PDE (6) for dual maximum likelihood training is the term−α(νt−νn), which draws νt closer to the empirical target measure νn. We have:
Proposition 3. In the limit α→∞, the equations for γσt in (10) reduce to
∂tγ σ t = σ∇θ · ( γσt ∇θV (γt)(θ) ) − γσt ( σV (γt)(θ)− V̄ (γt) ) , σ = ±1, γσt = γ±t (11)
where γt = γ+t − γ−t , V̄ (γ) = ∫ Θ V (γ)dγ, and V (γ)(θ) is the Frechet derivative of the score matching loss L :M(Θ)→ R defined in (4).
That is, in the large α limit, equation (10) is equivalent to the Wasserstein-Fisher-Rao gradient flow of a loss L which, remarkably, is the score matching loss for F1-EBMs. This means that adding the term −α (νt − νn) to the dual maximum likelihood measure dynamics and letting α → ∞ we recover the score matching dynamics. This additional term can be easily implemented at particle level by replacing each training sample X(i)t by some random target sample in {xi}ni=1 with probability pR = 1 − e−αt = αt + o(t) for every time interval of length t (proof in App. G).
Similar birth-death processes were used in (Rotskoff et al., 2019) in the context of neural network regression. Hence, the score matching scheme corresponds to setting the restart probability pR = sα in Algorithm 1. The restart probability acts as a knob that allows us to interpolate between score matching and maximum likelihood.
In summary, score matching differs from dual maximum likelihood in that the trained measure is being “pulled” towards the target measure at all times via particle restarting. Such constant pulling should be useful to alleviate sampling problems due to metastability issues which may arise with dual maximum likelihood. However, dual maximum likelihood has the upside of providing samples of the learned EBM as a byproduct of training, which score matching does not. A good balance between both algorithms may be to use a restart probability pR between sα ∧ 1 and 0, or even a pR that decreases with time from one value to the other, in such a way that at the beginning of training we avoid metastability issues by restarting the particles frequently, and at an advanced phase we perform little to no restarting to obtain faithful samples. It is also interesting to contrast our approach to score matching with the works Sutherland et al. (2018); Arbel & Gretton (2018), which using different techniques propose algorithms to train EBMs with RKHS energies via score matching. Finally, notice that a particle discretization of the flow (11) yields an alternative straightforward algorithm to train F1-EBMs via score matching; see Subsec. G.1. In Subsec. G.1 we show that this algorithm can be linked directly to Alg. 1 with particle restarts, without recurring to measure arguments.

heading: 5 EXPERIMENTS

text:
Setup. To illustrate Alg. 1 we perform numerical experiments on simple synthetic datasets generated by teacher models with energy f∗(x) = 1J ∑J j=1 w ∗ jσ(〈θ∗j , x〉), with θ∗j ∈ Sd for all j. The training is performed using Alg. 1 with the added detail that both the features θ(j)t and the particles X (i) t are constrained to remain on the sphere by adding a projection step in the update of their positions. The code, figures, and videos on the dynamics can be found in the supplementary material. In the main text we consider two planted teacher neurons (J = 2) with negative output weights w∗1 = w ∗ 2 = −10 in dimension d = 14 and m = 64 neurons for the student model, but we include additional experiments and videos in App. I and supplementary material. We study setups with two different choices of angles between the teacher neurons, which showcase different behaviors:
• Teacher neurons θ∗1 , θ ∗ 2 forming an angle of 2.87 rad (≈ 164 degrees), and output weights w∗1 =
w∗2 = −10. The teacher neurons are almost in opposite directions, and the resulting target distribution is bimodal, as the energy has two local minimizers around θ∗1 and θ ∗ 2 (see Figure 5).
• Teacher neurons θ∗1 , θ ∗ 2 forming an angle of 1.37 rad (≈ 78 degrees). The teacher neurons are
almost orthogonal and the resulting target distribution is monomodal; indeed, when the angle is less than π/2, the target energy has a unique minimizer at the geodesic average between θ∗1 and θ∗2 (see Figure 5 in App. I).
Monitoring convergence. In all our experiments, to monitor convergence we use a testing set of n∗ data points sampled from the teacher distribution: denoting these samples by {x∗i } n∗ i=1, we estimate
the KL divergence from the student to the teacher via log( 1n∗ ∑n∗ i=1 exp(−βft(x∗i ) + βf∗(x∗i ))) +
1 n∗ ∑n∗ i=1(ft(x ∗ i ) + βf ∗(x∗i )) where ft(x) = 1 m ∑m j=1 w (j) t σ(〈θ (j) t , x〉). Similarly, for the score
matching objective we use the estimate 1n∗ ∑n∗ i=1 ∣∣∇xft(x∗i )−∇xf∗(x∗i )∣∣2. Comparison of the primal algorithm and the dual algorithms. We defer the empirical study of tuning the restart probability to App. I, and in this section focus on comparing the dual algorithm for maximum likelihood F1-EBMs (i.e. with pR = 0) to the classical (primal) algorithm, which was the algorithm used in the experiments of Domingo-Enrich et al. (2021). The primal algorithm corresponds to Alg. 1 with α 1, while the dual algorithm uses α 1. To obtain a principled comparison of the two settings where numerical errors do not blow up, we set s to be the stepsize for the fastest process (particle evolution for the primal, neuron evolution for the dual), and min(α, 1/α)s the stepsize for the slow process.
The results are shown in Figure 6 for the two angle configurations between teacher neurons. We observe that the dual algorithm is several orders of magnitude faster at reaching KL and SM values close to the final ones, which showcases the main advantage of the dual approach. For the monomodal distribution, the final values obtained by the dual algorithm are slightly better than for the primal algorithm; for the bimodal distribution, the converse happens and the convergence is slower for both algorithms, most likely due to metastability. Interestingly, the decrease of the performance metrics seems to stall as soon as the hard F1-norm threshold is reached.

heading: 6 DISCUSSION AND OUTLOOK

text:
In this work we leverage a Fenchel duality result to recast the maximum likelihood loss for F1EBMs into a min-max problem on probability measures over the sample space. We provide meanfield dynamics at the measure level to solve this problem, which lead to a dual algorithm (Alg. 1) after discretization. We observe that if we restart particles at random target samples throughout training, we get an algorithm which is equivalent to training under the score matching loss. We perform experiments in which we learn planted distributions with two-layer ReLU networks, and we observe empirically that our dual algorithm is much faster than the classical one.
At theoretical level, one direction for future work is to obtain convergence results for the dynamics (6) and (10). Domingo-Enrich et al. (2020) study similar coupled Wasserstein-Fisher-Rao gradient flows, but their results only work in the case of weight learning rates much larger than position learning rates. We hypothesize that the additional term −α (νt − νn), which keeps νt close to νn, might help in the analysis.
At the numerical level, it would be interesting to further test the variant of Algorithm 1 with annealed pR decreasing in time, to understand under which parameter setup it captures the best features of maximum likelihood and score matching. One could also test Algorithm 1 using deeper neural architectures: while the analysis is more complicated in this case, the scheme itself can be straightforwardly generalized to deep networks.