Mean-field inference methods for neural networks

Inference in generative models. Generative models used for unsupervised learning are statistical models defining high-dimensional distributions with complex dependencies. As we have seen in section 2.2, the most common training objective in unsupervised learning is the maximization of the log-likelihood, i.e. the log of the probability assigned by the generative model to the training set ${\left\{{\underline{x}}_{\left(k\right)}\right\}}_{k=1}^{P}$. Computing the probability of observing a given sample ${\underline{x}}_{\left(k\right)}$ is an inference question. It requires to marginalize over all the hidden representations of the problem. For instance in the RBM (19),

$$\begin{equation}{p}_{\theta }\left({\underline{x}}_{\left(k\right)}\right)=\frac{1}{\mathcal{Z}}\sum _{\underline{t}\in {\left\{0,1\right\}}^{M}}{\mathrm{e}}^{{\underline{a}}^{\top }{\underline{x}}_{\left(k\right)}+{\underline{b}}^{\top }\underline{t}+{\underline{x}}_{\left(k\right)}^{\top }\underline{\underline{W}}\underline{t}}.\end{equation} \tag{ 20 }$$

While the numerator will be easy to evaluate, the partition function has no analytical expression and its exact evaluation requires to sum over all possible states of the network.

Learning as statistical inference: Bayesian inference and the teacher–student scenario. The practical problem of training neural networks from data as introduced in section 2 is not in general interpreted as inference. To do so, one needs to treat the learnable parameters as random variables, which is the case in Bayesian learning. For instance in supervised learning, an underlying prior distribution p_θ(θ) for the weights and biases of a neural network (5) and (6) is assumed, so that Bayes rule defines a posterior distribution given the training data $\mathcal{D}$,

$$\begin{equation}p\left(\theta \vert \mathcal{D}\right)=\frac{p\left(\mathcal{D}\vert \theta \right){p}_{\theta }\left(\theta \right)}{p\left(\mathcal{D}\right)}.\end{equation} \tag{ 21 }$$

Compared to the single output of risk minimization, we obtain an entire distribution for the learned parameters θ, which takes into account not only the training data but also some knowledge on the structure of the parameters (e.g. sparsity) through the prior. In practice, Bayesian learning and traditional empirical risk minimization may not be so different. On the one hand, the Bayesian posterior distribution is often summarized by a point estimate such as its maximum. On the other hand risk minimization is often biased toward desired properties of the weights through regularization techniques (e.g. promoting small norm) recalling the role of the Bayesian prior.

However, from a theoretical point of view, Bayesian learning is of particular interest in the teacher–student scenario. The idea here is to consider a toy model of the learning problem where parameters are effectively drawn from a prior distribution. Let us use as an illustration the case of the supervised learning of the perceptron model (3). We draw a weight vector ${\underline{w}}_{0}$, from a prior distribution p_w(⋅), along with a set of P inputs ${\left\{{\underline{x}}_{\left(k\right)}\right\}}_{k=1}^{P}$ i.i.d. from a data distribution p_x(⋅). Using this teacher perceptron model we also draw a set of possibly noisy corresponding outputs y_(k) from a teacher conditional probability $p\left(.\vert {\underline{w}}_{0}^{\top }{\underline{x}}_{\left(k\right)}\right)$. From the training set of the P pairs $\mathcal{D}=\left\{{\underline{x}}_{\left(k\right)},{y}_{\left(k\right)}\right\}$, one can attempt to rediscover the teacher rule by training a student perceptron model. The problem can equivalently be phrased as a reconstruction inference question: can we recover the value of ${\underline{w}}_{0}$ from the observations in $\mathcal{D}$? The Bayesian framework yields a posterior distribution of solutions

$$\begin{equation}p\left(\underline{w}\vert \mathcal{D}\right)=\prod _{k=1}^{P}p\left({y}_{\left(k\right)}\vert {\underline{w}}^{\top }{\underline{x}}_{\left(k\right)}\right){p}_{w}\left(\underline{w}\right)\;/\;\prod _{k=1}^{P}p\left({y}_{\left(k\right)}\vert {\underline{x}}_{\left(k\right)}\right).\end{equation} \tag{ 22 }$$

Note that the terminology of teacher–student applies for a generic inference problem of reconstruction: the statistical model used to generate the data along with the realization of the unknown signal is called the teacher; the statistical model assumed to perform the reconstruction of the signal is called the student. When the two models are identical or matched, the inference is Bayes optimal. When the teacher model is not perfectly known, the statistical models can also be different (from slightly differing prior distributions to entirely different models), in which case they are said to be mismatched, and the reconstruction is suboptimal.

Of course in practical machine learning applications of neural networks, one has only access to an empirical distribution of the data and it is unclear whether there should exist a formal rule underlying the input–output mapping. Yet the teacher–student setting is a modeling strategy of learning which offers interesting possibilities of analysis and we shall refer to numerous works resorting to the setup in section 6.

Many inference questions in the line of the ones mentioned in the previous section have no tractable exact solution. When there exists no analytical closed-form, computations of averages and marginals require summing over configurations. Their number typically scales exponentially with the size of the system, then becoming astronomically large for high-dimensional models. Hence it is necessary to design approximate inference strategies. They may require an algorithmic implementation but must run in finite (typically polynomial) time. An important cross-fertilization between statistical physics and information sciences have taken place over the past decades around the questions of inference. Two major classes of such algorithms are Monte Carlo Markov chains (MCMC), and mean-field methods. The former is nicely reviewed in the context of statistical physics in [Kra06]. The latter will be the focus of this short review, in the context of deep learning.

Note that representations of joint probability distributions through probabilistic graphical models and factor graphs are crucial tools to design efficient inference strategies. In appendix B, we quickly introduce for the unfamiliar reader these two formalisms that enable to encode and exploit independencies between random variables. As examples, figure 2 presents graphical representations of the RBM measure (19) and the posterior distribution in the Bayesian learning of the perceptron as discussed in the previous section.

Zoom In Zoom Out Reset image size

The naive mean-field method consists in approximating the joint probability distribution of interest by a fully factorized distribution. Therefore, it ignores correlations between random variables. Among multiple methods of derivation, we present here the variational method: it is the best known method across fields and it readily shows that, for any joint probability distribution interpreted as a Boltzmann distribution, the rather crude naive mean-field approximation yields an upper bound on the free energy. For the purpose of demonstration we consider a Boltzmann machine without hidden units (Ising model) with variables (spins) $\underline{x}=\left({x}_{1},\dots ,\;{x}_{N}\right)\in \mathcal{X}={\left\{0,1\right\}}^{N}$, and energy function

$$\begin{equation}E\left(\underline{x}\right)=-\sum _{i=1}^{N}{b}_{i}{x}_{i}-\sum _{\left(ij\right)}{W}_{ij}{x}_{i}{x}_{j}=-{\underline{b}}^{\top }\underline{x}-\frac{1}{2}{\underline{x}}^{\top }\underline{\underline{W}}\underline{x},\quad \underline{b}\in {\mathbb{R}}^{N},\;\underline{\underline{W}}\in {\mathbb{R}}^{N{\times}N},\end{equation} \tag{ 30 }$$

where the notation (ij) stands for pairs of connected spin-variables, and the weight matrix $\underline{\underline{W}}$ is symmetric. The choices for {0, 1} rather than {−1, +1} for the variable values, the notations $\underline{\underline{W}}$ for weights (instead of couplings), $\underline{b}$ for biases (instead of local fields), as well as the vector notation, are leaning toward the machine learning conventions. We denote by ${q}_{\underline{m}}$ a fully factorized distribution on {0,1}^N, which is a multivariate Bernoulli distribution parametrized by the mean values $\underline{m}=\left({m}_{1},\dots ,\;{m}_{N}\right)\in {\left[0,1\right]}^{N}$ of the marginals (denoted by ${q}_{{m}_{i}}$):

$$\begin{equation}{q}_{\underline{m}}\left(\underline{x}\right)=\prod _{i=1}^{N}{q}_{{m}_{i}}\left({x}_{i}\right)=\prod _{i=1}^{N}{m}_{i}\delta \left({x}_{i}-1\right)+\left(1-{m}_{i}\right)\delta \left({x}_{i}\right).\end{equation} \tag{ 31 }$$

We look for the optimal ${q}_{\underline{m}}$ distribution to approximate the Boltzmann distribution $p\left(\underline{x}\right)={\mathrm{e}}^{-\beta E\left(\underline{x}\right)}/\mathcal{Z}$ by minimizing the KL-divergence

$$\begin{equation}\underset{\underline{m}}{\mathrm{min}}\;\mathrm{K}\mathrm{L}\left({q}_{\underline{m}}{\Vert}p\right)=\underset{\underline{m}}{\mathrm{min}}\;\sum _{\underline{x}\in \mathcal{X}}{q}_{\underline{m}}\left(\underline{x}\right)\mathrm{log}\;\frac{{q}_{\underline{m}}\left(\underline{x}\right)}{p\left(\underline{x}\right)}\end{equation} \tag{ 32 }$$

$$\begin{equation}=\underset{\underline{m}}{\mathrm{min}}\;\sum _{\underline{x}\in \mathcal{X}}{q}_{\underline{m}}\left(\underline{x}\right)\mathrm{log}\;{q}_{\underline{m}}\left(\underline{x}\right)+\beta \sum _{\underline{x}\in \mathcal{X}}{q}_{\underline{m}}\left(\underline{x}\right)E\left(\underline{x}\right)+\mathrm{log}\;\mathcal{Z}\end{equation} \tag{ 33 }$$

$$\begin{equation}=\underset{\underline{m}}{\mathrm{min}}\;\;\beta G\left({q}_{\underline{m}}\right)-\beta F{\geqslant}0,\end{equation} \tag{ 34 }$$

where the last inequality comes from the positivity of the KL-divergence. For a generic distribution q, G(q) is the Gibbs free energy for the energy $E\left(\underline{x}\right)$,

$$\begin{equation}G\left(q\right)=\sum _{\underline{x}\in \mathcal{X}}q\left(\underline{x}\right)E\left(\underline{x}\right)+\frac{1}{\beta }\sum _{\underline{x}\in \mathcal{X}}q\left(\underline{x}\right)\mathrm{log}\;q\left(\underline{x}\right)=U\left(q\right)-H\left(q\right)/\beta {\geqslant}F,\end{equation} \tag{ 35 }$$

involving the average energy U(q) and the entropy H(q). It is greater than the true free energy F except when q = p, in which case they are equal. Note that this fact also means that the Boltzmann distribution minimizes the Gibbs free energy. Restricting to factorized ${q}_{\underline{m}}$ distributions, we obtain the naive mean-field approximations for the mean value of the variables (or magnetizations) and the free energy:

$$\begin{equation}{\underline{m}}^{{\ast}}=\underset{\underline{m}}{\text{arg\,min}}\;G\left({q}_{\underline{m}}\right)={\langle \underline{x}\rangle }_{{q}_{{\underline{m}}^{{\ast}}}},\end{equation} \tag{ 36 }$$

$$\begin{equation}{F}_{\mathrm{N}\mathrm{M}\mathrm{F}}=G\left({q}_{{\underline{m}}^{{\ast}}}\right){\geqslant}F.\end{equation} \tag{ 37 }$$

The choice of a very simple family of distributions ${q}_{\underline{m}}$ limits the quality of the approximation but allows for tractable computations of observables, for instance the two-spins correlations ${\langle {x}_{i}{x}_{j}\rangle }_{{q}^{{\ast}}}={m}_{i}^{{\ast}}{m}_{j}^{{\ast}}$ or variance of one spin ${\langle {x}_{i}^{2}\rangle }_{{q}^{{\ast}}}-{\langle {x}_{i}\rangle }_{{q}^{{\ast}}}^{2}={m}_{i}^{{\ast}}-{{m}_{i}^{{\ast}}}^{2}$.

In our example of the Boltzmann machine, it is easy to compute the Gibbs free energy for the factorized ansatz, we define functions of the magnetization vector:

$$\begin{equation}{U}_{\text{NMF}}\left(\underline{m}\right)={\langle E\left(\underline{x}\right)\rangle }_{{q}_{\underline{m}}}=-{\underline{b}}^{\top }\underline{m}-\frac{1}{2}{\underline{m}}^{\top }\underline{\underline{W}}\underline{m},\end{equation} \tag{ 38 }$$

$$\begin{equation}{H}_{\text{NMF}}\left(\underline{m}\right)=-{\langle \mathrm{log}\;{q}_{\underline{m}}\left(\underline{x}\right)\rangle }_{{q}_{\underline{m}}}=-\sum _{i=1}^{N}{m}_{i}\;\mathrm{log}\;{m}_{i}+\left(1-{m}_{i}\right)\mathrm{log}\left(1-{m}_{i}\right),\end{equation} \tag{ 39 }$$

$$\begin{equation}{G}_{\text{NMF}}\left(\underline{m}\right)=G\left({q}_{\underline{m}}\right)={U}_{\text{NMF}}\left(\underline{m}\right)-{H}_{\text{NMF}}\left(\underline{m}\right)/\beta .\end{equation} \tag{ 40 }$$

Looking for stationary points we find a closed set of non linear equations for the ${m}_{i}^{{\ast}}$,

$$\begin{equation}{\left.\frac{\partial {G}_{\mathrm{N}\mathrm{M}\mathrm{F}}}{\partial {m}_{i}}\right\vert }_{{\underline{m}}^{{\ast}}}=0\quad {\Rightarrow}\quad {m}_{i}^{{\ast}}=\sigma \left(\beta {b}_{i}+\sum _{j\in \partial i}\beta {W}_{ij}{m}_{j}^{{\ast}}\right)\quad \forall i=1,\dots ,\;N,\end{equation} \tag{ 41 }$$

where $\sigma \left(x\right)={\left(1+{\mathrm{e}}^{-x}\right)}^{-1}$. The solutions can be computed by iterating these relations from a random initialization until a fixed point is reached.

To understand the implication of the restriction to factorized distributions, it is instructive to compare this naive mean-field equation with the exact identity

$$\begin{equation}{\langle {x}_{i}\rangle }_{p}={\left\langle \sigma \left(\beta {b}_{i}+\sum _{j\in \partial i}\beta {W}_{ij}{x}_{j}\right)\right\rangle }_{p},\end{equation} \tag{ 42 }$$

derived in a few lines in appendix C. Under the Boltzmann distribution $p\left(\underline{x}\right)={\mathrm{e}}^{-\beta E\left(\underline{x}\right)}/\mathcal{Z}$, these averages are difficult to compute. The naive mean-field method is neglecting the fluctuations of the effective field felt by the variable x_i: ∑_j∈∂iW_ijx_j, keeping only its mean ∑_j∈∂iW_ijm_j. This incidentally justifies the name of mean-field methods.

The previous derivation shows that the naive mean-field approximation allows to bound the free energy. While this bound is expected to be rough in general, the approximation is reliable when the fluctuations of the local effective fields ∑_j∈∂iW_ijx_j are small. This may happen in particular in the thermodynamic limit N → ∞ in infinite range models, that is when weights or couplings are not only local but distributed in the entire system, or if each variable interacts directly with a non-vanishing fraction of the whole set of variables (e.g. [OS01] section 2). The influence on one given variable of the rest of the system can then be treated as an average background. Provided the couplings are weak enough, the naive mean-field method may even become asymptotically exact. This is the case of the Curie–Weiss model, which is the fully connected version of the model (30) with all W_ij = 1/N (see e.g. section 2.5.2 of [MM09]). The sum of weakly dependent variables then concentrates on its mean by the central limit theorem. We stress that it means that for finite dimensional models (more representative of a physical system, where for instance variables are assumed to be attached to the vertices of a lattice with nearest neighbors interactions), mean-field methods are expected to be quite poor. By contrast, infinite range models (interpreted as infinite-dimensional models by physicists) are thus traditionally called mean-field models.

