Deep learning via message passing algorithms based on belief propagation

1 INTRODUCTION . Belief Propagation is a method for computing marginals and entropies in probabilistic inference problems ( Bethe , 1935 ; Peierls , 1936 ; Gallager , 1962 ; Pearl , 1982 ) . These include optimization problems as well once they are written as zero temperature limit of a Gibbs distribution that uses the cost function as energy . Learning is one particular case , in which one wants to minimize a cost which is a data dependent loss function . These problems are generally intractable and message-passing techniques have been particularly successful at providing principled approximations through efficient distributed computations . A particularly compact representation of inference/optimization problems that is used to build massage-passing algorithms is provided by factor graphs . A factor graph is a bipartite graph composed of variables nodes and factor nodes expressing the interactions among variables . Belief Propagation is exact for tree-like factor graphs ( Yedidia et al. , 2003 ) ) , where the Gibbs distribution is naturally factorized , whereas it is approximate for graphs with loops . Still , loopy BP is routinely used with success in many real world applications ranging from error correcting codes , vision , clustering , just to mention a few . In all these problems , loops are indeed present in the factor graph and yet the variables are weakly correlated at long range and BP gives good results . A field in which BP has a long history is the statistical physics of disordered systems where it is known as Cavity Method ( Mézard et al. , 1987 ) . It has been used to study the typical properties of spin glass models which represent binary variables interacting through random interactions over a given graph . It is very well known that in spin glass models defined on complete graphs and in locally tree-like random graphs , which are both loopy , the weak correlation conditions between variables may hold and BP give asymptotic exact results ( Mézard & Montanari , 2009 ) . Here we will mostly focus on neural networks ±1 binary weights and sign activation functions , for which the messages and the marginals can be described simply by the difference between the probabilities associated with the +1 and -1 states , the so called magnetizations . The effectiveness of BP for deep learning has never been numerically tested in a systematic way , however there is clear evidence that the weak correlation decay condition does not hold and thus BP convergence and approximation quality is unpredictable . In this paper we explore the effectiveness of a variant of BP that has shown excellent convergence properties in hard optimization problems and in non-convex shallow networks . It goes under the name of focusing BP ( fBP ) and is based on a probability distribution , a likelihood , that focuses on highly entropic wide minima , neglecting the contribution to marginals from narrow minima even when they are the majority ( and hence dominate the Gibbs distribution ) . This version of BP is thus expected to give good results only in models that have such wide entropic minima as part of their energy landscape . As discussed in ( Baldassi et al. , 2016a ) , a simple way to define fBP is to add a `` reinforcement '' term to the BP equations : an iteration-dependent local field is introduced for each variable , with an intensity proportional to its marginal probability computed in the previous iteration step . This field is gradually increased until the entire system becomes fully biased on a configuration . The first version of reinforced BP was introduced in ( Braunstein & Zecchina , 2006 ) as a heuristic algorithm to solve the learning problem in shallow binary networks . Baldassi et al . ( 2016a ) showed that this version of BP is a limiting case of fBP , i.e. , BP equations written for a likelihood that uses the local entropy function instead of the error ( energy ) loss function . As discussed in depth in that study , one way to introduce a likelihood that focuses on highly entropic regions is to create y coupled replicas of the original system . fBP equations are obtained as BP equations for the replicated system . It turns out that the fBP equations are identical to the BP equations for the original system with the only addition of a self-reinforcing term in the message passing scheme . The fBP algorithm can be used as a solver by gradually increasing the effect of the reinforcement : one can control the size of the regions over which the fBP equations estimate the marginals by tuning the parameters that appear in the expression of the reinforcement , until the high entropy regions reduce to a single configuration . Interestingly , by keeping the size of the high entropy region fixed , the fBP fixed point allows one to estimate the marginals and entropy relative to the region . In this work , we present and adapt to GPU computation a family of fBP inspired message passing algorithms that are capable of training multi-layer neural networks with generalization performance and computational speed comparable to SGD . This is the first work that shows that learning by message passing in deep neural networks 1 ) is possible and 2 ) is a viable alternative to SGD . Our version of fBP adds the reinforcement term at each mini-batch step in what we call the Posterioras-Prior ( PasP ) rule . Furthermore , using the message-passing algorithm not as a solver but as an estimator of marginals allows us to make locally Bayesian predictions , averaging the predictions over the approximate posterior . The resulting generalization error is significantly better than those of the solver , showing that , although approximate , the marginals of the weights estimated by messagepassing retain useful information . Consistently with the assumptions underlying fBP , we find that the solutions provided by the message passing algorithms belong to flat entropic regions of the loss landscape and have good performance in continual learning tasks and on sparse networks as well . We also remark that our PasP update scheme is of independent interest and can be combined with different posterior approximation techniques . The paper is structured as follows : in Sec . 2 we give a brief review of some related works . In Sec . 3 we provide a detailed description of the message-passing equations and of the high level structure of the algorithms . In Sec . 4 we compare the performance of the message passing algorithms versus SGD based approaches in different learning settings . 2 RELATED WORKS . The literature on message passing algorithms is extensive , we refer to Mézard & Montanari ( 2009 ) and Zdeborová & Krzakala ( 2016 ) for a general overview . More related to our work , multilayer message-passing algorithms have been developed in inference contexts ( Manoel et al. , 2017 ; Fletcher et al. , 2018 ) , where they have been shown to produce exact marginals under certain statistical assumptions on ( unlearned ) weight matrices . The properties of message-passing for learning shallow neural networks have been extensively studied ( see Baldassi et al . ( 2020 ) and reference therein ) . Barbier et al . ( 2019 ) rigorously show that message passing algorithms in generalized linear models perform asymptotically exact inference under some statistical assumptions . Dictionary learning and matrix factorization are harder problems closely related to deep network learning problems , in particular to the modelling of a single intermediate layer . They have been approached using message passing in Kabashima et al . ( 2016 ) and Parker et al . ( 2014 ) , although the resulting predictions are found to be asymptotically inexact ( Maillard et al. , 2021 ) . The same problem is faced by the message passing algorithm recently proposed for a multi-layer matrix factorization scenario ( Zou et al. , 2021 ) . Unfortunately , our framework as well doesn ’ t yield asymptotic exact predictions . Nonetheless , it gives a message passing heuristic that for the first time is able to train deep neural networks on natural datasets , therefore sets a reference for the algorithmic applications of this research line . A few papers advocate the success of SGD to the geometrical structure ( smoothness and flatness ) of the loss landscape in neural networks ( Baldassi et al. , 2015 ; Chaudhari et al. , 2017 ; Garipov et al. , 2018 ; Li et al. , 2018 ; Pittorino et al. , 2021 ; Feng & Tu , 2021 ) . These considerations do not depend on the particular form of the SGD dynamics and should extend also to other types of algorithms , although SGD is by far the most popular choice among NNs practitioners due to its simplicity , flexibility , speed , and generalization performance . While our work focuses on message passing schemes , some of the ideas presented here , such as the PasP rule , can be combined with algorithms for Bayesian neural networks ’ training ( HernándezLobato & Adams , 2015 ; Wu et al. , 2018 ) . Recent work extends BP by combining it with graph neural networks ( Kuck et al. , 2020 ; Satorras & Welling , 2021 ) . Finally , some work in computational neuroscience shows similarities to our approach ( Rao , 2007 ) . 3 LEARNING BY MESSAGE PASSING . 3.1 POSTERIOR-AS-PRIOR UPDATES . We consider a multi-layer perceptron with L hidden neuron layers , having weight and bias parameters W = { W ` , b ` } L ` =0 . We allow for stochastic activations P ` ( x ` +1|z ` ) , where z ` is the neuron ’ s preactivation vector for layer ` , and P ` is assumed to be factorized over the neurons . If no stochasticity is present , P ` just encodes an element-wise activation function . The probability of output y given an input x is then given by P ( y |x , W ) = ∫ dx1 : L L∏ ` =0 P ` +1 ( x ` +1 | W ` x ` + b ` ) , ( 1 ) where for convenience we defined x0 = x and xL+1 = y . In a Bayesian framework , given a training set D = { ( xn , yn ) } n and a prior distribution over the weights qθ ( W ) in some parametric family , the posterior distribution is given by P ( W |D , θ ) ∝ ∏ n P ( yn |xn , W ) qθ ( W ) . ( 2 ) Here ∝ denotes equality up to a normalization factor . Using the posterior one can compute the Bayesian prediction for a new data-point x through the average P ( y |x , D , θ ) =∫ dW P ( y |x , W ) P ( W |D , θ ) . Unfortunately , the posterior is generically intractable due to the hard-to-compute normalization factor . On the other hand , we are mainly interested in training a distribution that covers wide minima of the loss landscape that generalize well ( Baldassi et al. , 2016a ) and in recovering pointwise estimators within these regions . The Bayesian modeling becomes an auxiliary tool to set the stage for the message passing algorithms seeking flat minima . We also need a formalism that allows for mini-batch training to speed-up the computation and deal with large datasets . Therefore , we devise an update scheme that we call Posterior-as-Prior ( PasP ) , where we evolve the parameters θt of a distribution qθt ( W ) computed as an approximate mini-batch posterior , in such a way that the outcome of the previous iteration becomes the prior in the following step . In the PasP scheme , θt retains the memory of past observations . We also add an exponential factor ρ , that we typically set close to 1 , tuning the forgetting rate and playing a role similar to the learning rate in SGD . Given a mini-batch ( Xt , yt ) sampled from the training set at time t and a scalar ρ > 0 , the PasP update reads qθt+1 ( W ) ≈ [ P ( W |yt , Xt , θt ) ] ρ , ( 3 ) where ≈ denotes approximate equality and we do not account for the normalization factor . A first approximation may be needed in the computation of the posterior , a second to project the approximate posterior onto the distribution manifold spanned by θ ( Minka , 2001 ) . In practice , we will consider factorized approximate posterior in an exponential family and priors qθ in the same family , although Eq . 3 generically allow for more refined approximations . Notice that setting ρ = 1 , the batch-size to 1 , and taking a single pass over the dataset , we recover the Assumed Density Filtering algorithm ( Minka , 2001 ) . For large enough ρ ( including ρ = 1 ) , the iterations of qθt will concentrate on a pointwise estimator . This mechanism mimics the reinforcement heuristic commonly used to turn Belief Propagation into a solver for constrained satisfaction problems ( Braunstein & Zecchina , 2006 ) and related to flat-minima discovery ( see focusing-BP in Baldassi et al . ( 2016a ) ) . A different prior updating mechanism which can be understood as empirical Bayes has been used in Baldassi et al . ( 2016b ) .

Summary: This manuscript provides an interesting try on alterative training algorithms for deep neural networks, based on (approximate) message-passing algorithms based on the well-known belief propagation (BP) algorithm.  In particular, the binary neural network is considered and four algorithms (BP, three variants of BP, i.e., BPI, MF, AMP) are proposed within a unified Posterior-As-Prior Update framework. Experiments are conducted on standard supervised classification tasks and continual learning settings, which shows comparable performances as standard SGD based methods.   ========================== After rebuttal: =========================  I have read the authors' feedback (many thanks for the detailed point-to-point feedback) and other reviewers' comments and modified the score accordingly. Overall, the proposed scheme is interesting, though strictly speaking the results are not very advantageous (at least from its current results) compared to traditional ones, and some of the comparisons seem not very reasonable/fair.