Perturb-and-max-product: Sampling and learning in discrete energy-based models

1 Introduction . In this work , we concern ourselves with the problem of black-box parameter estimation for a very general class of discrete energy-based models ( EBMs ) , possibly with hidden variables . The golden measure of estimation quality is the likelihood of the parameters given the data , according to the model . However , the likelihood of EBMs is not computable in polynomial time due to the partition function , which in turn makes its optimization ( parameter estimation ) difficult . This is an active research problem , with multiple approaches having been proposed over the last decades [ 37 , 29 , 33 , 2 , 13 , 14 , 12 , 37 , 17 , 38 , 19 ] . These methods typically fall into one or more of the following categories ( a ) they depart from the maximum likelihood ( ML ) criterion to obtain a tractable objective , but in doing so they fail when the data does not match the model ( e.g. , pseudolikelihood [ 13 ] ) ; ( b ) they can only be applied to fully observed models ( e.g. , piecewise training [ 29 ] ) ; ( c ) they fail catastrophically for simple models ( e.g. , Bethe approximation [ 12 ] ) ; ( d ) they require additional designs ( such as a custom encoder ) for new models , i.e. , they are not black-box ( e.g , Advil [ 19 ] ) . A general approach to discrete EBM learning that does not fall in any of those pitfalls is sampling : as we review in Section 2.1 , sampling from an EBM results in efficient learning . Exact sampling is unfeasible in general , but approximate sampling works well enough in multiple models . Gibbs sampling is the most popular because it is simple , general , and can be applied in a black-box manner : the sampler can be generated directly from the model , even automatically [ 5 ] . Its main limitation is that for some models , mixing can be too slow to be usable . Recently the Gibbs-with-gradients ( GWG ) [ 8 ] sampler was proposed , which significantly improves mixing speed , enabling tractability of more models . A completely different approach to learning is Perturb-and-MAP [ 22 ] , but this approach has not seen widespread adoption because the general MAP problem could not be solved fast enough or with good enough quality , see Section 2.4.1 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this work , we show that max-product , a relatively under-utilized algorithm , can be used to perform good enough MAP inference in many non-tractable cases very fast . Thus , perturb-and-max-product not only expands the applicability of perturb-and-MAP to new scenarios ( we show that it can sample in highly entangled and multimodal EBMs in which the state-of-the-art GWG and traditional Perturb-and-MAP fail ) but can also be much faster in practice than the state-of-the-art GWG . 2 Background . We review here ML learning , Perturb-and-MAP , belief propagation , and how they connect . 2.1 Sampling and learning in discrete EBMs . Notation : Consider an EBM with variables x described by a set of A factors { φa ( xa ; θa ) = θ > a φa ( xa ) } Aa=1 and I unary terms { φi ( xi ; θi ) = θ > i φi ( xi ) } Ii=1 : xa is the subset of variables used in factor a , θa is a vector of factor parameters and φa ( xa ; θa ) is a vector of factor sufficient statistics . For the unary terms , the sufficient statistics φi ( xi ; θi ) are the indicator functions , which are set to 1 when the variable takes a given value and to zero otherwise . The energy of the model can be expressed as E ( x ) = − ∑A a=1 φa ( xa ; θa ) − ∑I i=1 φi ( xi ; θi ) or , collecting the parameters and sufficient statistics in corresponding vectors , E ( x ) = −Θ > a Φa ( x ) −Θ > i Φi ( x ) = −Θ > Φ ( x ) . Given a discrete data set D ≡ { x ( n ) } Ni=1 , where x ∈ { 0 , . . . , M − 1 } D , the log-likelihood of the parameters for a generic EBM in which some variables ( xh ) are hidden can be expressed as L ( Θ|D ) = 1 N N∑ n=1 log p ( x ( n ) ; Θ ) = 1 N N∑ n=1 log ∑ xh exp ( Θ > Φ ( x ( n ) , xh ) ) − log ∑ x exp ( Θ > Φ ( x ) ) ( 1 ) The learning problem can be described as efficiently maximizing the key quantity Eq . ( 1 ) . However , Eq . ( 1 ) can only be computed for small models , due to the summation over x ( the partition function ) , which involves an exponential number of summands . In order to use gradient ascent methods , we need the gradient of Eq . ( 1 ) w.r.t . Θ , which is∇ΘL ( Θ|D ) = 1N ∑N n=1〈Φ ( x ( n ) , xh ) 〉p ( xh|x ( n ) ; Θ ) − 〈Φ ( x ) 〉p ( x ; Θ ) . From this gradient we can make two standard observations : First , at a maximum of Eq . ( 1 ) , the expectation of the sufficient statistics according to the data and the model match , since the gradient needs to be zero . Second , sampling enables efficient learning : we can efficiently approximate the expected sufficient statistics , with no bias , by sampling from the model ( both unconditionally and conditioned on x ( n ) ) . Exact sampling is also unfeasible , so we turn to approximate sampling . 2.2 Perturb-and-MAP . In [ 22 ] , a family of samplers , including the following convenient form , was proposed : x ( s ) = argmax x { Θ > a Φa ( x ) + ( Θi + ε ( s ) ) > Φi ( x ) } ( 2 ) where ε ( s ) is a random perturbation vector with as many entries as Θi ( i.e. , ∑ i |xi| entries , where |xi| is the cardinality of variable xi ) . This means that to generate a sample , first we draw ε ( s ) from some distribution , use it to additively perturb the unary terms , and then we solve a standard maximum a posteriori ( MAP ) problem on the perturbed model . While MAP inference is NP-complete , it is easier than partition function computation which is # P-complete . In fact , in log-supermodular models , such as attractive Ising models , the exact MAP can be computed using graph-cuts , but the partition function still needs exponential time . The distribution of the samples { x ( s ) } will of course depend on the perturbation density . We exclusively use independent , zero-mean densities ε ∼ Gumbel ( ε ; −c , 1 ) , where c ≈ 0.5772 is the Euler-Mascheroni constant . With this choice , models with only unary terms result in the exact Gibbs distribution p ( x ( s ) ) ∝ exp ( Θ > Φ ( x ( s ) ) ) ( see Lemma 1 , [ 22 ] ) . More realistic models containing non-unary terms , will in general not match the Gibbs density , but will approximate it . The Perturb-and-MAP approach also allows for the computation of an upper bound on the partition function , log ∑ x exp ( Θ > Φ ( x ) ) ≤ 〈 maxx ( Θ > a Φa ( x ) + ( Θi + ε ) > Φi ( x ) ) 〉 p ( ε ) , ( see Corollary 1 , [ 10 ] ) . If we use this approximation in the log-likelihood ( 1 ) , we get a much easier target for optimization L̂ ( Θ|D ) = 1 N N∑ n=1 〈 max xh ( Θ > a Φa ( x ( n ) , xh ) + ( Θi + ε ) > Φi ( x ( n ) , xh ) ) 〉 p ( ε ) − 〈 max x ( Θ > a Φa ( x ) + ( Θi + ε ) > Φi ( x ) ) 〉 p ( ε ) . ( 3 ) whose gradient is ∇ΘL̂ ( Θ|D ) = 1N ∑N n=1 Φ ( x ( n ) , xh ) p̂ ( xh|x ( n ) ; Θ ) − 〈Φ ( x ) 〉p̂ ( x ; Θ ) , where p̂ ( · ; Θ ) is a ( possibly conditional ) sampler of the form of Eq . ( 2 ) . Note that this gradient has the same form as the gradient of the true likelihood , but with the exact sampler replaced by the approximate one . So this can be thought of as the approximate gradient of the true log-likelihood , or as the true gradient of the approximate log-likelihood . The expectations in Eq . ( 3 ) and its gradient can be easily approximated , without bias , using the sample mean . Thus , we can use stochastic gradient descent ( SGD ) and minibatches , to get a generally applicable technique for graphical models if we can solve the MAP problem efficiently . So far , the application of perturb-and-MAP has been restricted to models whose structure allowed exact MAP solutions or allowed LP relaxations of the MAP problem to produce accurate results . As we show in Section 5.2 , even simple Ising models violate these assumptions . 2.3 The Bethe approximation and belief propagation . To make perturb-and-MAP more widely useful , we need a more general ( even if approximate ) approach to MAP inference . Here we describe a way to achieve this leveraging the Bethe free energy , which uses a variational formulation and a modified entropy to approximate the partition function . Using an arbitrary variational distribution q ( x ) we can lower bound the partition function using the expectation of the free energy . At the maximum with respect to this q ( x ) , we obtain the equality : β−1 log ∑ x exp ( βΘ > Φ ( x ) ) = max q ( x ) { 〈Θ > Φ ( x ) 〉q ( x ) + β−1H ( q ( x ) ) } ( 4 ) whereH ( q ( x ) ) = − ∑ x q ( x ) log q ( x ) is the entropy and β is an inverse temperature parameter . At β = 1 , the l.h.s . of Eq . ( 4 ) becomes the log-partition function . At β →∞ and with perturbing Θ , the problem posed by Eq . ( 4 ) matches the one in Eq . ( 2 ) . We keep this temperature parameter throughout our derivation . The maximization in ( 4 ) is still intractable , so we introduce two approximations and we obtain max qa ( xa ) , qi ( xi ) A∑ a=1 〈φa ( xa ; θa ) 〉qa ( xa ) + I∑ i=1 〈φi ( xi ; θi ) 〉qi ( xi ) + β−1HB ( qa ( xa ) , qi ( xi ) ) s.t . ∑ xa\i qa ( xa ) = qi ( xi ) , ∑ xi qi ( xi ) = 1 , qa ( xa ) ≥ 0 , which is the Bethe approximation of the log-partition function [ 39 ] . These approximations are ( a ) the use of the Bethe free entropy HB ( qa ( xa ) , qi ( xi ) ) = ∑ aH ( qa ( xa ) ) + ∑ i ( 1 − fi ) H ( qi ( xi ) ) ( which is exact for trees , approximate for loopy models ) instead of the exact entropy ; and ( b ) the use of locally consistent pseudo-marginals qa ( xa ) and qi ( xi ) , which might not correspond to any joint density q ( x ) . To convert this into an unconstrained optimization problem , we move to the dual . The dual variables are called “ messages ” , since there is one of them associated to each factor-variable pair . See [ 40 ] for details . Setting the gradient to 0 , we recover the belief propagation fixed-point equations mi→a ( xi ) = φi ( xi ; θi ) + ∑ b∈nb ( i ) \a mb→i ( xi ) ma→i ( xi ) = β −1 log ∑ xk\i eβ ( φa ( xa ; θa ) + ∑ k∈nb ( a ) \imk→a ( xk ) ) , where nb ( · ) denotes the neighbors of a factor or variable . Updates must proceed sequentially , factor by factor . Convergence to a fixed point is not guaranteed for loopy models . We recover the primal quantities from the messages . In particular , the unary marginals are qi ( xi ) ∝ exp ( β ( φi ( xi ; θi ) +∑ b∈nb ( i ) mb→i ( xi ) ) . The recovered primal value in Eq . ( 5 ) , for β = 1 , is an approximation of the log-partition function , but can not be used for learning in general , see [ 23 , 12 ] , and Section 5.1 .

Summary: Perturb & Max is an established method for sampling (and therefore learning) a certain class of energy-based models, namely latent-variable models whose joint distribution over visibles & latents is exponential family. However, the max operation poses computational difficulties and existing approaches based on LP solvers are rather limited. This paper proposes an alternative approach, called PMP, that is based on max-product messaging passing. PMP shows favourable empirical performance compared to LP-based solvers, as well as a recent state-of-the-art Gibbs-with-Gradients.