In the next section we will recover the naive mean-field equations through a different method. The following derivation will also allow to compute corrections to the rather crude approximation we just discussed by taking into account some of the correlations it neglects.

Going back to the variational formulation (34), we shall now perform a minimization in two steps. Consider first the family of distributions ${q}_{\underline{m}}$ enforcing ${\langle \underline{x}\rangle }_{{q}_{\underline{m}}}=\underline{m}$ for a fixed vector of magnetizations $\underline{m}$, but without any factorization constraint. The corresponding Gibbs free energy is

$$\begin{equation}G\left({q}_{\underline{m}}\right)=U\left({q}_{\underline{m}}\right)-H\left({q}_{\underline{m}}\right)/\beta .\end{equation} \tag{ 43 }$$

A first minimization at fixed $\underline{m}$ over the ${q}_{\underline{m}}$ defines another auxiliary free energy

$$\begin{equation}{G}_{\text{TAP}}\left(\underline{m}\right)=\underset{{q}_{\underline{m}}}{\mathrm{min}}\;G\left({q}_{\underline{m}}\right).\end{equation} \tag{ 44 }$$

A second minimization over $\underline{m}$ would recover the overall unconstrained minimum of the variational problem (34) which is the exact free energy

$$\begin{equation}F=-\mathrm{log}\;\mathcal{Z}/\beta =\underset{\underline{m}}{\mathrm{min}}\;{G}_{\text{TAP}}\left(\underline{m}\right).\end{equation} \tag{ 45 }$$

Yet the actual value of ${G}_{\text{TAP}}\left(\underline{m}\right)$ turns out as complicated to compute as F itself. Fortunately, $\beta {G}_{\text{TAP}}\left(\underline{m}\right)$ can be easily approximated by a Taylor expansion around β = 0 due to interactions vanishing at high temperature, as noticed by Plefka, Georges and Yedidia [GY99, Ple82]. After expanding, the minimization over ${G}_{\text{TAP}}\left(\underline{m}\right)$ yields a set of self consistent equations on the magnetizations $\underline{m}$, called the TAP equations, reminiscent of the naive mean-field equations (41). Here again, the consistency equations are typically solved by iterations. Plugging the solutions ${\underline{m}}^{{\ast}}$ back into the expanded expression yields the TAP free energy ${F}_{\text{TAP}}={G}_{\text{TAP}}\left({\underline{m}}^{{\ast}}\right)$. Note that ultimately the approximation lies in the truncation of the expansion. At first order the naive mean-field approximation is recovered. Historically, the expansion was first stopped at the second order. This choice was model dependent, it results from the fact that the mean-field theory is already exact at the second order for the SK model [MH76, Ple82, TAP77].

For the Boltzmann machine (30), the TAP equations and TAP free energy (truncated at second order) are [TAP77],

$$\begin{equation}{m}_{i}^{{\ast}}=\sigma \left(\beta {b}_{i}+\sum _{j\in \partial i}\beta {W}_{ij}{m}_{j}^{{\ast}}-{\beta }^{2}{W}_{ij}^{2}\left({m}_{j}^{{\ast}}-\frac{1}{2}\right)\left({m}_{i}^{{\ast}}-{{m}_{i}^{{\ast}}}^{2}\right)\right)\forall i\end{equation} \tag{ 46 }$$

$$\begin{align}\hfill \beta {G}_{\mathrm{T}\mathrm{A}\mathrm{P}}\left({\underline{m}}^{{\ast}}\right)& =-{H}_{\mathrm{N}\mathrm{M}\mathrm{F}}\left({\underline{m}}^{{\ast}}\right)-\beta \sum _{i=1}^{N}{b}_{i}{m}_{i}^{{\ast}}-\beta \sum _{\left(ij\right)}{m}_{i}^{{\ast}}{W}_{ij}{m}_{j}^{{\ast}}\hfill \\ \hfill & \hfill \quad -\frac{{\beta }^{2}}{2}\sum _{\left(ij\right)}{W}_{ij}^{2}\left({m}_{i}^{{\ast}}-{{m}_{i}^{{\ast}}}^{2}\right)\left({m}_{j}^{{\ast}}-{{m}_{j}^{{\ast}}}^{2}\right),\end{align} \tag{ 47 }$$

where the naive mean-field entropy H_NMF was defined in (39). For this model, albeit with {+1, −1} variables instead of {0, 1}, several references pedagogically present the details of the derivation sketched in the previous paragraph. The interested reader should check in particular [OS01, Zam10]. We also present a more general derivation in appendix D, see section 4.2.3.

Onsager reaction term. Compared to the naive mean-field approximation the TAP equations include a correction to the effective field called the Onsager reaction term. The idea is that, in the effective field at variable i, we should consider corrected magnetizations of neighboring spins j ∈ ∂i, that would correspond to the absence of variable i. This intuition echoes at two other derivations of the TAP approximation: the cavity method [MPV86] that will not be covered here and the message passing which will be discussed in the next section.

As far as the SK model is concerned, this second order correction is enough in the thermodynamic limit as the statistics of the weights imply that higher orders will typically be subleading. Yet in general, the correct TAP equations for a given model will depend on the statistics of interactions and there is no guarantee that there exists a finite order of truncation leading to an exact mean-field theory. In section 5.2 we will discuss models beyond SK where a conjectured exact TAP approximation can be derived.

Single instance. Although the selection of the correct TAP approximation relies on the statistics of the weights, the derivation of the expansion outlined above does not require to average over them, i.e. it does not require an average over the disorder. Consequently, the approximation method is well defined for a single instance of the random disordered model and the TAP free energy and magnetizations can be computed for a given (realization of the) set of weights ${\left\{{W}_{ij}\right\}}_{\left(ij\right)}$ as explained in the following paragraph. In other words, it means that the approximation can be used to design practical inference algorithms in finite-sized problems and not only for theoretical predictions on average over the disordered class of models. Crucially, these algorithms may provide approximations of disorder-dependent observables, such as correlations, and not only of self averaging quantities.

Finding solutions. The self-consistent equations on the magnetizations (46) are usually solved by turning them into an iteration scheme and looking for fixed points. This generic recipe leaves nonetheless room for interpretation: which exact form should be iterated? How should the updates for the different equations be scheduled? Which time indexing should be used? While the following scheme may seem natural

$$\begin{equation}{{m}_{i}}^{\left(t+1\right)}{\leftarrow}\sigma \left(\beta {b}_{i}+\sum _{j\in \partial i}\beta {W}_{ij}{{m}_{j}}^{\left(t\right)}-{W}_{ij}^{2}\left({{m}_{j}}^{\left(t\right)}-\frac{1}{2}\right)\left({{m}_{i}}^{\left(\mathbf{t}\right)}-{{{m}_{i}}^{\left(\mathbf{t}\right)}}^{2}\right)\right),\end{equation} \tag{ 48 }$$

it typically has more convergence issues than the following alternative scheme including the time index t − 1

$$\begin{equation}{{m}_{i}}^{\left(t+1\right)}{\leftarrow}\sigma \left(\beta {b}_{i}+\sum _{j\in \partial i}\beta {W}_{ij}{{m}_{j}}^{\left(t\right)}-{W}_{ij}^{2}\left({{m}_{j}}^{\left(t\right)}-\frac{1}{2}\right)\left({{m}_{i}}^{\left(\mathbf{t}-\mathbf{1}\right)}-{{{m}_{i}}^{\left(\mathbf{t}-\mathbf{1}\right)}}^{2}\right)\right).\end{equation} \tag{ 49 }$$

This issue was discussed in particular in [Bol14, Kab03]. Remarkably, this last scheme, or algorithm, is actually the one obtained by the approximate message passing derivation that will be discussed in the upcoming section 4.3.

Solutions of the TAP equations. The TAP equations can admit multiple solutions with either equal or different TAP free energy. While the true free energy F corresponds to the minimum of the Gibbs free energy, reached for the Boltzmann distribution, the TAP derivation consists in performing an effectively unconstrained minimization in two steps, but with an approximation through a Taylor expansion in between. The truncation of the expansion therefore breaks the correspondence between the discovered minimizer and the unique Boltzmann distribution, hence the possible multiplicity of solutions. For the SK model for instance, the number of solutions of the TAP equations increases rapidly as β grows [MPV86]. While the different solutions can be accessed using different initializations of the iterative scheme, it is notably hard in phases where they are numerous to find exhaustively all the TAP solutions. In theory, they should be weighted according to their free energy density and averaged to recover the thermodynamics predicted by the replica computation [DY83], another mean-field approximation discussed in section 4.4.

In the derivation outlined above for binary variables, x_i = 0 or 1, the mean of each variable m_i was fixed. This is enough to parametrize the corresponding marginal distribution ${q}_{{m}_{i}}\left({x}_{i}\right)$. Yet the expansion can actually be generalized to Potts variables (taking multiple discrete values) or even real valued variables by introducing appropriate parameters for the marginals. A general derivation fixing arbitrary real valued marginal distribution was proposed in appendix B of [LKZ17] for the problem of low rank matrix factorization. Alternatively, another level of approximation can be introduced for real valued variables by restricting the set of marginal distributions tested to a parametrized family of distributions. By choosing a Gaussian parametrization, one recovers TAP equations equivalent to the approximate message passing algorithm that will be discussed in the next section. In appendix D, we present a derivation for real-valued Boltzmann machines with a Gaussian parametrization as proposed in [TGM⁺18].

Definition. We introduce the generalized linear model (GLM) which is a fairly simple model to illustrate message passing algorithms and which is also an elementary brick for a large range of interesting inference questions on neural networks. It falls under the teacher–student set up: a student model is used to reconstruct a signal from a teacher model producing indirect observations. In the GLM, the product of an unknown signal ${\underline{x}}_{0}\in {\mathbb{R}}^{N}$ and a known weight matrix $\underline{\underline{W}}\in {\mathbb{R}}^{N{\times}M}$ is observed as $\underline{y}$ through a noisy channel p_out,0,

$$\begin{equation}\begin{cases}\underline{\underline{W}}\sim {p}_{W}\left(\underline{\underline{W}}\right)\quad \hfill \\ {\underline{x}}_{0}\sim {p}_{{x}_{0}}\left({\underline{x}}_{0}\right)=\prod _{i=1}^{N}{p}_{{x}_{0}}\left({x}_{0,i}\right)\quad \hfill \end{cases}\quad {\Rightarrow}\underline{y}\sim {p}_{\mathrm{o}\mathrm{u}\mathrm{t},0}\left(\underline{y}\vert \underline{\underline{W}}{\underline{x}}_{0}\right)=\prod _{\mu =1}^{M}{p}_{\mathrm{o}\mathrm{u}\mathrm{t},0}\left({y}_{\mu }\vert {\underline{w}}_{\mu }^{\top }{\underline{x}}_{0}\right).\end{equation} \tag{ 50 }$$

The probabilistic graphical model corresponding to this teacher is represented in figure 3. The prior over the signal ${p}_{{x}_{0}}$ is supposed to be factorized, and the channel p_out,0 likewise. The inference problem is to produce an estimator $\hat{\underline{x}}$ for the unknown signal ${\underline{x}}_{0}$ from the observations $\underline{y}$. Given the prior p_x and the channel p_out of the student, not necessarily matching the teacher, the posterior distribution is

$$\begin{equation}p\left(\underline{x}\vert \underline{y},\underline{\underline{W}}\right)=\frac{1}{\mathcal{Z}\left(\underline{y},\underline{\underline{W}}\right)}\;\prod _{\mu =1}^{M}{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu }\vert \sum _{i=1}^{N}{W}_{\mu i}{x}_{i}\right)\;\prod _{i=1}^{N}{p}_{x}\left({x}_{i}\right),\end{equation} \tag{ 51 }$$

$$\begin{equation}\mathcal{Z}\left(\underline{y},\underline{\underline{W}}\right)=\int \mathrm{d}\underline{x}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(\underline{y}\vert \underline{x},\underline{\underline{W}}\right){p}_{x}\left(\underline{x}\right),\end{equation} \tag{ 52 }$$

represented as a factor graph also in figure 3. The difficulty of the reconstruction task of ${\underline{x}}_{0}$ from $\underline{y}$ is controlled by the measurement ratio α = M/N and the amplitude of the noise possibly present in the channel.

Zoom In Zoom Out Reset image size

Applications. The generic GLM underlies a number of applications. In the context of neural networks of particular interest in this technical review, the channel p_out generating observations $\underline{y}\in {\mathbb{R}}^{M}$ can equivalently be seen as a stochastic activation function $g\left(\cdot ;\underline{{\epsilon}}\right)$ incorporating a noise $\underline{{\epsilon}}\in {\mathbb{R}}^{M}$ component-wise to the output,

$$\begin{equation}{y}_{\mu }=g\left({\underline{w}}_{\mu }^{\top }\underline{x};\;{\underline{{\epsilon}}}_{\mu }\right).\end{equation} \tag{ 53 }$$

The inference of the teacher signal in a GLM has then two possible interpretations. On the one hand, it can be interpreted as the reconstruction of the input $\underline{x}$ of a stochastic single-layer neural network from its output $\underline{y}$. For example, this inference problem can arise in the maximum likelihood training of a one-layer VAE (see corresponding paragraph in section 2.2). On the other hand, the same question can also correspond to the Bayesian learning of a single-layer neural network with a single output—the perceptron—where this time $\left\{\underline{\underline{W}},\underline{y}\right\}$ are interpreted as the collection of training input–output pairs and ${\underline{x}}_{0}$ plays the role of the unknown weight vector of the teacher (as cited as an example in section 3.1.1). However, note that one of the most important applications of the GLM, compressed sensing (CS) [Don06], does not involve neural networks.

Statistical physics treatment, random weights and scaling. From the statistical physics perspective, the effective energy functional is read from the posterior (52) seen as a Boltzmann distribution with energy

$$\begin{equation}E\left(\underline{x}\right)=-\mathrm{log}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(\underline{y}\vert \underline{x},\underline{\underline{W}}\right){p}_{x}\left(\underline{x}\right)=-\sum _{\mu =1}^{M}\mathrm{log}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu }\vert \sum _{i=1}^{N}{W}_{\mu i}{x}_{i}\right)-\sum _{i=1}^{N}\mathrm{log}\;{p}_{x}\left({x}_{i}\right).\end{equation} \tag{ 54 }$$

The inverse temperature β has here no formal equivalent and can be thought as being equal to 1. The energy is a function of the random realizations of $\underline{\underline{W}}$ and $\underline{y}$, playing the role of the disorder. Furthermore, the validity of the approximation presented below require additional assumptions. Crucially, the weight matrix is assumed to have i.i.d. Gaussian entries with zero mean and variance 1/N, much like in the SK model. The prior of the signal is chosen so as to ensure that the x_i–s (and consequently the y_μ–s) remain of order 1. Finally, the thermodynamic limit N → ∞ is taken for a fixed measurement ratio α = M/N.

Recall that inference in high-dimensional problems consists in marginalizations over complex joint distributions, typically in the view of computing partition functions, averages or marginal probabilities for sampling. Belief propagation (BP) is an inference algorithm, sometimes exact and sometimes approximate as we will see, leveraging the known factorization of a distribution, which encodes the precious information of (in)depencies between the random variables in the joint distribution. For a generic joint probability distribution p over $\underline{x}\in {\mathbb{R}}^{N}$ factorized as

$$\begin{equation}p\left(\underline{x}\right)=\frac{1}{\mathcal{Z}}\prod _{\mu =1}^{M}{\psi }_{\mu }\left({\underline{x}}_{\partial \mu }\right),\end{equation} \tag{ 55 }$$

ψ_μ are called potential functions taking as arguments the variables x_i–s involved in the factor μ shortened as ${\underline{x}}_{\partial \mu }$.

