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: This paper proposes to use max-product belief propagation to provide an approximate solution to the MAP problem encountered in the perturb-and-MAP sampling method. The paper argues that, despite its lack of convergence guarantees, max-product is faster (due in large part to its parallelizability) and better for sampling than linear-programming based methods (the standard approach) in problems with multiple modes. Experiments in a number of synthetic tasks show that the proposed approach, perturb-and-max-product (PMP), is an efficient and effective sampling method that can also be used for parameter estimation of discrete energy-based models in a black-box manner.