Definition of messages. Let us first write the BP equations and then explain the origin of these definitions. The underlying factor graph of (55) has N nodes carrying variables x_i–s and M factors associated with the potential functions ψ_μ–s (see appendix B for a quick reminder). BP acts on messages variables which are tied to the edges of the factor graph. Specifically, the sum-product version of the algorithm (as opposed to the max-sum, see e.g. [MM09]) consists in the update equations

$$\begin{equation}{\tilde {m}}_{\mu \to i}^{\left(t\right)}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{\mu \to i}}\int \prod _{{i}^{\prime }\in \partial \mu {\backslash}i}\mathrm{d}{x}_{{i}^{\prime }}\;{\psi }_{\mu }\left({\underline{x}}_{\partial \mu }\right)\prod _{{i}^{\prime }\in \partial \mu {\backslash}i}{m}_{{i}^{\prime }\to \mu }^{\left(t\right)}\left({x}_{{i}^{\prime }}\right),\end{equation} \tag{ 56 }$$

$$\begin{equation}{m}_{i\to \mu }^{\left(t+1\right)}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{i\to \mu }}{p}_{x}\left({x}_{i}\right)\prod _{{\mu }^{\prime }\in \partial i{\backslash}\mu }{\tilde {m}}_{{\mu }^{\prime }\to i}^{\left(t\right)}\left({x}_{i}\right)\end{equation} \tag{ 57 }$$

where again the i–s index the variable nodes and the μ–s index the factor nodes. The notation ∂μ\i designate the set of neighbor variables of the factor μ except the variable i (and reciprocally for ∂i\μ). The partition functions ${\mathcal{Z}}_{i\to \mu }$ and ${\mathcal{Z}}_{\mu \to i}$ are normalization factors ensuring that the messages can be interpreted as probabilities.

For acyclic (or tree-like) factor graphs, the BP updates are guaranteed to converge to a fixed point, that is a set of time independent messages $\left\{{m}_{i\to \mu },{\tilde {m}}_{\mu \to i}\right\}$ solution of the system of equations (56) and (57). Starting at a leaf of the tree, these messages communicate beliefs of a given node variable taking a given value based on the nodes and factors already visited along the tree. More precisely, ${\tilde {m}}_{\mu \to i}\left({x}_{i}\right)$ is the marginal probability of x_i in the factor graph before visiting the factors in ∂i except for μ, and m_i→μ(x_i) is equal the marginal probability of x_i in the factor graph before visiting the factor μ, see figure 4.

Zoom In Zoom Out Reset image size

Thus, at convergence of the iterations, the marginals can be computed as

$$\begin{equation}{m}_{i}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{i}}{p}_{x}\left({x}_{i}\right)\prod _{\mu \in \partial i}{\tilde {m}}_{\mu \to i}\left({x}_{i}\right),\end{equation} \tag{ 58 }$$

which can be seen as the main output of the BP algorithm. These marginals will only be exact on trees where incoming messages, computed from different part of the graph, are independent. Nonetheless, the algorithm (56) and (57), occasionally then called loopy-BP, can sometimes be converged on graphs with cycles and in some cases will still provide high quality approximations. For instance, graphs with no short loops are locally tree like and BP is an efficient method of approximate inference, provided correlations decay with distance (i.e. incoming messages at each node are still effectively independent). BP will also appear principled for some infinite range mean-field models previously discussed; an example of which being our case-study the GLM discussed below. While this is the only example that will be discussed here in the interest of conciseness, getting fluent in BP generally requires more than one application. The interested reader could also consult [YFW02] and [MM09] section 14.1. for simple concrete examples.

The Bethe free energy. The BP algorithm can also be recovered from a variational argument. Let us consider both the single variable marginals m_i(x_i) and the marginals of the neighborhood of each factor ${\tilde {m}}_{\mu }\left({\underline{x}}_{\partial \mu }\right)$. On tree graphs, the joint distribution (55) can be re-expressed as

$$\begin{equation}p\left(\underline{x}\right)=\frac{\prod _{\mu =1}^{M}{\tilde {m}}_{\mu }\left({\underline{x}}_{\partial \mu }\right)}{\prod _{i=1}^{N}{m}_{i}{\left({x}_{i}\right)}^{{n}_{i}-1}},\end{equation} \tag{ 59 }$$

where n_i is the number of neighbor factors of the ith variable. Abusively, we can use this form as an ansatz for loopy graph and plug it in the Gibbs free energy to derive an approximation of the free energy, similarly to the naive mean-field derivation of section 4.1. This time the variational parameters will be the distributions m_i and ${\tilde {m}}_{\mu }$ (see e.g. [MM09, YFW02] for additional details). The corresponding functional form of the Gibbs free energy is called the Bethe free energy:

$$\begin{equation}{F}_{\mathrm{B}\mathrm{e}\mathrm{t}\mathrm{h}\mathrm{e}}\left({m}_{i},{\tilde {m}}_{\mu }\right)=-\int \mathrm{d}{\underline{x}}_{\partial \mu }\;{\tilde {m}}_{\mu }\left({\underline{x}}_{\partial \mu }\right)\mathrm{ln}\;{\psi }_{\mu }\left({\underline{x}}_{\partial \mu }\right)+\left({n}_{i}-1\right)H\left({m}_{i}\right)-H\left({\tilde {m}}_{\mu }\right),\end{equation} \tag{ 60 }$$

where H(q) is the entropy of the distribution q. Optimization of the Bethe free energy with respect to its arguments under the constraint of consistency

$$\begin{equation}\int \mathrm{d}{\underline{x}}_{\partial \mu {\backslash}i}\;{\tilde {m}}_{\mu }\left({\underline{x}}_{\partial \mu }\right)={m}_{i}\left({x}_{i}\right)\end{equation} \tag{ 61 }$$

involves Lagrange multipliers which can be shown to be related to the messages defined in (56) and (57). Eventually, one can verify that marginals defined as (58) and

$$\begin{equation}{\tilde {m}}_{\mu }\left({\underline{x}}_{{\partial }_{\mu }}\right)=\frac{1}{{\mathcal{Z}}_{\mu }}\psi \left({\underline{x}}_{\partial \mu }\right)\prod _{i\in \partial \mu }{m}_{i\to \mu }\left({x}_{i}\right),\end{equation} \tag{ 62 }$$

are stationary point of the Bethe free energy for messages that are BP solutions. In other words, the BP fixed points are consistent with the stationary point of the Bethe free energy. Using the normalizing constants of the messages, the Bethe free energy can also be re-written as

$$\begin{equation}{F}_{\mathrm{B}\mathrm{e}\mathrm{t}\mathrm{h}\mathrm{e}}=-\sum _{i\in V}\mathrm{log}\;{\mathcal{Z}}_{i}-\sum _{\mu \in F}\mathrm{log}\;{\mathcal{Z}}_{\mu }+\sum _{\left(i\mu \right)\in E}\mathrm{log}\;{\mathcal{Z}}_{\mu i},\end{equation} \tag{ 63 }$$

with

$$\begin{equation}{\mathcal{Z}}_{i}=\int \mathrm{d}{x}_{i}\;{p}_{x}\left({x}_{i}\right)\prod _{\mu \in \partial i}{\tilde {m}}_{\mu \to i}\left({x}_{i}\right),\end{equation} \tag{ 64 }$$

$$\begin{equation}{\mathcal{Z}}_{\mu }=\int \prod _{i\in \partial \mu }\mathrm{d}{x}_{i}\;\psi \left({\underline{x}}_{\partial \mu }\right)\prod _{i\in \partial \mu }{m}_{i\to \mu }\left({x}_{i}\right),\end{equation} \tag{ 65 }$$

$$\begin{equation}{\mathcal{Z}}_{\mu i}=\int \mathrm{d}{x}_{i}\;{\tilde {m}}_{\mu \to i}\left({x}_{i}\right){m}_{i\to \mu }\left({x}_{i}\right).\end{equation} \tag{ 66 }$$

As for the marginals, the Bethe free energy, will only be exact if the underlying factor graph is a tree. Otherwise it is an approximation of the free energy, that is not generally an upper bound.

Belief propagation for the GLM. The writing of the BP-equations for our case-study is schematized on the right of figure 3. There are 2 × N × M updates:

$$\begin{equation}{\tilde {m}}_{\mu \to i}^{\left(t\right)}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{\mu \to i}}\int \prod _{{i}^{\prime }\ne i}\mathrm{d}{x}_{{i}^{\prime }}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu }\vert {{w}_{\mu }}^{\top }\underline{x}\right)\prod _{{i}^{\prime }\ne i}{m}_{{i}^{\prime }\to \mu }^{\left(t\right)}\left({x}_{{i}^{\prime }}\right),\end{equation} \tag{ 67 }$$

$$\begin{equation}{m}_{i\to \mu }^{\left(t+1\right)}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{i\to \mu }}{p}_{x}\left({x}_{i}\right)\prod _{{\mu }^{\prime }\ne \mu }{\tilde {m}}_{{\mu }^{\prime }\to i}^{\left(t\right)}\left({x}_{i}\right),\end{equation} \tag{ 68 }$$

for all i–μ pairs. Despite a relatively concise formulation, running BP in practice turns out intractable since for a signal $\underline{x}$ taking continuous values it would entail keeping track of distributions on continuous variables. In this case, BP is approximated by the (G)AMP algorithm presented in the next section.

The name of approximate message passing (AMP) was fixed by Donoho et al [DMM09] who derived the algorithm in the context of compressed sensing. Several works from statistical physics had nevertheless already proposed related algorithmic procedures and made connections with the TAP equations for different systems [Kab03, KS98, OS01]. The algorithm was derived systematically for any channel of the GLM by Rangan [Ran11] and became generalized-AMP (GAMP), yet again it seems that [KU04] proposed the first generalized derivation.

The systematic procedure to write AMP for a given joint probability distribution consists in first writing BP on the factor graph, second project the messages on a parametrized family of functions to obtain the corresponding relaxed-BP and third close the equations on a reduced set of parameters by keeping only leading terms in the thermodynamic limit. We will quickly review and justify these steps for the GLM. Note that here a relevant algorithm for approximate inference will be derived from message passing on a fully connected graph of interactions. As it tuns out, the high connectivity limit and the introduction of short loops does not break the assumption of independence of incoming messages in this specific case thanks to the small scale $O\left(1/\sqrt{N}\right)$ and the independence of the weight entries. The statistics of the weights are here crucial.

Relaxed belief propagation. In the thermodynamic limit M, N → +∞, one can show that the scaling $1/\sqrt{N}$ of the W_ij and the extensive connectivity of the underlying factor graph imply that messages are approximately Gaussian. Without giving all the details of the computation which can be cumbersome, let us try to provide some intuitions. We drop the time indices for simplicity and start with (67). Consider the intermediate reconstruction variable ${z}_{\mu }={\underline{w}}_{\mu }^{\top }\underline{x}={\sum }_{{i}^{\prime }\ne i}{W}_{\mu {i}^{\prime }}{x}_{{i}^{\prime }}+{W}_{\mu i}{x}_{i}$. Under the statistics of the messages m_i'→μ(x_i'), the x_i' are independent such that by the central limit theorem z_μ − W_μix_i is a Gaussian random variable with respectively mean and variance

$$\begin{equation}{\omega }_{\mu \to i}=\sum _{{i}^{\prime }\ne i}{W}_{\mu {i}^{\prime }}{\hat{x}}_{{i}^{\prime }\to \mu },\end{equation} \tag{ 69 }$$

$$\begin{equation}{V}_{\mu \to i}=\sum _{{i}^{\prime }\ne i}{W}_{\mu {i}^{\prime }}^{2}{C}_{{i}^{\prime }\to \mu }^{x},\end{equation} \tag{ 70 }$$

where we defined the mean and the variance of the messages m_i'→μ(x_i'),

$$\begin{equation}{\hat{x}}_{{i}^{\prime }\to \mu }=\int \mathrm{d}{x}_{{i}^{\prime }}\;{x}_{{i}^{\prime }}\;{m}_{{i}^{\prime }\to \mu }\left({x}_{{i}^{\prime }}\right),\end{equation} \tag{ 71 }$$

$$\begin{equation}{C}_{{i}^{\prime }\to \mu }^{x}=\int \mathrm{d}{x}_{{i}^{\prime }}\;{x}_{{i}^{\prime }}^{2}\;{m}_{{i}^{\prime }\to \mu }\left({x}_{{i}^{\prime }}\right)-{\hat{x}}_{{i}^{\prime }\to \mu }^{2}.\end{equation} \tag{ 72 }$$

Using these new definitions, (67) can be rewritten as

$$\begin{equation}{\tilde {m}}_{\mu \to i}\left({x}_{i}\right)\propto \int \mathrm{d}{z}_{\mu }\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu }\vert {z}_{\mu }\right){\mathrm{e}}^{-\frac{{\left({z}_{\mu }-{W}_{\mu i}{x}_{i}-{\omega }_{\mu \to i}\right)}^{2}}{2{V}_{\mu \to i}}},\end{equation} \tag{ 73 }$$

where the notation ∝ omits the normalization factor for distributions. Considering that W_μi is of order $1/\sqrt{N}$, the development of (73) shows that at leading order ${\tilde {m}}_{\mu \to i}\left({x}_{i}\right)$ is Gaussian:

$$\begin{equation}{\tilde {m}}_{\mu \to i}\left({x}_{i}\right)\propto {\mathrm{e}}^{{B}_{\mu \to i}{x}_{i}+\frac{1}{2}{A}_{\mu \to i}{x}_{i}^{2}}\end{equation} \tag{ 74 }$$

where the details of the computations yield

$$\begin{equation}{B}_{\mu \to i}={W}_{\mu i}\;{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu },{\omega }_{\mu \to i},{V}_{\mu \to i}\right)\end{equation} \tag{ 75 }$$

$$\begin{equation}{A}_{\mu \to i}=-{W}_{\mu i}^{2}\;{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu },{\omega }_{\mu \to i},{V}_{\mu \to i}\right)\end{equation} \tag{ 76 }$$

using the output update functions

$$\begin{equation}{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y,\omega ,V\right)=\frac{1}{{\mathcal{Z}}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\int \mathrm{d}z\;\frac{\left(z-\omega \right)}{V}{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y\vert z\right)\mathcal{N}\left(z;\omega ,V\right),\end{equation} \tag{ 77 }$$

$$\begin{equation}{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y,\omega ,V\right)=\frac{1}{{\mathcal{Z}}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\int \mathrm{d}z\;\frac{{\left(z-\omega \right)}^{2}}{{V}^{2}}{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y\vert z\right)\mathcal{N}\left(z;\omega ,V\right)-\frac{1}{V}-{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}{\left(y,\omega ,V\right)}^{2},\end{equation} \tag{ 78 }$$

$$\begin{equation}{\mathcal{Z}}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y,\omega ,V\right)=\int \mathrm{d}z\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y\vert z\right)\mathcal{N}\left(z;\omega ,V\right).\end{equation} \tag{ 79 }$$

These arguably complicated functions, again coming out of the development of (73), can be interpreted as the estimation of the mean and the variance of the gap between two different estimate of z_μ considered by the algorithm: the mean estimate ω_μ→i given incoming messages m_i'→μ(x_i') and the same mean estima-te updated to incorporate the information coming from the channel p_out and observation y_μ. Finally, the Gaussian parametrization (74) of ${\tilde {m}}_{\mu \to i}\left({x}_{i}\right)$ serves to rewrite the other type of messages m_i→μ(x_i) (68),

$$\begin{equation}{m}_{i\to \mu }\left({x}_{i}\right)\propto {p}_{x}\left({x}_{i}\right){\mathrm{e}}^{-\frac{{\left({\lambda }_{i\to \mu }-{x}_{i}\right)}^{2}}{2{\sigma }_{i\to \mu }}},\end{equation} \tag{ 80 }$$

with

$$\begin{equation}{\sigma }_{i\to \mu }={\left(\sum _{{\mu }^{\prime }\ne \mu }{A}_{{\mu }^{\prime }\to i}\right)}^{-1}\end{equation} \tag{ 81 }$$

$$\begin{equation}{\lambda }_{i\to \mu }={\sigma }_{i\to \mu }\left(\sum _{{\mu }^{\prime }\ne \mu }{B}_{{\mu }^{\prime }\to i}\right).\end{equation} \tag{ 82 }$$

The set of equations can finally be closed by recalling the definitions (71) and (72):

$$\begin{equation}{\hat{x}}_{i\to \mu }={f}_{1}^{x}\left({\lambda }_{i\to \mu },{\sigma }_{i\to \mu }\right)\end{equation} \tag{ 83 }$$

$$\begin{equation}{C}_{i\to \mu }^{x}={f}_{2}^{x}\left({\lambda }_{i\to \mu },{\sigma }_{i\to \mu }\right)\end{equation} \tag{ 84 }$$

with now the input update functions

$$\begin{equation}{\mathcal{Z}}^{x}=\int \mathrm{d}x\;{p}_{x}\left(x\right){\mathrm{e}}^{-\frac{{\left(x-\lambda \right)}^{2}}{2\sigma }},\end{equation} \tag{ 85 }$$

$$\begin{equation}{f}_{1}^{x}\left(\lambda ,\sigma \right)=\frac{1}{{\mathcal{Z}}^{x}}\int \mathrm{d}x\;x\;{p}_{x}\left(x\right){\mathrm{e}}^{-\frac{{\left(x-\lambda \right)}^{2}}{2\sigma }},\end{equation} \tag{ 86 }$$

$$\begin{equation}{f}_{2}^{x}\left(\lambda ,\sigma \right)=\frac{1}{{\mathcal{Z}}^{x}}\int \mathrm{d}x\;{x}^{2}\;{p}_{x}\left(x\right){\mathrm{e}}^{-\frac{{\left(x-\lambda \right)}^{2}}{2\sigma }}-{f}_{1}^{x}{\left(\lambda ,\sigma \right)}^{2}.\end{equation} \tag{ 87 }$$

The input update functions can be interpreted as updating the estimation of the mean and variance of the signal x_i based on the information coming from the incoming messages grasped by λ_i→μ and σ_i→μ with the information of the prior p_x.

To sum-up, by considering the leading order terms in the thermodynamic limit, the BP equations can be self-consistently re-written as a closed set of equations over mean and variance variables (69), (70), (75), (76), (81)–(84). Eventually, r-BP can equivalently be thought of as the projection of BP onto the following parametrizations of the messages

$$\begin{equation}{\tilde {m}}_{\mu \to i}^{\left(t\right)}\left({x}_{i}\right)\propto {\mathrm{e}}^{{B}_{\mu \to i}^{\left(t\right)}{x}_{i}+\frac{1}{2}{A}_{\mu \to i}^{\left(t\right)}{x}_{i}^{2}}\propto \int \mathrm{d}{z}_{\mu }\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu }\vert {z}_{\mu }\right){\mathrm{e}}^{-\frac{{\left({z}_{\mu }-{W}_{\mu i}{x}_{i}-{\omega }_{\mu \to i}^{\left(t\right)}\right)}^{2}}{2{V}_{\mu \to i}^{\left(t\right)}}},\end{equation} \tag{ 88 }$$

$$\begin{equation}{m}_{i\to \mu }^{\left(t+1\right)}\left({x}_{i}\right)\propto {\mathrm{e}}^{-\frac{{\left({\hat{x}}_{i\to \mu }^{\left(t+1\right)}-{x}_{i}\right)}^{2}}{2{{C}^{x}}_{i\to \mu }^{\left(t+1\right)}}}\propto {p}_{x}\left({x}_{i}\right){\mathrm{e}}^{-\frac{{\left({\lambda }_{i\to \mu }^{\left(t\right)}-{x}_{i}\right)}^{2}}{2{\sigma }_{i\to \mu }^{\left(t\right)}}}.\end{equation} \tag{ 89 }$$

Note that, at convergence, an approximation of the marginals is recovered from the projection on the parametrization (89) of (58),

$$\begin{equation}{m}_{i}\left({x}_{i}\right)=\frac{1}{{\mathcal{Z}}_{i}}{p}_{x}\left({x}_{i}\right){\mathrm{e}}^{-{\left({\lambda }_{i}-{x}_{i}\right)}^{2}/2{\sigma }_{i}}={f}_{1}^{x}\left({\lambda }_{i},{\sigma }_{i}\right),\end{equation} \tag{ 90 }$$

$$\begin{equation}{\sigma }_{i\to \mu }={\left(\sum _{\mu }{A}_{\mu \to i}\right)}^{-1},\end{equation} \tag{ 91 }$$

$$\begin{equation}{\lambda }_{i\to \mu }={\sigma }_{i\to \mu }\left(\sum _{\mu }{B}_{\mu \to i}\right).\end{equation} \tag{ 92 }$$

Nonetheless, r-BP is scarcely used as such as the computational cost can be readily reduced with little more approximation. Because the parameters in (88) and (89) take the form of messages on the edges of the factor graph there are still O(M × N) quantities to track to solve the self-consistency equations by iterations. Yet, in the thermodynamic limit, the messages are closely related to the marginals as the contribution of the missing message between (57) and (58) is to a certain extent negligible. Careful book keeping of the order of contributions of these small differences leads to a set of closed equations on parameters of the marginals, i.e. O(N) variables, corresponding to the GAMP algorithm.

A detailed derivation and developed algorithm of r-BP for the GLM can be found for example in [ZK16] (section 6.3.1). In section 5.3 of the present paper, we also present the derivation in a slightly more general setting where the variables x_i and y_μ are possibly vectors instead of scalars.

Generalized approximate message passing. The GAMP algorithm with respect to marginal parameters, analogous to the messages parameters introduced above (summarized in (88) and (89)), is given in algorithm 1. The origin of GAMP is again the development of the r-BP message-like equations around marginal quantities. The details of this derivation for the GLM can be found for instance in [ZK16] (section 6.3.1). For a random initialization, the algorithm can be decomposed in 4 steps per iteration which refine the estimate of the signal $\underline{x}$ and the intermediary variable $\underline{z}$ by incorporating the different sources of information. Steps (2) and (4) involve the update functions relative to the prior and output channel defined above. Steps (1) and (3) are general for any GLM with a random Gaussian weight matrix, as they result from the consistency of the two alternative parametrizations introduced for the same messages in (88) and (89).

Algorithm 1. Generalized approximate message passing

Input: vector $\underline{y}\in {\mathbb{R}}^{M}$ and matrix $\underline{\underline{W}}\in {\mathbb{R}}^{M{\times}N}$:

Initialize: ${\hat{x}}_{i}$, ${C}_{i}^{x}\quad \forall i$ and ${{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }$, ${{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }\quad \forall \mu$

repeat

(1) Estimate mean ${\omega }_{\mu }^{\left(t\right)}$ and variance ${V}_{\mu }^{\left(t\right)}$ of zμ given current ${\hat{\underline{x}}}^{\left(t\right)}$

$$\begin{equation}{V}_{\mu }^{\left(t\right)}=\sum _{i=1}^{N}{W}_{\mu i}^{2}{{C}_{i}^{x}}^{\left(t\right)}\end{equation} \tag{ 93 }$$

$$\begin{equation}{\omega }_{\mu }^{\left(t\right)}=\sum _{i=1}^{N}{W}_{\mu i}{\hat{x}}_{i}^{\left(t\right)}-\sum _{i=1}^{N}{W}_{\mu i}^{2}{{C}^{x}}_{i}^{\left(t\right)}{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t-1\right)}\end{equation} \tag{ 94 }$$

(2) Estimate mean ${{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}$ and variance ${{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}$ of the gap between optimal zμ and ${\omega }_{\mu }^{\left(t\right)}$ given yμ

$$\begin{equation}{{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}={\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu },{\omega }_{\mu }^{\left(t\right)},{V}_{\mu }^{\left(t\right)}\right)\end{equation} \tag{ 95 }$$

$$\begin{equation}{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}={g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left({y}_{\mu },{\omega }_{\mu }^{\left(t\right)},{V}_{\mu }^{\left(t\right)}\right)\end{equation} \tag{ 96 }$$

(3) Estimate mean $\underline{\lambda }$ and variance $\underline{\sigma }$ of $\underline{x}$ given current gap between optimal $\underline{z}$ and $\underline{\omega }$

$$\begin{equation}{\sigma }_{i}^{\left(t\right)}={\left(-\sum _{\mu =1}^{M}{W}_{\mu i}^{2}{{\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}\right)}^{-1}\end{equation} \tag{ 97 }$$

$$\begin{equation}{\lambda }_{i}^{\left(t\right)}={\hat{x}}_{i}^{\left(t\right)}+{\sigma }_{i}^{\left(t\right)}\left(\sum _{\mu =1}^{M}{W}_{\mu i}{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}_{\mu }^{\left(t\right)}\right)\end{equation} \tag{ 98 }$$

(4) Estimate of mean ${\hat{\underline{x}}}^{\left(t+1\right)}$ and ${{\underline{C}}^{x}}^{\left(t+1\right)}$ variance of $\underline{x}$ updated with information from the prior

$$\begin{equation}{{C}_{i}^{x}}^{\left(t+1\right)}={f}_{2}^{x}\left({\lambda }_{i}^{\left(t\right)},{\sigma }_{i}^{\left(t\right)}\right)\end{equation} \tag{ 99 }$$

$$\begin{equation}{\hat{x}}_{i}^{\left(t+1\right)}={f}_{1}^{x}\left({\lambda }_{i}^{\left(t\right)},{\sigma }_{i}^{\left(t\right)}\right)\end{equation} \tag{ 100 }$$

until convergence

Relation to TAP equations. Historically the main difference between the AMP algorithm and the TAP equations is that the latter was first derived for binary variables with two-body interactions (SK model) while the former was proposed for continuous random variables with N-body interactions (compressed sensing). The details of the derivation (described in [ZK16] or in a more general case in section 5.3), rely on the knowledge of the statistics of the disordered variable $\underline{\underline{W}}$ but do not require a disorder average, as in the Georges–Yedidia expansion yielding the TAP equations. By focusing on the GLM with a random Gaussian weight matrix scaling as $O\left(1/\sqrt{N}\right)$ (similarly to the couplings of the SK model) we naturally obtained TAP equations at second order, with an Onsager term in the update (94) of ω_μ. Yet an advantage of the AMP derivation from BP over the high-temperature expansion is that it explicitly provides 'correct' time indices in the iteration scheme to solve the self consistent equations [Bol14].

Reconstruction with AMP. AMP is therefore a practical reconstruction algorithm which can be run on a single instance (the disorder is not averaged) to estimate an unknown signal ${\underline{x}}_{0}$. Note that the prior p_x and channel p_out used in the algorithm correspond to the student statistical model and they may be different from the true underlying teacher model that generates ${\underline{x}}_{0}$ and $\underline{y}$. In other words, the AMP algorithm may be used either in the Bayes optimal or in the mismatched setting defined in section 3.1.1. Remarkably, it is also possible to consider a disorder average in the thermodynamic limit to study the average case computational hardness, here of the GLM inference problem, in either of these matched or mismatched configurations.

State evolution. The statistical analysis of the AMP equations for compressed sensing in the average case and in the thermodynamic limit N → ∞ lead to another closed set of equations that was called state evolution (SE) in [DMM09]. Such an analysis can be generalized to other problems of application of approximate message passing algorithms. The derivation of SE starts from the r-BP equations and relies on the assumption of independent incoming messages to invoke the central limit theorem. It is therefore only necessary to follow the evolution of a set of means and variances parametrizing Gaussian distributions. When the different variables and factors are statistically equivalent, as it is the case of the GLM, SE reduces to a few scalar equations. The interested reader should refer to appendix F for a detailed derivation in a more general setting.

Mismatched setting. In the general mismatched setting we need to carefully differentiate the teacher and the student. We note ${p}_{{x}_{0}}$ the prior used by the teacher. We also rewrite its channel ${p}_{\mathrm{o}\mathrm{u}\mathrm{t},0}\left(y\vert {\underline{w}}^{\top }\underline{x}\right)$ as the explicit function $y={g}_{0}\left({\underline{w}}^{\top }\underline{x};{\epsilon}\right)$ assuming the noise to be distributed according to the standard normal distribution. The tracked quantities are the overlaps,

$$\begin{equation}q=\underset{N\to \infty }{\mathrm{lim}}\frac{1}{N}\sum _{i=1}^{N}{\hat{x}}_{i}^{2},\quad m=\underset{N\to \infty }{\mathrm{lim}}\frac{1}{N}\sum _{i=1}^{N}{\hat{x}}_{i}{x}_{0,i},\quad {q}_{0}=\underset{N\to \infty }{\mathrm{lim}}\frac{1}{N}\sum _{i=1}^{N}{x}_{0,i}^{2}={\mathbb{E}}_{{p}_{{x}_{0}}}\left[{x}_{0}^{2}\right],\end{equation} \tag{ 101 }$$

along with the auxiliary V, $\hat{q}$, $\hat{m}$ and $\hat{\chi }$:

$$\begin{equation}{\hat{q}}^{\left(t\right)}=\int \mathcal{D}{\epsilon}\;\int \mathrm{d}\omega \;\mathrm{d}z\;\mathcal{N}\left(z,\omega ;0,{\underline{\underline{Q}}}^{\left(t\right)}\right){g}_{\mathrm{o}\mathrm{u}\mathrm{t}}{\left(\omega ,{g}_{0}\left(z;{\epsilon}\right),{V}^{\left(t\right)}\right)}^{2},\end{equation} \tag{ 102 }$$

$$\begin{equation}{\hat{m}}^{\left(t\right)}=\int \mathcal{D}{\epsilon}\;\int \mathrm{d}\omega \;\mathrm{d}z\;\mathcal{N}\left(z,\omega ;0,{\underline{\underline{Q}}}^{\left(t\right)}\right){\partial }_{z}{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(\omega ,{g}_{0}\left(z;{\epsilon}\right),{V}^{\left(t\right)}\right),\end{equation} \tag{ 103 }$$

$$\begin{equation}{\hat{\chi }}^{\left(t\right)}=-\int \mathcal{D}{\epsilon}\;\int \mathrm{d}\omega \;\mathrm{d}z\;\mathcal{N}\left(z,\omega ;0,{\underline{\underline{Q}}}^{\left(t\right)}\right){\partial }_{\omega }{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(\omega ,{g}_{0}\left(z;{{\epsilon}}_{\mu }\right),{V}^{\left(t\right)}\right),\end{equation} \tag{ 104 }$$

$$\begin{equation}{q}^{\left(t+1\right)}=\int \mathrm{d}{x}_{0}\;{p}_{{x}_{0}}\left({x}_{0}\right)\int \mathcal{D}\xi \;{f}_{1}^{x}{\left({\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\left(\sqrt{\alpha {\hat{q}}^{\left(t\right)}}\xi +\alpha {\hat{m}}^{\left(t\right)}{\underline{x}}_{0}\right);{\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\right)}^{2},\end{equation} \tag{ 105 }$$

$$\begin{equation}{m}^{\left(t+1\right)}=\int \mathrm{d}{x}_{0}\;{p}_{{x}_{0}}\left({x}_{0}\right)\int \mathcal{D}\xi \;{x}_{0}{f}_{1}^{x}\left({\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\left(\sqrt{\alpha {\hat{q}}^{\left(t\right)}}\xi +\alpha {\hat{m}}^{\left(t\right)}{\underline{x}}_{0}\right);{\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\right),\end{equation} \tag{ 106 }$$

$$\begin{equation}{V}^{\left(t+1\right)}=\int \mathrm{d}{x}_{0}\;{p}_{{x}_{0}}\left({x}_{0}\right)\int \mathcal{D}\xi \;{f}_{2}^{x}\left({\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\left(\sqrt{\alpha {\hat{q}}^{\left(t\right)}}\xi +\alpha {\hat{m}}^{\left(t\right)}{x}_{0}\right);{\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\right),\end{equation} \tag{ 107 }$$

where we use the notation $\mathcal{N}\left(\cdot ;\cdot ,\cdot \right)$ for the normal distribution, $\mathcal{D}\xi \;$ for the standard normal measure and the covariance matrix ${\underline{\underline{Q}}}^{\left(t\right)}$ is given at each time step by

$$\begin{equation*}{\underline{\underline{Q}}}^{\left(t\right)}=\left[\begin{bmatrix}{cc}\hfill {q}_{0}\hfill & \hfill {m}^{\left(t\right)}\hfill \\ \hfill {m}^{\left(t\right)}\hfill & \hfill {q}^{\left(t\right)}\hfill \end{bmatrix}\right].\end{equation*}$$

Due to the self-averaging property, the performance of the reconstruction by the AMP algorithm on an instance of size N can be tracked along the iterations given

$$\begin{equation}\mathrm{M}\mathrm{S}\mathrm{E}\left(\hat{x}\right)=\frac{1}{N}\sum _{i=1}^{N}{\left({\hat{x}}_{i}-{x}_{0,i}\right)}^{2}=q-2m+{q}_{0},\end{equation} \tag{ 108 }$$

with only minor differences coming from finite-size effects. State evolution also provides an efficient procedure to study from the theoretical perspective the AMP fixed points for a generic model, such as the GLM, as a function of some control parameters. It reports the average results for running the complete AMP algorithm on O(N) variables using a few scalar equations. Furthermore, the state evolution equations simplify further in the Bayes optimal setting.

Bayes optimal setting. When the prior and channel are identical for the student and the teacher, the true unknown signal ${\underline{x}}_{0}$ is in some sense statistically equivalent to the estimate $\hat{\underline{x}}$ coming from the posterior. More precisely one can prove the Nishimori identities [Iba99, Nis01, OH91] (or [KKM⁺16] for a concise demonstration and discussion) implying that q = m, V = q₀ − m and $\hat{q}=\hat{m}=\hat{\chi }$. Only two equations are then necessary to track the performance of the reconstruction:

$$\begin{equation}{\hat{q}}^{\left(t\right)}=\int \mathrm{d}{\epsilon}\;{p}_{{{\epsilon}}_{0}}\left({\epsilon}\right)\int \mathrm{d}\omega \;\mathrm{d}z\;\mathcal{N}\left(z,\omega ;0,{\underline{\underline{Q}}}^{\left(t\right)}\right){g}_{\mathrm{o}\mathrm{u}\mathrm{t}}{\left(\omega ,{g}_{0}\left(z;{\epsilon}\right),{V}^{\left(t\right)}\right)}^{2}\end{equation} \tag{ 109 }$$

$$\begin{equation}{q}^{\left(t+1\right)}=\int \mathrm{d}{x}_{0}\;{p}_{{x}_{0}}\left({x}_{0}\right)\int \mathcal{D}\xi \;{f}_{1}^{x}{\left({\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\left(\sqrt{\alpha {\hat{q}}^{\left(t\right)}}\xi +\alpha {\hat{m}}^{\left(t\right)}{\underline{x}}_{0}\right);{\left(\alpha {\hat{\chi }}^{\left(t\right)}\right)}^{-1}\right)}^{2}.\end{equation} \tag{ 110 }$$

The basic idea of the replica computation is to compute the average over the disorder of $\mathrm{log}\;\mathcal{Z}$ by considering the identity $\mathrm{log}\;\mathcal{Z}={\mathrm{lim}}_{n\to 0}\left({\mathcal{Z}}^{n}-1\right)/n$. First the expectation of ${\mathcal{Z}}^{n}$ is evaluated for $n\in \mathbb{N}$, then the n → 0 limit is taken by 'analytic continuation'. Thus the method takes advantage of the fact that the average of a power of $\mathcal{Z}$ is sometimes easier to compute than the average of a logarithm. We illustrate the key steps of the calculation for the partition function of the GLM (52).

Disorder average for the replicated system: coupling of the replicas. The average of ${\mathcal{Z}}^{n}$ for $n\in \mathbb{N}$ can be seen as the partition function of a system with n + 1 non interacting replicas of $\underline{x}$ indexed by a ∈ {0, ..., n}, where the first replica a = 0 is representative of the teacher and the n other replicas are identically distributed as the student:

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]={\mathbb{E}}_{\underline{\underline{W}}}\left[\int \mathrm{d}\underline{y}\;\mathrm{d}{\underline{x}}_{0}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t},0}\left(\underline{y}\vert \underline{\underline{W}}{\underline{x}}_{0}\right){p}_{{\underline{x}}_{0}}\left({\underline{x}}_{0}\right){\left(\int \mathrm{d}\underline{x}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(\underline{y}\vert \underline{\underline{W}}\underline{x}\right){p}_{x}\left(\underline{x}\right)\right)}^{n}\right]\end{equation} \tag{ 111 }$$

$$\begin{equation}={\mathbb{E}}_{\underline{\underline{W}}}\left[\int \mathrm{d}\underline{y}\;\prod _{a=0}^{n}\left(\mathrm{d}{\underline{x}}_{a}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t},a}\left(\underline{y}\vert \underline{\underline{W}}{\underline{x}}_{a}\right){p}_{{x}_{a}}\left({\underline{x}}_{a}\right)\right)\right]\end{equation} \tag{ 112 }$$

$$\begin{equation}={\mathbb{E}}_{\underline{\underline{W}}}\left[\int \mathrm{d}\underline{y}\;\prod _{a=0}^{n}\left(\mathrm{d}{\underline{x}}_{a}\;\mathrm{d}{\underline{z}}_{a}\;\delta \left({\underline{z}}_{a}-\underline{\underline{W}}{\underline{x}}_{a}\right){p}_{\mathrm{o}\mathrm{u}\mathrm{t},a}\left(\underline{y}\vert {\underline{z}}_{a}\right){p}_{{x}_{a}}\left({\underline{x}}_{a}\right)\right)\right].\end{equation} \tag{ 113 }$$

To perform the average over the disordered interactions $\underline{\underline{W}}$ we consider the statistics of ${\underline{z}}_{a}=\underline{\underline{W}}{\underline{x}}_{a}$. Recall that ${W}_{\mu i}\sim \mathcal{N}\left({W}_{\mu i};0,1/N\right)$, independently for all μ and i. Consequently, the ${\underline{z}}_{a}$ are jointly Gaussian in the thermodynamic limit with means and covariances

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{z}_{a,\mu }\right]={\mathbb{E}}_{\underline{\underline{W}}}\left[\sum _{i=1}^{N}{W}_{\mu i}{x}_{a,i}\right]=0,\quad {E}_{\underline{\underline{W}}}\left[{z}_{a,\mu }{z}_{b,\nu }\right]=\sum _{i=1}^{N}{x}_{a,i}{x}_{b,i}/N={q}_{ab}.\end{equation} \tag{ 114 }$$

The overlaps, that we already introduced in the SE formalism, naturally re-appear. We introduce the notation $\underline{\underline{q}}$ for the (n + 1) × (n + 1) overlap matrix. Integrating out the disorder $\underline{\underline{W}}$ shared by the n + 1 replicas will therefore leave us with an effective system of now coupled replicas:

$$\begin{align}\hfill {\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]=& \int \prod _{a,b}\mathrm{d}N{q}_{ab}\;\int \mathrm{d}\underline{y}\;\prod _{a=0}^{n}\mathrm{d}{\underline{z}}_{a}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t},a}\left(\underline{y}\vert {\underline{z}}_{a}\right)\mathrm{exp}\left(-\frac{1}{2}\sum _{\mu =1}^{M}\sum _{a,b}{z}_{a,\mu }{z}_{b,\mu }{\left({\underline{\underline{q}}}^{-1}\right)}_{ab}-MC\left(\underline{\underline{q}},n\right)\right)\hfill \\ \hfill & {\times}\;\int \prod _{a=1}^{n}\mathrm{d}{\underline{x}}_{a}\;{p}_{{x}_{a}}\left({\underline{x}}_{a}\right)\delta \left(N{q}_{ab}-\sum _{i=1}^{N}{x}_{a,i}{x}_{b,i}\right).\hfill \end{align} \tag{ 115 }$$

Change of variable for the overlaps: decoupling of the variables. We consider the Fourier representation of the Dirac distribution fixing the consistency between overlaps and replicas,

$$\begin{equation}\delta \left(N{q}_{ab}-\sum _{i=1}^{N}{x}_{a,i}{x}_{b,i}\right)=\int \frac{\mathrm{d}{\hat{q}}_{ab}}{2i\pi }\;{\mathrm{e}}^{{\hat{q}}_{ab}\left(N{q}_{ab}-\sum _{i=1}^{N}{x}_{a,i}{x}_{b,i}\right)},\end{equation} \tag{ 116 }$$

where ${\hat{q}}_{ab}$ is purely imaginary, which yields

$$\begin{align}\hfill {\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]=& \int \prod _{a,b}\mathrm{d}N{q}_{ab}\;\int \prod _{a,b}\frac{\mathrm{d}{\hat{q}}_{ab}}{2i\pi }\;\mathrm{exp}\left(N{\hat{q}}_{ab}{q}_{ab}\right)\hfill \\ \hfill & {\times}\int \mathrm{d}\underline{y}\;\prod _{a=0}^{n}\mathrm{d}{\underline{z}}_{a}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t},a}\left(\underline{y}\vert {\underline{z}}_{a}\right)\mathrm{exp}\left(-\frac{1}{2}\sum _{\mu =1}^{M}\sum _{a,b}{z}_{a,\mu }{z}_{b,\mu }{\left({\underline{\underline{q}}}^{-1}\right)}_{ab}-MC\left(\underline{\underline{q}},n\right)\right)\hfill \\ \hfill & {\times}\;\int \prod _{a=1}^{n}\mathrm{d}{\underline{x}}_{a}\;{p}_{{x}_{a}}\left({\underline{x}}_{a}\right)\mathrm{exp}\left(-{\hat{q}}_{ab}\sum _{i=1}^{N}{x}_{a,i}{x}_{b,i}\right)\hfill \end{align} \tag{ 117 }$$

where $C\left(\underline{\underline{q}},n\right)$ is related to the normalization of the Gaussian distributions over the ${\underline{z}}_{a}$ variables, and the integrals can be factorized over the i–s and μ–s. Thus we obtain

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]=\int \prod _{a,b}\mathrm{d}N{q}_{ab}\;\int \prod _{a,b}\mathrm{d}{\hat{q}}_{ab}\;\;{\mathrm{e}}^{N{\hat{q}}_{ab}{q}_{ab}}\;{\mathrm{e}}^{M\mathrm{log}{\hat{\mathcal{I}}}_{z}\left(\underline{\underline{q}}\right)}\;{\mathrm{e}}^{N\mathrm{log}{\hat{\mathcal{I}}}_{x}\left(\hat{\underline{\underline{q}}}\right)},\end{equation} \tag{ 118 }$$

with

$$\begin{equation}{\hat{\mathcal{I}}}_{z}\left(\underline{\underline{q}}\right)=\int \mathrm{d}y\;\prod _{a=0}^{n}\mathrm{d}{z}_{a}\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t},a}\left(y\vert {z}_{a}\right)\mathrm{exp}\left(-\frac{1}{2}\sum _{a,b}{z}_{a}{z}_{b}{\left({q}_{ab}\right)}^{-1}-C\left(\underline{\underline{q}},n\right)\right),\end{equation} \tag{ 119 }$$

$$\begin{equation}{\hat{\mathcal{I}}}_{x}\left(\hat{\underline{\underline{q}}}\right)=\int \prod _{a=1}^{n}\mathrm{d}{x}_{a}\;{p}_{{x}_{a}}\left({x}_{a}\right)\mathrm{exp}\left(-{\hat{q}}_{ab}{x}_{a}{x}_{b}\right),\end{equation} \tag{ 120 }$$

where we introduce the notation $\hat{\underline{\underline{q}}}$ for the auxiliary overlap matrix with entries ${\left(\hat{\underline{\underline{q}}}\right)}_{ab}={\hat{q}}_{ab}$ and we omitted the factor 2iπ which is eventually subleading as N → +∞. The decoupling of the x_i and the z_μ of the infinite range system yields pre-factors N and M in the exponential arguments. In the thermodynamic limit, we recall that both N and M tend to +∞ while the ratio α = M/N remains fixed. Hence, the integral for the replicated average is easily computed in this limit by the saddle point method:

$$\begin{equation}\mathrm{log}\;{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]\simeq N\;{\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{\underline{\underline{q}}\hat{\underline{\underline{q}}}}\left[\mathcal{\phi }\left(\underline{\underline{q}},\hat{\underline{\underline{q}}}\right)\right],\quad \mathcal{\phi }\left(\underline{\underline{q}},\hat{\underline{\underline{q}}}\right)=\sum _{a,b}{\hat{q}}_{ab}{q}_{ab}+\alpha {\hat{\mathcal{I}}}_{z}\left(\underline{\underline{q}}\right)+{\hat{\mathcal{I}}}_{x}\left(\hat{\underline{\underline{q}}}\right),\end{equation} \tag{ 121 }$$

where we defined the replica potential ϕ.

Exchange of limits: back to the quenched average. The thermodynamic average of the log-partition is recovered through an a priori risky mathematical manipulation: (i) perform an analytical continuation from $n\in \mathbb{N}$ to n → 0

$$\begin{equation}\frac{1}{N}{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[\mathrm{log}\;\mathcal{Z}\right]=\underset{n\to 0}{\mathrm{lim}}\frac{1}{nN}{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}-1\right]=\underset{n\to 0}{\mathrm{lim}}\frac{1}{nN}\;\mathrm{log}\;{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]\end{equation} \tag{ 122 }$$

and (ii) exchange limits

$$\begin{equation}-f=\underset{N\to \infty }{\mathrm{lim}}\underset{n\to 0}{\mathrm{lim}}\frac{1}{n}\frac{1}{N}\;\mathrm{log}\;{\mathbb{E}}_{\underline{\underline{W}},\underline{y},{\underline{x}}_{0}}\left[{\mathcal{Z}}^{n}\right]=\underset{n\to 0}{\mathrm{lim}}\frac{1}{n}\;{\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{\underline{\underline{q}}\hat{\underline{\underline{q}}}}\left[\mathcal{\phi }\left(\underline{\underline{q}},\hat{\underline{\underline{q}}}\right)\right].\end{equation} \tag{ 123 }$$

Despite the apparent lack of rigor in taking these last steps, the replica method has been proven to yield exact predictions in the thermodynamic limit for different problems and in particular for the GLM [BKM⁺18, RP16].

Saddle point solution: choice of a replica ansatz. At this point, we are still left with the problem of computing the extrema of $\mathcal{\phi }\left(\underline{\underline{q}},\hat{\underline{\underline{q}}}\right)$. To solve this optimization problem over $\underline{\underline{q}}$ and $\hat{\underline{\underline{q}}}$, a natural assumption is that replicas, that are a pure artifact of the calculation, are equivalent. This is reflected in a special structure for overlap matrices between replicas that only depend on three parameters each,

$$\begin{equation}\underline{\underline{q}}=\left[\begin{bmatrix}{cccc}\hfill {q}_{0}\hfill & \hfill m\hfill & \hfill m\hfill & \hfill m\hfill \\ \hfill m\hfill & \hfill q\hfill & \hfill {q}_{12}\hfill & \hfill {q}_{12}\hfill \\ \hfill m\hfill & \hfill {q}_{12}\hfill & \hfill q\hfill & \hfill {q}_{12}\hfill \\ \hfill m\hfill & \hfill {q}_{12}\hfill & \hfill {q}_{12}\hfill & \hfill q\hfill \end{bmatrix}\right],\quad \hat{\underline{\underline{q}}}=\left[\begin{bmatrix}{cccc}\hfill \hat{{q}_{0}}\hfill & \hfill \hat{m}\hfill & \hfill \hat{m}\hfill & \hfill \hat{m}\hfill \\ \hfill \hat{m}\hfill & \hfill \hat{q}\hfill & \hfill {\hat{q}}_{12}\hfill & \hfill {\hat{q}}_{12}\hfill \\ \hfill \hat{m}\hfill & \hfill {\hat{q}}_{12}\hfill & \hfill \hat{q}\hfill & \hfill {\hat{q}}_{12}\hfill \\ \hfill \hat{m}\hfill & \hfill {\hat{q}}_{12}\hfill & \hfill {\hat{q}}_{12}\hfill & \hfill \hat{q}\hfill \end{bmatrix}\right],\end{equation} \tag{ 124 }$$

here given as an example for n = 3 replicas. Plugging this replica symmetric (RS) ansatz in the expression of $\mathcal{\phi }\left(\underline{\underline{q}},\hat{\underline{\underline{q}}}\right)$, taking the limit n → 0 and looking for the stationary points as a function of the parameters q, m, q₁₂ and $\hat{m}$, $\hat{q}$, ${\hat{q}}_{12}$ recovers a set of equations equivalent to SE (7), albeit without time indices. Hence the two a priori different heuristics of BP and the replica method are remarkably consistent under the RS assumption.

Nevertheless, the replica symmetry can be spontaneously broken in the large N limit and the dominating saddle point does not necessarily correspond to the RS overlap matrix. This replica symmetry breaking (RSB) corresponds to substantial changes in the structure of the examined Boltzmann distribution. It is among the great strengths of the replica formalism to naturally capture it. Yet for inference problems falling under the teacher–student scenario, the correct ansatz is always replica symmetric in the Bayes optimal setting [CCFM05, Nis01, ZK16], and we will not investigate here this direction further. The interested reader can refer to the classical references for an introduction to replica symmetry breaking [CCFM05, MPV86, Nis01] in the context of the theory of spin-glasses.

Bayes optimal setting. As in SE the equations simplify in the matched setting, where the first replica corresponding to the teacher becomes equivalent to all the others. The replica free energy of the GLM is then given as the extremum of a potential over two scalar variables:

$$\begin{equation}-f={\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{q\hat{q}}\left[-\frac{1}{2}q\hat{q}+{\mathcal{I}}_{x}\left(\hat{q}\right)+\alpha {\mathcal{I}}_{z}\left({q}_{0},q\right)\right]\end{equation} \tag{ 125 }$$

$$\begin{equation}{\mathcal{I}}_{x}\left(\hat{q}\right)=\int \mathcal{D}\xi \;\mathrm{d}x\;{p}_{x}\left(x\right){\mathrm{e}}^{-\hat{q}\frac{{x}^{2}}{2}+\sqrt{\hat{q}}\xi x}\;\mathrm{log}\left(\int \mathrm{d}{x}^{\prime }\;{p}_{x}\left({x}^{\prime }\right){\mathrm{e}}^{-\hat{q}\frac{{x}^{\prime 2}}{2}+\sqrt{\hat{q}}\xi {x}^{\prime }}\right)\end{equation} \tag{ 126 }$$

$$\begin{equation}{\mathcal{I}}_{z}\left(q,{q}_{0}\right)=\int \mathcal{D}\xi \;\mathrm{d}y\;\mathrm{d}z\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y\vert z\right)\mathcal{N}\left(z;\sqrt{q}\xi ,{q}_{0}-q\right)\mathrm{log}\left(\int \mathrm{d}{z}^{\prime }\;{p}_{\mathrm{o}\mathrm{u}\mathrm{t}}\left(y\vert {z}^{\prime }\right)\mathcal{N}\left({z}^{\prime };\sqrt{q}\xi ,{q}_{0}-q\right)\right).\end{equation} \tag{ 127 }$$

The saddle point equations corresponding to the extremization (125), fixing the values of q and $\hat{q}$, would again be found equivalent to the Bayes optimal SE (109) and (110). This Bayes optimal result is derived in [KMS⁺12] for the case of a linear channel and Gauss–Bernoulli prior, and can also be recovered as a special case of the low-rank matrix factorization formula (where the measurement matrix is in fact known) [KKM⁺16].

A crucial point in the above derivation of the replica formula is the extensivity of the interactions of the infinite range model that allowed the factorization of the N scaling of the argument of the exponential integrand in (118). The statistics of the disorder $\underline{\underline{W}}$ and in particular the independence of all the W_μi was also necessary. While this is an important assumption for the technique to go through, it can be possible to relax it for some types of correlation statistics, as we will see in section 5.2.

Note that the replica method directly enforces the disorder averaging and does not provide a prediction at the level of the single instance. Therefore it cannot be turned into a practical algorithm of reconstruction. Nonetheless, we have seen that the saddle point equations of the replica derivation, under the RS assumption, matches the SE equations derived from BP. This is sufficient to theoretically study inference questions under a teacher–student scenario in the Bayes optimal setting, and in particular predict the MSE following (108).

In the mismatched setting however, the predictions of the replica method under the RS assumption and the equivalent BP conclusions can be wrong. By introducing the symmetry breaking between replicas, the method can sometimes be corrected. It is an important endeavor of the replica formalism to grasp the importance of the overlaps and connect the form of the replica ansatz to the properties of the joint probability distribution examined. When BP fails on loopy graphs, correlations between variables are not decaying with distance, which manifests into an RSB phase. Note that there also exists message passing algorithms operating in this regime [AFUZ19, AKUZ19, MM09, MP01, MPZ02, SLL19].

Consider the ensemble of orthogonally invariant weight matrices $\underline{\underline{W}}$ with spectral density ${\sum }_{i=1}^{N}\delta \left(\lambda -{\lambda }_{i}\right)/N$ of their 'square' $\underline{\underline{W}}{\underline{\underline{W}}}^{\top }$ converging in the thermodynamic limit N → +∞ to a given density ρ_λ(λ). The quenched free energy of the GLM in the Bayes optimal setting derived in [Kab08, SK09] writes

$$\begin{equation}-f={\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{q\hat{q}}\left[-\frac{1}{2}q\hat{q}+{\mathcal{I}}_{x}\left(\hat{q}\right)+{\mathcal{J}}_{z}\left({q}_{0},q,\alpha ,{\rho }_{\lambda }\right)\right],\end{equation} \tag{ 131 }$$

$$\begin{equation}{\mathcal{J}}_{z}\left({q}_{0},q,\alpha ,{\rho }_{\lambda }\right)={\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{u\hat{u}}\left[{F}_{{\rho }_{\lambda },\alpha }\left({q}_{0}-q,\hat{u}/{\lambda }_{0}\right)+\frac{\hat{u}{q}_{0}}{2}-\frac{\alpha \hat{u}u}{2{\lambda }_{0}}+\alpha {\mathcal{I}}_{z}\left({q}_{0}{\lambda }_{0}/\alpha ,u\right)\right],\end{equation} \tag{ 132 }$$

where ${\mathcal{I}}_{x}$ and ${\mathcal{I}}_{z}$ were defined as (126) and (127) and the spectral density ρ_λ(λ) appears via its mean ${\lambda }_{0}={\mathbb{E}}_{\lambda }\left[\lambda \right]$ and in the definition of

$$\begin{align}\hfill {F}_{{\rho }_{\lambda },\alpha }\left(q,u\right)& =\frac{1}{2}{\mathrm{e}\mathrm{x}\mathrm{t}\mathrm{r}}_{{{\Lambda}}_{q},{{\Lambda}}_{u}}\left[-\left(\alpha -1\right)\mathrm{log}\;{{\Lambda}}_{u}-{\mathbb{E}}_{\lambda }\;\mathrm{log}\left({{\Lambda}}_{u}{{\Lambda}}_{q}+\lambda \right)+{{\Lambda}}_{q}q+\alpha {{\Lambda}}_{u}u\right]\hfill \\ \hfill & \hfill -\frac{1}{2}\left(\mathrm{log}\;q+1\right)+\frac{\alpha }{2}\left(\mathrm{log}\;u+1\right).\end{align} \tag{ 133 }$$

Gaussian random matrices are a particular case of the considered ensemble. Their singular values are characterized asymptotically by the Marcenko–Pastur distribution [MP67]. In this case, one can check that the above expression reduces to (125). More generally, note that ${\mathcal{J}}_{z}$ generalizes ${\mathcal{I}}_{z}$.

The VAMP algorithm consists in writing EP [Min01] with Gaussian messages on the factor graph representation of the GLM posterior distribution given in figure 5. The estimation of the signal $\underline{x}$ is decomposed onto four variables, two duplicates of $\underline{x}$ itself and two duplicates of the linear transformation $\underline{z}=\underline{\underline{W}}\underline{x}$. The potential functions ψ_x and ψ_z of factors connecting copies of the same variable are Dirac distributions enforcing their equality. The factor node linking ${\underline{z}}^{\left(2\right)}$ and ${\underline{x}}^{\left(2\right)}$ is assumed Gaussian with variance going to zero.The procedure of derivation, equivalent to the projection of the BP algorithm on Gaussian messages, is recalled in appendix E and leads to algorithm 2. Like for AMP, the algorithm features some auxiliary variables introduced along the derivation. At convergence the duplicated ${\hat{\underline{x}}}_{1}$, ${\hat{\underline{x}}}_{2}$ (and ${\hat{\underline{z}}}_{1}$, ${\hat{\underline{z}}}_{2}$) are equal and either can be returned by the algorithm as an estimator. For readability, we omitted the time indices in the iterations that here simply follow the indicated update.

Zoom In Zoom Out Reset image size

Algorithm 2. Vector approximate message passing

Input: vector of observations $\underline{y}\in {\mathbb{R}}^{M}$ and weight matrix $\underline{\underline{W}}\in {\mathbb{R}}^{M{\times}N}$:

Initialize: ${\underline{\underline{A}}}_{1}^{x}$, ${\underline{B}}_{1}^{x}$, ${\underline{\underline{A}}}_{1}^{z}$, ${\underline{B}}_{1}^{z}$

repeat

$$\begin{equation}{\hat{\underline{x}}}_{1}={f}_{1}^{x}\left({\underline{B}}_{1}^{x},{\underline{\underline{A}}}_{1}^{x}\right),\quad {\underline{\underline{{C}^{x}}}}_{1}={f}_{2}^{x}\left({\underline{B}}_{1}^{x},{\underline{\underline{A}}}_{1}^{x}\right)\end{equation} \tag{ 134 }$$

$$\begin{equation}{\underline{\underline{A}}}_{2}^{x}={{\underline{\underline{{C}^{x}}}}_{1}}^{-1}-{\underline{\underline{A}}}_{1}^{x},\quad {\underline{B}}_{2}^{x}={{\underline{\underline{{C}^{x}}}}_{1}}^{-1}{\hat{\underline{x}}}_{2}-{\underline{B}}_{1}^{x}\end{equation} \tag{ 135 }$$

$$\begin{equation}{\hat{\underline{z}}}_{1}={f}_{1}^{z}\left({\underline{B}}_{1}^{z},{\underline{\underline{A}}}_{1}^{z}\right),\quad {{\underline{\underline{C}}}^{z}}_{1}={f}_{2}^{z}\left({\underline{B}}_{1}^{z},{\underline{\underline{A}}}_{1}^{z}\right)\end{equation} \tag{ 136 }$$

$$\begin{equation}{\underline{\underline{A}}}_{2}^{z}={{\underline{\underline{{C}^{x}}}}_{1}}^{-1}-{\underline{\underline{A}}}_{1}^{z},\quad {\underline{B}}_{2}^{z}={{\underline{\underline{{C}^{x}}}}_{1}}^{-1}{\hat{\underline{x}}}_{2}-{\underline{B}}_{1}^{z}\end{equation} \tag{ 137 }$$

$$\begin{equation}{\hat{\underline{x}}}_{2}={g}_{1}^{x}\left({\underline{B}}_{2}^{x},{A}_{2}^{x},{\underline{B}}_{2}^{z},{A}_{2}^{z}\right),\quad {\underline{\underline{{C}^{x}}}}_{2}={g}_{2}^{x}\left({\underline{B}}_{2}^{x},{A}_{2}^{x},{\underline{B}}_{2}^{z},{A}_{2}^{z}\right)\end{equation} \tag{ 138 }$$

$$\begin{equation}{\underline{\underline{A}}}_{1}^{x}={{\underline{\underline{{C}^{x}}}}_{2}}^{-1}-{\underline{\underline{A}}}_{2}^{x},\quad {\underline{B}}_{1}^{x}={{\underline{\underline{{C}^{x}}}}_{2}}^{-1}{\hat{\underline{x}}}_{1}-{\underline{B}}_{2}^{x}\end{equation} \tag{ 139 }$$

$$\begin{equation}{\hat{\underline{z}}}_{2}={g}_{1}^{z}\left({\underline{B}}_{2}^{x},{A}_{2}^{x},{\underline{B}}_{2}^{z},{A}_{2}^{z}\right),\quad {{\underline{\underline{C}}}^{z}}_{2}={g}_{2}^{z}\left({\underline{B}}_{2}^{x},{A}_{2}^{x},{\underline{B}}_{2}^{z},{A}_{2}^{z}\right)\end{equation} \tag{ 140 }$$

$$\begin{equation}{\underline{\underline{A}}}_{1}^{z}={{\underline{\underline{{C}^{x}}}}_{2}}^{-1}-{\underline{\underline{A}}}_{2}^{z},\quad {\underline{B}}_{1}^{z}={{\underline{\underline{{C}^{x}}}}_{2}}^{-1}{\hat{\underline{x}}}_{1}-{\underline{B}}_{2}^{z}\end{equation} \tag{ 141 }$$

until convergence

Output: signal estimate ${\hat{\underline{x}}}_{1}\in {\mathbb{R}}^{N}$, and estimated covariance ${\underline{\underline{{C}^{x}}}}_{1}\in {\mathbb{R}}^{N{\times}N}$

For a given instance of the GLM inference problem, i.e. a given weight matrix $\underline{\underline{W}}$, one can always launch either the AMP algorithm or the VAMP algorithm to attempt the reconstruction. If the weight matrix has i.i.d. zero mean Gaussian entries, the two strategies are conjectured to be equivalent and GAMP can be provably convergent for certain settings [RSF14]. If the weight matrix is not Gaussian but orthogonally invariant, then only VAMP is expected to always converge. More generally, even in cases where none of these assumptions are verified, VAMP has been observed to have less convergence issues than AMP.

Like for AMP, a state evolution can also be derived for VAMP (which was actually directly proposed for the multi-layer GLM [FRS18]). It rigorously characterizes the behavior of the algorithm when $\underline{\underline{W}}$ is orthogonally invariant. One can also verify that the SE fixed points can be mapped to the solutions of the saddle point equations of the replica free energy (131) (see section 1 of supplementary material of [GML⁺18]); so that the robust algorithmic procedure can advantageously be used to compute the fixed points to be plugged in the replica potential to approximate the free energy.

Fundamental study of learning. Given their similarity with the Ising model, restricted Boltzmann machines have unsurprisingly attracted a lot of interest. Studying an ensemble of RBMs with random parameters using the replica method, Tubiana and Monasson [TM17] evidenced different regimes of typical pattern of activations in the hidden units and identified control parameters as the sparsity of the weights, their strength (playing the role of an effective temperature) or the type of prior for the hidden layer. Their study contributes to the understanding of the conditions under which the RBMs can represent high-order correlations between input units, albeit without including data and learning in their model. Barra and collaborators [BGST17, BGST18], exploited the connections between the Hopfield model and RBMs to characterize RBM learning understood as an associative memory. Relying again on replica calculations, they characterize the retrieval phase of RBMs. Mézard [Méz17] also re-examined retrieval in the Hopfield model using its RBM representation and message passing, showing in particular that the addition of correlations between memorized patterns could still allow for a mean-field treatment at the price of a supplementary hidden layer in the Boltzmann machine representation. This result remarkably draws a theoretical link between correlations in the training data and the necessity of depth in neural network models.

While the above results do not include characterization of the learning driven by data, a few others were able to discuss the dynamics of training. Huang [Hua17] studied with the replica method and TAP equations the Bayesian leaning of an RBM with a single hidden unit and binary weights. Barra and collaborators [BGST17] empirically studied a teacher–student scenario of unsupervised learning by maximum likelihood on samples of an Hopfield model which they could compare to their theoretical characterization of the retrieval phase. Decelle and collaborators [DFF17, DFF18] introduced an ensemble of RBMs characterized by the spectral properties of the weight matrix and derived the typical dynamics of the corresponding order parameters during learning driven by data. Beyond RBMs, analyses of the learning in other generative models are starting to appear [WHLP18].

Training algorithm based on mean-field methods. Beyond bringing theoretical insights, mean-field methods are also found useful to build tractable estimators of the likelihood in generative models, which in turn serves to design novel training algorithms.

For Boltzmann machines, this direction was already investigated in the 80s and 90s, [Gal93, Hin89, KR98, PA87], albeit in small models with binary units and for artificial data sets very different from modern machine learning benchmarks. More recently, a deterministic training based on naive mean-field was tested on RBMs [Tie08, WH02]. On toy deep learning data sets, the algorithm was found to perform poorly when compared to both CD and PCD, the commonly employed approximate Monte Carlo methods. However going beyond naive mean-field, considering the second order TAP expansion, allows to bridge the gap in efficiency [GTK15, TGM⁺18]. Additionally, the deterministic mean-field framework offers a tractable way of evaluating the learning success by exploiting the mean-field observables to visualize the representation learned by RBMs. Interestingly, high temperature expansions of estimators different from the maximum likelihood have also been recently proposed as efficient inference method for the inverse Ising problem [LVMC18].

By construction, variational auto-encoders (VAEs) rely on a variational approximation of the likelihood. In practice, the posterior distribution of the latent representation given an input (see section 2.2) is typically approached by a factorized Gaussian distribution with mean and variance parametrized by neural networks. The factorization assumption relates the method to a naive mean-field approximation.

Structured Bayesian priors. With the progress of unsupervised learning, the idea of using generative models as expressive priors has emerged.

For reconstruction tasks, in the event where a set of typical signals is available a priori, the latter can serve as a training set to learn a model of the underlying data distribution with a generative model. Subsequently, in the reconstruction of a new signal of the same type, the generative model can serve as a Bayesian prior. In particular, the idea to exploit RBMs in CS applications was pioneered by [DHD12] and [TDK16], who trained binary RBMs using contrastive divergence (to locate the support of the non-zero entries of sparse signals) and combined it with an AMP reconstruction. They demonstrated drastic improvements in the reconstruction with structured learned priors compared to the usual sparse unstructured priors. The approach, requiring to combine the AMP reconstruction for CS and the RBM TAP inference, was further generalized in [TGM+18, TMC⁺16] to real valued distributions. In the line of these applications, several works have also investigated using feed forward generative models for inference tasks. Using this time multi-layer VAMP inference, Rangan and co-authors [PSRF19] showed that VAEs could help for in-painting partially observed images. Note also that a different line of works, mainly considering GANs, examined the same type of applications without resorting to mean-field algorithms [BJPD17, HLV18, HV18, MV18]. Instead they performed the inference via gradient descent and back-propagation.

Another application of generative priors is to model synthetic data sets with structure. In [ALM⁺19, GMKZ19, GML⁺18], the authors designed learning problems amenable to a mean-field theoretical treatment by assuming the inputs to be drawn from a generative prior (albeit with untrained weights so far). This approach goes beyond the vanilla teacher–student scenario where input data is typically unstructured with i.i.d. components. This is a crucial direction of research as the role of structure in data appears as an important component to understand the puzzling efficiency of deep learning.

New results in the replica analysis of learning. The classical replica analysis of learning with simple architectures, following bases set by Gardner and Derrida 30 years ago, continues to be explored. Among the most prominent results, Kabashima and collaborators [Kab08, SK08, SK09] extended the mean-field treatment of the perceptron from data matrices with i.i.d. entries to random orthogonal matrices. It is a much larger class of random matrices where matrix entries can be correlated. More recently, a series of works explored in depth the specific case of the perceptron with binary weight values for classification on random inputs. Replica computations showed that the space of solutions is dominated in the thermodynamic limit by isolated solutions [HK14, HWK13], but also that subdominant dense clusters of solutions exist with good generalization properties in the teacher–student scenario case [BBC⁺16, BGK⁺18, BIL⁺15]. This observation inspired a novel training algorithm [CCS⁺17]. The simple two-layer architecture of the committee machine was also reexamined recently [AMB⁺18]. In the teacher–student scenario, a computationally hard phase of learning was evidenced by comparing a message passing algorithm (believed to be optimal) and the replica prediction. In this work, the authors also proposed a strategy of proof of the replica prediction.

Signal propagation in depth. Mean-field approximations can also help understand the role and implications of depth by characterizing signal propagation in neural networks. The following papers consider the width of each layer to go to infinity. In this limit, Sompolinsky and collaborators characterized how neural networks manage to progressively separate data manifolds fed as inputs [CCLS19, CLS18, KS16]. Another line of works focused on the initialization of neural networks (i.e. with random weights), and found an order-to-chaos transition in the signal propagation as a function of hyperparameters of training [PLR⁺16, SBG⁺17]. As a result, the authors could formulate recommendations for combinations of hyperparameters to practitioners. This type of analysis could furthermore be generalized to convolutional networks [NXL⁺19], recurrent networks [GCC⁺19] and networks with batch-normalization regularization [YPR⁺19]. The space of functions spanned by deep random networks in the infinite-size limit was also studied by [LS18, LS20], using the different but related approach of the generating functional analysis. Yet another mean-field argument, this time relying on a replica computation, allowed to compute the mutual information between layers of large non-linear deep neural networks with orthogonally invariant weight matrices [GML⁺18]. Using this method, mutual informations can be followed precisely along the learning for an appropriate teacher–student scenario. The strategy offers an experimental test bed to characterize possible links between the generalization ability of deep neural networks and information compression phases in the training (see [SBD⁺18, SZT17, TZ15]).

Dynamics of SGD learning in simple networks and generalization. A number of different mean-field limits led to interesting analyses of the dynamics of gradient descent learning. In particular, the below mentioned works contribute to shed light on the generalization power of neural networks in the so-called overparametrized regimes, that is where the number of parameters exceeds largely either the number of training points or the underlying degrees of freedom of the teacher rule. In linear networks first, an exact description in the high-dimensional limit was obtained for the teacher-student setup by [AS17] using random matrix theory. The generalization was predicted to improve with the overparametrization of the student. Non-linear networks with one infinitely wide hidden layer were considered by [CB18b, MMN18, RVE18, SS18] who showed that gradient descent converges to a finite generalization error. Their results are related to others obtained in a slightly different limit of infinitely large neural networks [JGH18]. For arbitrarily deep networks, Jacot and collaborators [JGH18] showed that, in a certain setting, gradient descent was effectively performing a kernel regression with a kernel function converging to a fixed value for the entire training as the size of the layers increases. In both related limits, the absence of divergence is accounting for generalization not deteriorating despite of the explosion of the number of parameters. The relationship between the two above limits was discussed in [CB18a, GSJW19, MMM19]. Subsequent works, leveraged the formalism introduced in [JGH18]. Scaling for the generalization error as a function of network sizes were derived by [GJS⁺19]. Other authors focused on the characterization of the network output function in this limit, which takes the form of a Gaussian process [LXS⁺19]. This fact was probably first noticed by Opper and Winther with one hidden layer [OW99b], to whom it inspired a TAP based Bayesian classification method using Gaussian processes. Finally, yet another limit was analyzed by [GAS⁺19], considering a finite number of hidden units with an infinitely wide input. Following classical works on the mean-field analysis of online learning (not covered in the previous sections [BS95, Saa99, SS95a, SS95b]), a closed set of equations can be derived and analyzed for a collection of overlaps. Note that these are the same order parameters as in replica computations. The resulting learning curves evidence the necessity of multi-layer learning to observe the improvement of generalization with overparametrization. An interplay between optimization, architecture and data sets seems necessary to explain the phenomenon.

Definition of the overlaps. The important quantities to follow the dynamic of iterations and fixed points of AMP are the overlaps. Here, they are the P × P matrices

$$\begin{equation}\underline{\underline{q}}=\frac{1}{N}\sum _{i=1}^{N}{\hat{\underline{x}}}_{i}{\hat{\underline{x}}}_{i}^{\mathrm{T}},\quad \underline{\underline{m}}=\frac{1}{N}\sum _{i=1}^{N}{\hat{\underline{x}}}_{i}{{\underline{x}}_{0,i}}^{\mathrm{T}},\quad {\underline{\underline{q}}}_{0}=\frac{1}{N}\sum _{i=1}^{N}{\underline{x}}_{0,i}{{\underline{x}}_{0,i}}^{\mathrm{T}}.\end{equation} \tag{ F.32 }$$

Output parameters. Under independent statistics of the entries of $\underline{\underline{W}}$ and under the assumption of independent incoming messages, the variable ${\underline{\omega }}_{\mu \to i}$ defined in (F.6) is a sum of independent variables and follows a Gaussian distribution by the central limit theorem. Its first and second moments are

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{\omega }}_{\mu \to i}\right]=\frac{1}{\sqrt{N}}\sum _{{i}^{\prime }\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{W}_{\mu {i}^{\prime }}\right]{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }=0,\end{equation} \tag{ F.33 }$$

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{\omega }}_{\mu \to i}{\underline{\omega }}_{\mu \to i}^{\mathrm{T}}\right]=\frac{1}{N}\sum _{{i}^{\prime }\ne i}\sum _{{i}^{{\prime\prime}}\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{W}_{\mu {i}^{{\prime\prime}}}{W}_{\mu {i}^{\prime }}\right]{\hat{\underline{x}}}_{{i}^{{\prime\prime}}\to \mu }{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }^{\mathrm{T}}\end{equation} \tag{ F.34 }$$

$$\begin{align}\hfill & =\frac{1}{N}\sum _{{i}^{\prime }\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{W}_{\mu {i}^{\prime }}^{2}\right]{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }^{\mathrm{T}}=\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }^{\mathrm{T}}+O\left(1/N\right)\hfill \\ \hfill & =\frac{1}{N}\sum _{i=1}^{N}{\hat{\underline{x}}}_{{i}^{\prime }}{\hat{\underline{x}}}_{{i}^{\prime }}^{\mathrm{T}}-{\partial }_{\underline{\lambda }}{\underline{f}}_{1}^{x}{\underline{\underline{\sigma }}}_{i}{B}_{\mu \to i}{\hat{\underline{x}}}_{i}^{\mathrm{T}}-{\left({\partial }_{\underline{\lambda }}{\underline{f}}_{1}^{x}{\underline{\underline{\sigma }}}_{i}{B}_{\mu \to i}{\hat{\underline{x}}}_{i}^{\mathrm{T}}\right)}^{\mathrm{T}}+O\left(1/N\right)\hfill \\ \hfill & =\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\hat{\underline{x}}}_{{i}^{\prime }}{\hat{\underline{x}}}_{\prime i}^{\mathrm{T}}+O\left(1/\sqrt{N}\right)\hfill \end{align} \tag{ F.35 }$$

where we used the facts that the W_μi–s are independent with zero mean, and that B_μ→i, defined in (F.18), is of order $O\left(1/\sqrt{N}\right)$. Similarly, the variable ${\underline{z}}_{\mu \to i}={\sum }_{{i}^{\prime }\ne i}\frac{{W}_{\mu {i}^{\prime }}}{\sqrt{N}}{\underline{x}}_{{i}^{\prime }}$ is Gaussian with first and second moments

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{z}}_{\mu \to i}\right]=\frac{1}{\sqrt{N}}\sum _{{i}^{\prime }\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{W}_{\mu {i}^{\prime }}\right]{\underline{x}}_{0,{i}^{\prime }}=0,\end{equation} \tag{ F.36 }$$

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{z}}_{\mu \to i}{\underline{z}}_{\mu \to i}^{\mathrm{T}}\right]=\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\underline{x}}_{0,{i}^{\prime }}{{\underline{x}}_{0,{i}^{\prime }}}^{\mathrm{T}}+O\left(1/\sqrt{N}\right).\end{equation} \tag{ F.37 }$$

Furthermore, their covariance is

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{z}}_{\mu \to i}{\underline{\omega }}_{\mu \to i}^{\mathrm{T}}\right]=\frac{1}{N}\sum _{{i}^{\prime }\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{W}_{\mu {i}^{\prime }}^{2}\right]{\underline{x}}_{0,{i}^{\prime }}{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }^{\mathrm{T}}=\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\underline{x}}_{0,{i}^{\prime }}{\hat{\underline{x}}}_{{i}^{\prime }\to \mu }^{\mathrm{T}}+O\left(1/N\right)\end{equation} \tag{ F.38 }$$

$$\begin{equation}=\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\underline{x}}_{0,{i}^{\prime }}{\hat{\underline{x}}}_{{i}^{\prime }}^{\mathrm{T}}-{\underline{x}}_{0,i}{\partial }_{\underline{\lambda }}{\underline{f}}_{1}^{x}{\underline{\underline{\sigma }}}_{i}{\underline{B}}_{\mu \to i}^{\mathrm{T}}+O\left(1/N\right)\end{equation} \tag{ F.39 }$$

$$\begin{equation}=\frac{1}{N}\sum _{{i}^{\prime }=1}^{N}{\underline{x}}_{0,{i}^{\prime }}{\hat{\underline{x}}}_{{i}^{\prime }}^{\mathrm{T}}+O\left(1/\sqrt{N}\right).\end{equation} \tag{ F.40 }$$

Hence we find that for all μ–s and all i–s, ${\underline{\omega }}_{\mu \to i}$ and ${\underline{z}}_{\mu \to i}$ are approximately jointly Gaussian in the thermodynamic limit following a unique distribution $\mathcal{N}\left({\underline{z}}_{\mu \to i},{\underline{\omega }}_{\mu \to i};\;\underline{0},\;\underline{\underline{Q}}\right)$ with the block covariance matrix

$$\begin{equation}\underline{\underline{Q}}=\left[\begin{bmatrix}{cc}\hfill {\underline{\underline{q}}}_{0}\hfill & \hfill \underline{\underline{m}}\hfill \\ \hfill {\underline{\underline{m}}}^{\top }\hfill & \hfill \underline{\underline{q}}\hfill \end{bmatrix}\right].\end{equation} \tag{ F.41 }$$

For the variance message ${\underline{\underline{V}}}_{\mu \to i}$, defined in (F.7), we have

$$\begin{equation}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{\underline{V}}}_{\mu \to i}\right]=\sum _{{i}^{\prime }\ne i}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\frac{{W}_{\mu {i}^{\prime }}}{N}}^{2}\right]{\underline{\underline{{C}^{x}}}}_{{i}^{\prime }\to \mu }=\sum _{{i}^{\prime }=1}^{N}\frac{1}{N}{\underline{\underline{{C}^{x}}}}_{{i}^{\prime }\to \mu }+O\left(1/N\right)\end{equation} \tag{ F.42 }$$

$$\begin{equation}=\sum _{{i}^{\prime }=1}^{N}\frac{1}{N}{\underline{\underline{{C}^{x}}}}_{{i}^{\prime }}+O\left(1/\sqrt{N}\right),\end{equation} \tag{ F.43 }$$

where using the developments of ${\underline{\lambda }}_{i\to \mu }$ and ${\underline{\underline{\sigma }}}_{i\to \mu }$ (F.21) and (F.22), along with the scaling of ${\underline{B}}_{\mu \to i}$ in $O\left(1/\sqrt{N}\right)$ we replaced

$$\begin{equation}{\underline{\underline{{C}^{x}}}}_{i\to \mu }={\underline{\underline{f}}}_{2}^{x}\left({\underline{\lambda }}_{i\to \mu },{\underline{\underline{\sigma }}}_{i\to \mu }\right)={\underline{\underline{f}}}_{2}^{x}\left({\underline{\lambda }}_{i},{\underline{\underline{\sigma }}}_{i}\right)-{\partial }_{\underline{\lambda }}{\underline{\underline{f}}}_{2}^{x}{\underline{\underline{\sigma }}}_{i}{\underline{B}}_{\mu \to i}^{T}={\underline{\underline{f}}}_{2}^{x}\left({\underline{\lambda }}_{i},{\underline{\underline{\sigma }}}_{i}\right)+O\left(1/\sqrt{N}\right).\end{equation} \tag{ F.44 }$$

Futhermore, we can check that

$$\begin{equation}\underset{N\to +\infty }{\mathrm{lim}}{\mathbb{E}}_{\underline{\underline{W}}}\left[{\underline{\underline{V}}}_{\mu \to i}^{2}-{\mathbb{E}}_{\underline{\underline{W}}}{\left[{\underline{\underline{V}}}_{\mu \to i}\right]}^{2}\right]=0,\end{equation} \tag{ F.45 }$$

meaning that all ${\underline{\underline{V}}}_{\mu \to i}$ concentrate on their identical mean in the thermodynamic limit, which we note

$$\begin{equation}\underline{\underline{V}}=\sum _{i=1}^{N}\frac{1}{N}{\underline{\underline{{C}^{x}}}}_{i}.\end{equation} \tag{ F.46 }$$

Input parameters. Here we use the re-parametrization trick to express ${\underline{y}}_{\mu }$ as a function g₀(⋅) taking a noise ${\underline{{\epsilon}}}_{\mu }\sim {p}_{{\epsilon}}\left({\underline{{\epsilon}}}_{\mu }\right)$ as inputs: ${\underline{y}}_{\mu }={g}_{0}\left({\underline{w}}_{\mu }^{\top }{\underline{\underline{X}}}_{0},{\underline{{\epsilon}}}_{\mu }\right)$. Following (F.19), (F.18) and (F.28),

$$\begin{equation}{\underline{\underline{\sigma }}}_{i}^{-1}{\underline{\lambda }}_{i}=\sum _{\mu =1}^{M}\frac{{W}_{\mu i}}{\sqrt{N}}\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({\underline{y}}_{\mu },{\underline{\omega }}_{\mu \to i},{\underline{\underline{V}}}_{\mu \to i}\right)\end{equation} \tag{ F.47 }$$

$$\begin{equation}=\sum _{\mu =1}^{M}\frac{{W}_{\mu i}}{\sqrt{N}}\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({g}_{0}\left(\sum _{{i}^{\prime }\ne i}\frac{{W}_{\mu {i}^{\prime }}}{\sqrt{N}}{\underline{x}}_{0,{i}^{\prime }}+\frac{{W}_{\mu i}}{\sqrt{N}}{\underline{x}}_{0,i},{\underline{{\epsilon}}}_{\mu }\right),{\underline{\omega }}_{\mu \to i},{\underline{\underline{V}}}_{\mu \to i}\right)\end{equation} \tag{ F.48 }$$

$$\begin{align}\hfill & =\sum _{\mu =1}^{M}\frac{{W}_{\mu i}}{\sqrt{N}}\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({g}_{0}\left(\sum _{{i}^{\prime }\ne i}\frac{{W}_{\mu {i}^{\prime }}}{\sqrt{N}}{\underline{x}}_{0,{i}^{\prime }},{\underline{{\epsilon}}}_{\mu }\right),{\underline{\omega }}_{\mu \to i},{\underline{\underline{V}}}_{\mu \to i}\right)\hfill \\ \hfill & +\sum _{\mu =1}^{M}\frac{{W}_{\mu i}^{2}}{N}{\partial }_{z}\underline{\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}}\left({g}_{0}\left({\underline{z}}_{\mu \to i},{\underline{{\epsilon}}}_{\mu }\right),{\underline{\omega }}_{\mu \to i},{\underline{\underline{V}}}_{\mu \to i}\right){\underline{x}}_{0,i}.\hfill \end{align} \tag{ F.49 }$$

The first term is again a sum of independent random variables, given the W_μi are i.i.d. with zero mean, of which the messages of type μ → i are assumed independent. The second term has non-zero mean and can be shown to concentrate. Finally recalling that all ${\underline{\underline{V}}}_{\mu \to i}$ also concentrate on $\underline{\underline{V}}$ we obtain the distribution

$$\begin{equation}{\underline{\underline{\sigma }}}_{i}^{-1}{\underline{\lambda }}_{i}\sim \mathcal{N}\left({\underline{\underline{\sigma }}}_{i}^{-1}{\underline{\lambda }}_{i};\;\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0,i},\sqrt{\alpha \underline{\underline{\hat{q}}}}{I}_{P}\right)\end{equation} \tag{ F.50 }$$

with

$$\begin{equation}\underline{\underline{\hat{q}}}=\int \mathrm{d}\underline{{\epsilon}}\;{p}_{{\epsilon}}\left(\underline{{\epsilon}}\right)\mathrm{d}{s}_{0}\;{p}_{{s}_{0}}\left({s}_{0}\right)\int \mathrm{d}\underline{\omega }\;\mathrm{d}\underline{z}\;\mathcal{N}\left(\underline{z},\underline{\omega };\underline{0},\underline{\underline{Q}}\right)\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },\underline{\underline{V}}\right)\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}{\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },\underline{\underline{V}}\right)}^{\mathrm{T}},\end{equation} \tag{ F.51 }$$

$$\begin{equation}\underline{\underline{\hat{m}}}=\int \mathrm{d}\underline{{\epsilon}}\;{p}_{{\epsilon}}\left(\underline{{\epsilon}}\right)\mathrm{d}{s}_{0}\;{p}_{{s}_{0}}\left({s}_{0}\right)\int \mathrm{d}\underline{\omega }\;\mathrm{d}\underline{z}\;\mathcal{N}\left(\underline{z},\underline{\omega };\underline{0},\underline{\underline{Q}}\right){\partial }_{\underline{z}}\underline{\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}}\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },\underline{\underline{V}}\right).\end{equation} \tag{ F.52 }$$

For the inverse variance ${\underline{\underline{\sigma }}}_{i}^{-1}$ one can check again that it concentrates on its mean

$$\begin{equation}{\underline{\underline{\sigma }}}_{i}^{-1}=\sum _{\mu =1}^{M}\frac{{{W}_{\mu i}}^{2}}{N}{\partial }_{\underline{\omega }}\underline{\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}}\left({\underline{y}}_{\mu },{\underline{\omega }}_{\mu \to i},{\underline{\underline{V}}}_{\mu \to i}\right)\simeq \alpha \underline{\underline{\hat{\chi }}},\end{equation} \tag{ F.53 }$$

$$\begin{equation}\underline{\underline{\hat{\chi }}}=-\int \mathrm{d}\underline{{\epsilon}}\;{p}_{{\epsilon}}\left(\underline{{\epsilon}}\right)\mathrm{d}{s}_{0}\;{p}_{{s}_{0}}\left({s}_{0}\right)\int \mathrm{d}\underline{{\epsilon}}\;\mathrm{d}\underline{z}\;\mathcal{N}\left(\underline{z},\underline{\omega };\underline{0},\underline{\underline{Q}}\right){\partial }_{\omega }\underline{\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}}\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },\underline{\underline{V}}\right).\end{equation} \tag{ F.54 }$$

Closing the equations. These statistics of the input parameters must ensure that consistently

$$\begin{equation}\underline{\underline{V}}=\frac{1}{N}\sum _{i=1}^{N}{\underline{\underline{{C}^{x}}}}_{i}={\mathbb{E}}_{\underline{\lambda },\underline{\underline{\sigma }}}\left[{\underline{\underline{f}}}_{2}^{x}\left(\underline{\lambda },\underline{\underline{\sigma }}\right)\right],\end{equation} \tag{ F.55 }$$

$$\begin{equation}\underline{\underline{q}}=\frac{1}{N}\sum _{i=1}^{N}{\hat{\underline{x}}}_{i}{\hat{\underline{x}}}_{i}^{\top }={\mathbb{E}}_{\underline{\lambda },\underline{\underline{\sigma }}}\left[{\underline{\underline{f}}}_{1}^{x}\left(\underline{\lambda },\underline{\underline{\sigma }}\right){\underline{\underline{f}}}_{1}^{x}{\left(\underline{\lambda },\underline{\underline{\sigma }}\right)}^{\top }\right],\end{equation} \tag{ F.56 }$$

$$\begin{equation}\underline{\underline{m}}=\frac{1}{N}\sum _{i=1}^{N}{\hat{\underline{x}}}_{i}{{\underline{x}}_{0,i}}^{\top }={\mathbb{E}}_{\underline{\lambda },\underline{\underline{\sigma }}}\left[{\underline{\underline{f}}}_{1}^{x}\left(\underline{\lambda },\underline{\underline{\sigma }}\right){{\underline{x}}_{0,i}}^{\top }\right],\end{equation} \tag{ F.57 }$$

which gives upon expressing the computation of the expectations

$$\begin{equation}\underline{\underline{V}}=\int \mathrm{d}{\underline{x}}_{0}\;{p}_{{x}_{0}}\left({\underline{x}}_{0}\right)\int \mathcal{D}\underline{\xi }\;{\underline{\underline{f}}}_{2}^{x}\left({\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\left(\sqrt{\alpha \underline{\underline{\hat{q}}}}\underline{\xi }+\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0}\right);{\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\right),\end{equation} \tag{ F.58 }$$

$$\begin{equation}\underline{\underline{m}}=\int \mathrm{d}{\underline{x}}_{0}\;{p}_{{x}_{0}}\left({\underline{x}}_{0}\right)\int \mathcal{D}\underline{\xi }\;{\underline{f}}_{1}^{x}\left({\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\left(\sqrt{\alpha \underline{\underline{\hat{q}}}}\underline{\xi }+\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0}\right);{\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\right){{\underline{x}}_{0}}^{\top },\end{equation} \tag{ F.59 }$$

$$\begin{align}\hfill \underline{\underline{q}}& =\int \mathrm{d}{\underline{x}}_{0}\;{p}_{{x}_{0}}\left({\underline{x}}_{0}\right)\int \mathcal{D}\underline{\xi }\;{\underline{f}}_{1}^{x}\left({\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\left(\sqrt{\alpha \underline{\underline{\hat{q}}}}\underline{\xi }+\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0}\right);{\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\right)\hfill \\ \hfill & {\times}{\underline{f}}_{1}^{x}{\left({\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\left(\sqrt{\alpha \underline{\underline{\hat{q}}}}\underline{\xi }+\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0}\right);{\left(\alpha \underline{\underline{\hat{\chi }}}\right)}^{-1}\right)}^{\top }.\hfill \end{align} \tag{ F.60 }$$

The state evolution analysis of the GLM on the vector variables finally consists in iterating alternatively the equations (F.50), (F.51), (F.53), and the equations (F.57), (F.58) (F.59) until convergence.

Performance analysis. The mean squared error (MSE) on the reconstruction of $\underline{\underline{X}}$ by the AMP algorithm is then predicted by

$$\begin{equation}\mathrm{M}\mathrm{S}\mathrm{E}\left(\underline{\underline{X}}\right)=q-2m+{q}_{0},\end{equation} \tag{ F.61 }$$

where the scalar values used here correspond to the (unique) value of the diagonal elements of the corresponding overlap matrices. This MSE can be computed throughout the iterations of state evolution. Remarkably, the state evolution MSEs follow precisely the MSE of the cal-AMP predictors along the iterations of the algorithm provided the procedures are initialized consistently. A random initialization of ${\hat{\underline{x}}}_{i}$ in cal-AMP corresponds to an initialization of zero overlap m = 0, ν = 0, with variance of the priors q = q₀ in the state evolution.

The SE equations can be greatly simplified in the Bayes optimal setting where the statistical model used by the student (priors p_x and p_s, and channel p_out) is known to match the teacher. In this case, the true unknown signal ${\underline{\underline{X}}}_{0}$ is in some sense statistically equivalent to the estimate $\underline{\underline{\hat{X}}}$ coming from the posterior. More precisely one can prove the Nishimori identities [Iba99, Nis01, OH91] (or [KKM⁺16] for a concise demonstration and discussion) implying that

$$\begin{equation}\underline{\underline{q}}=\underline{\underline{m}},\quad \underline{\underline{V}}={\underline{\underline{q}}}_{0}-\underline{\underline{m}},\quad \underline{\underline{\hat{q}}}=\underline{\underline{\hat{m}}}=\underline{\underline{\hat{\chi }}}\quad \text{and}\quad r=\nu .\end{equation} \tag{ F.62 }$$

As a result the state evolution reduces to a set of two equations

$$\begin{equation}\underline{\underline{m}}=\int \mathrm{d}{\underline{x}}_{0}\;{p}_{{x}_{0}}\left({\underline{x}}_{0}\right)\int \mathcal{D}\underline{\xi }\;{\underline{f}}_{1}^{x}\left({\left(\alpha \underline{\underline{\hat{m}}}\right)}^{-1}\left(\sqrt{\alpha \underline{\underline{\hat{m}}}}\underline{\xi }+\alpha \underline{\underline{\hat{m}}}{\underline{x}}_{0}\right);{\left(\alpha \underline{\underline{\hat{m}}}\right)}^{-1}\right){{\underline{x}}_{0}}^{\top },\end{equation} \tag{ F.63 }$$

$$\begin{align}\hfill \underline{\underline{\hat{m}}}& =\int \mathrm{d}\underline{{\epsilon}}\;{p}_{{\epsilon}}\left(\underline{{\epsilon}}\right)\mathrm{d}{s}_{0}\;{p}_{{s}_{0}}\left({s}_{0}\right)\int \mathrm{d}\underline{\omega }\;\mathrm{d}\underline{z}\;\mathcal{N}\left(\underline{z},\underline{\omega };\underline{0},\underline{\underline{Q}}\right)\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },{\underline{\underline{q}}}_{0}-\underline{m}\left.\right)\right)\hfill \\ \hfill & {\times}\underline{{g}_{\mathrm{o}\mathrm{u}\mathrm{t}}}\left({g}_{0}\left(\underline{z},\underline{{\epsilon}}\right),\underline{\omega },{\underline{\underline{q}}}_{0}-\underline{m}\left.\right)\right),\hfill \end{align} \tag{ F.64 }$$

with the block covariance matrix

$$\begin{equation}\underline{\underline{Q}}=\left[\begin{bmatrix}{cc}\hfill {\underline{\underline{q}}}_{0}\hfill & \hfill \underline{\underline{m}}\hfill \\ \hfill {\underline{\underline{m}}}^{\top }\hfill & \hfill \underline{\underline{m}}\hfill \end{bmatrix}\right].\end{equation} \tag{ F.65 }$$