Zilu Zhao
email: zilu.zhao@eurecom.fr

Dirk Slock
email: dirk.slock@eurecom.fr

IMPROVED VARIANCE PREDICTIONS IN APPROXIMATE MESSAGE PASSING

Keywords: Approximate Message Passing, AMP, VAMP, reVAMP, AMB UAMP, Haar Large System Analysis

In the Generalized Linear Model (GLM), the unknowns may be non-identically independent distributed (niid), as for instance in the Sparse Bayesian Learning (SBL) problem. The Generalized Approximate Message Passing (GAMP) algorithm performs computationally efficient belief propagation for Bayesian inference. The GAMP algorithms predicts the posterior variances correctly in the case of measurement matrices with (n)iid entries. In order to cover more illconditioned measurement matrices, the (right) rotationally invariant (RRI) model was introduced in which the (right) singular vectors are Haar distributed, leading to Vector AMP VAMP however assumes iid priors and posteriors. Here we introduce a convergent version of AMB (AMBAMP) applied to Unitarily transformed data, with a variance correction based on Haar Large System Analysis (LSA). The recently introduced reVAMP perspective shows that the resulting AM-BUAMP algorithm has an underlying multivariate Gaussian posterior approximation, that does not get computed but that allows the LSA. The individual variance predictions are exact asymptotically in the RRI setting, as illustrated by a Gaussian Mixture Model example.

INTRODUCTION

The recovery of sparse signal vectors is a fundamental problem in signal processing and has wide-ranging applications, including compressed sensing, image and speech processing, and machine learning. Sparse Bayesian Learning (SBL) [START_REF] Michael | Sparse Bayesian Learning and the Relevance Vector Machine[END_REF] is one popular method for sparse signal recovery. In spite of the a priori non sparsifying Gaussian prior, when the prior variances need to be estimated also, this hierarchical Gaussian setting becomes sparsifying. In the Gaussian noise case, the signal model for the recovery of a sparse signal vector x can be formulated as, y = Ax + v, where y are the observations or data, A is called the known measurement or sensing matrix of dimension M × N with M < N in the compressed sens-EURECOM's research is partially supported by its industrial members: ORANGE, BMW, SAP, iABG, Norton LifeLock, and by the FrancoGerman ANR-DFG project CellFree6G. ing case. For a sparse model, x contains only K non-zero (or significant) entries, with K < M < N . However, SBL can be computationally complex, especially when dealing with high-dimensional data. This complexity arises from the need to perform matrix inversion in each iteration (for the hyperparameter estimation). To overcome this challenge, approximate inference methods have been developed, with Approximate Message Passing (AMP) being a popular and efficient approach [START_REF] Zou | A Concise Tutorial on Approximate Message Passing[END_REF]. AMP has been shown to be effective in recovering high-dimensional signals, and its dynamics can be fully characterized by state evolution [START_REF] Bayati | The Dynamics of Message Passing on Dense Graphs, with Applications to Compressed Sensing[END_REF]. However, the convergence of AMP can be problematic when dealing with ill-conditioned measurement matrices A.

To address this issue, the Vector Approximate Message Passing (VAMP) algorithm has been proposed [START_REF] Schniter | Vector Approximate Message Passing for the Generalized Linear Model[END_REF]. It generates the factor graph by splitting one vector variable node x into two variable instances x 1 = x 2 = x. An Expectation-Propagation (EP) like message passing algorithm [START_REF] Minka | Expectation Propagation for Approximate Bayesian Inference[END_REF] is then applied to the factor graph with vector valued messages. VAMP has been shown to perform well under Right Rotationally Invariant (RRI) A, and the (variance) state evolution of VAMP has been rigorously established [START_REF] Schniter | Vector Approximate Message Passing for the Generalized Linear Model[END_REF].

Prior Work

In previous research [START_REF] Schniter | Vector Approximate Message Passing for the Generalized Linear Model[END_REF], the optimality of VAMP has been analyzed using the replica method [START_REF] Tulino | Support Recovery with Sparsely Sampled Free Random Matrices[END_REF]. The replica method gives a system of equations that describes the fixed point of the VAMP state evolution and the optimal (sum) Mean Squared Error (MSE). VAMP (implicitly) assumes identically independently distributed (iid) priors and posteriors. AMP provides individual MSEs, but with convergence issues. Unitary Transformed AMP (UTAMP) [START_REF] Guo | Approximate Message Passing with Unitary Transformation[END_REF] applies AMP to a transformed linear model to make AMP more robust to A.

In [START_REF] Kurisummoottil | Towards Convergent Approximate Message Passing by Alternating Constrained Minimization of Bethe Free Energy[END_REF] we have introduced a convergent version of AMP, AM-BGAMP, which works with arbitrary priors to accommodate e.g. LASSO and many other compressed sensing problem formulations. The variance predictions in AMP are correct in the Large System Limit (LSL) for a matrix A considered random with non iid (niid) elements (i.e. independent but not necessarily identically distributed, e.g. an A with deterministic magnitudes but iid signs). In [START_REF] Zhao | Variance Predictions in VAMP/UAMP with Right Rotationally Invariant Measurement Matrices for niid Generalized Linear Models[END_REF] we have introduced UAMP which assumes a RRI model for A. Such a model covers a much wider class of matrices since only the right singular vectors are considered Haar distributed [START_REF] Couillet | Random Beamforming Over Quasi-Static and Fading Channels: A Deterministic Equivalent Approach[END_REF] (the row subspace of A is uniformly oriented in space), whereas left singular vectors and singular values are deterministic. The RRI model requires a variance correction to be correct in the LSL, based on Haar Large System Analysis (LSA). However, UAMP in [START_REF] Zhao | Variance Predictions in VAMP/UAMP with Right Rotationally Invariant Measurement Matrices for niid Generalized Linear Models[END_REF] is limited to Gaussian priors.

Main Contribution

In this paper, we extend UAMP to handle arbitrary priors, leading to the Generalized Linear Model (GLM) (here still with Gaussian noise though). We introduce the Haar LSA based variance corrections to the convergent AMBAMP applied to a unitarily tranformed data model, resulting in the AMBUAMP algorithm. Since the LSA is based on Gaussian signal models, we invoke the recently introduced reVAMP perspective [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF] to show that these AMP algorithms with arbitrary priors implicitly construct a multivariate Gaussian posterior approximation, on which the LSA can be based, and which is consistent with the marginal posterior first and second-order moments constructed by the algorithm. We show that without correction term the optimal MSE is not reached, but the Haar LSA allows to distill the proper correction term which only depends on quantities already appearing in the algorithm. To illustrate the proper functioning of the AMBUAMP algorithm, we work out the details for and simulate the case of Gaussian Mixture Model (GMM) priors.

UNITARILY TRANSFORMED LINEAR MODEL

The data model considered in AMP is essentially a linear mixing model

y = A x + v , p x (x) , p v (v) (1)
with (possibly) non identically independently distributed (niid) prior

p x (x) = N i=1 p xi (x i ) and iid measurement noise p v (v) = N i=1 p v (v i ), v ∼ N (0, σ 2 v I).
A is the M × N measurement matrix and we assume M < N . In order to handle possibly ill-conditioned A , consider the economy singular value decomposition (SVD) of A :

A = UΣ V T (2) where U ∈ R M ×M is an orthogonal matrix, Σ ∈ R M ×M is diagonal and V ∈ R M ×N is semi-orthogonal (V T V = I).
As in the computationally efficient version of VAMP [START_REF] Guo | Approximate Message Passing with Unitary Transformation[END_REF], [START_REF] Schniter | Vector Approximate Message Passing for the Generalized Linear Model[END_REF], we can take advantage of the simple white Gaussian noise model to transform the linear model as:

U T y = Σ V T x + U T v .
(3) Since we assume v to be iid Gaussian, we still have U T v ∼ N (0, σ 2 v I). For simplicity, let's denote

y ′ = U T y, v ′ = U T v, σ 2 v = σ 2 v 1, A ′ = ΣV T , S ′ = A ′ .A ′ and λ = Σ 2 1, leading to the transformed linear model y ′ = A ′ x + v ′ (4
) where z = A ′ x is the result of the linear mixture.

PROPOSED AMBUAMP

The abbreviation AMB stands for ACM-LSL-BFE, which stands for Alternating Constrained Minimization of the LSL of the Bethe Free Energy (BFE). AMBGAMP employs most of the same updates as GAMP, but GAMP does not apply a strict alternating minimization (block coordinate descent) principle, particularly in the presence of constraints. AM-BGAMP has been derived in [START_REF] Kurisummoottil | Towards Convergent Approximate Message Passing by Alternating Constrained Minimization of Bethe Free Energy[END_REF] of which we recall the key steps. We shall apply here AMBGAMP to the unitarily transformed data (4), which will lead to AMBUAMP after a correction of the variances. GAMP corresponds to the constrained minimization of a LSL of the BFE (see also [START_REF] Rangan | Fixed Points of Generalized Approximate Message Passing with Arbitrary Matrices[END_REF] and references therein): min qx,qz,τp

J BF E (q x , q z , τ p ) s.t. E(z|q z ) = A ′ E(x|q x ) τ p = S ′ var(x|q x ), (5)
where the LSL BFE is given by

J BF E (q x , q z , τ p ) = D(q x ||e -fx )+D(q z ||e -fz )+H G (q z ,τ p ) with H G (q z , τ p ) = 1 2 M k=1 var(z k |q z k ) τ p k +ln(2πτ p k ) (6)
and where D(q||p) = E q (ln( q p )) is the Kullback-Leibler Divergence (KLD) and H G (q z , τ p ) is a sum of a KLD and an entropy of Gaussians with identical means but different variances. The LSL BFE optimization problem ( 6) can be reformulated with the following augmented Lagrangian min qx,qz,τp,u max s,τs L(q x , q z , τ p , u, s, τ s ) with

L = D(q x ||e -fx )+D(q z ||e -fz )+H G (q z , τ p ) +s T (E(z|q z ) -A ′ E(x|q x )) -1 2 τ T s (τ p -S ′ var(x|q x )) + 1 2 ∥ E(x|q x ) -u∥ 2 τr + 1 2 ∥ E(z|q z ) -A ′ u∥ 2 τp , (7)
where s, τ s are Lagrange multipliers, and τ r = 1./(S ′T τ s ) is just a short-hand notation for a quantity that depends on τ s . We also use the notations: ∥u∥ 2 τ = i u 2 i /τ i , elementwise multiplication as in s.τ and element-wise division as in 1./τ , and 1 is a vector of ones. In [START_REF] Kurisummoottil | Towards Convergent Approximate Message Passing by Alternating Constrained Minimization of Bethe Free Energy[END_REF] we apply an alternating optimization strategy, combined with a gradient update with line search for the auxiliary quantity u, an ADMM update of the Lagrange multipliers s, and fixed point iterations for τ p and τ s . We detail two key steps.

Update of u

To update u, we use a gradient descent method with line search optimized step-size. From [START_REF] Guo | Approximate Message Passing with Unitary Transformation[END_REF] we get at iteration t L(q t-1

x , q t-1 z , τ t-1 p , u, s t-1 , τ t-1 s ) = 1 2 ∥ x t-1 -u∥ 2 τ t-1 r + 1 2 ∥ z t-1 -A ′ u∥ 2 τ t-1 p + const. (8
) where const. denotes constants w.r.t. u. The minimizing update can be obtained as

u t = u t-1 -η t g t (9)
with gradient g t = g t (u t-1 ) where

g t (u) = ∇ u L(q t-1 x , q t-1 z , τ t-1 p , u, s t-1 , τ t-1 s ) = -A ′T (( z t-1 -A ′ u)./τ t-1 p ) -( x t-1 -u)./τ t-1 r = g t (0) + H t u, H t = D(1./τ t-1 r ) + A ′T D(1./τ t-1 p )A ′ (10)
where D(τ ) denotes a diagonal matrix with diagonal elements τ . The step-size η t gets optimized for maximum descent :

∂L(q t-1 x , q t-1 z , τ t-1 p , u t , s t-1 , τ t-1 s ) ∂η t = 0 ⇒ η t = ∥g t ∥ 2 /g t T H t g t . (11)

Update of q x

For the update of q x , consider the relevant terms in the augmented Lagrangian (and remember that 1.

/τ t-1 r = S ′T τ t-1 s ) L(q x , q t-1 z , τ t-1 p , u t , s t-1 , τ t-1 s ) = D(q x ||e -fx ) -s t-1 T A ′ E(x|q x ) + 1 2 τ t-1 T s S ′ var(x|q x ) + 1 2 ∥ E(x|q x ) -u t ∥ 2 τ t-1 r + const. = D(q x ||e -fx ) + 1 2 (1./τ t-1 r ) T E(x.x|q x ) -s t-1 T A ′ E(x|q x ) -(u t ./τ t-1 r )) T E(x|q x ) + const. = D(q x ||e -fx ) + 1 2 (1./τ t-1 r ) T E(x.x|q x ) -(u t + τ t-1 r .A ′T s t-1 ) T (E(x|q x )./τ t-1 r ) + const. = D(q x ||e -fx ) + 1 2 E(∥x -r t ∥ 2 τ t r |q x ) + const. (12)
where const. denotes constants w.r.t. x, and r t = u t + τ t-1 r .A ′T s t-1 . The Lagrangian in [START_REF] Rangan | Fixed Points of Generalized Approximate Message Passing with Arbitrary Matrices[END_REF] is separable. We get per component min

qx i D(q xi ||g t xi /Z t xi ) ⇒ q t xi = g t xi /Z t xi , Z t xi = g t xi (x i ) dx i , -ln g t xi (x i ) = f xi (x i ) + 1 2τ t r i [(x k -r t i ) 2 -r t 2 i ] . (13)
The partition function Z t xi acts as cumulant generating function:

τ t ri ∂ ln Z t xi ∂r t i = E(x i |q t xi ) = E(x i |r t i , τ t ri ) = x t i (τ t ri ) 2 ∂ 2 ln Z t xi ∂r t 2 i = var(x i |r t i , τ t ri ) = τ t xi . (14)
In the Gaussian prior case, we get a Gaussian posterior

q t x with 1./τ t x = 1./τ t-1 r +1./σ 2 x , x t = τ t x .(r t ./τ t-1 r
) .

(15) The AMBUAMP algorithm appears in the Algorithm 1 table. The variance correction term appears in red, where the introduction of the nonnegative part ⌊.⌋ + may avoid transient problems.

REVAMP INTERPRETATION

In [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF], we introduced the reVAMP algorithm, which corresponds to a revisited VAMP approach to MMSE estimation in the GLM where the only asymptotic approximation that is made is to invoke the Central Limit Theorem (CLT) to approximate the extrinsic pdfs by Gaussians. A MMSE estimate corresponds to the mean of the posterior pdf, and decouples

Algorithm 1 AMBUAMP Require: y ′ , A ′ , S ′ = A ′ .A ′ , fx(x), σ 2 v 1: Initialize: t = 0, u 0 = 0, x 0 = 0, z 0 = 0, s 0 = 0, τ 0 r = 1, τ 0 p = 1
2: repeat (t=1,2,. . . ) 3:

u t = u t-1 -η t g t , with g t , η t from (10), [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF] 4:

[Input node update] 5:

r t = u t + τ t-1 r .(A ′T s t-1 )
6:

x t = E(x|r t , τ t-1 r ), Gaussian px : 1./τ t x = 1./τ t-1 r + 1./σ 2 x 7: τ t x =var(x|r t , τ t-1 r
), Gaussian px : x t = τ t x .(r t ./τ t-1 r ) 8:

τ t p = S ′ τ t x 9:
[Output node update] 10:

p t = A ′ u t -s t-1 .τ t p 11:
1./τ t z = 1./τ t p + 1./σ 2 v 12:

z t = τ t z .(y ′ ./σ 2 v +p t ./τ t p )
13:

s t = s t-1 + ( z t -A ′ u t
)./τ t p 14:

τ t s = 1./(τ t p + 1 -1 N 2 1 T τ t x λ T τ t-1 s + σ 2 v )
15:

τ t r = 1./(S ′T τ t s )
16: until Convergence between the components of a vector. So, consider the joint pdf

p(x, y ′ ) = p(y ′ |x) N i=1 p xi (x i ) . ( 16
)
The true posterior for x i can be written as

p(x i |y ′ ) = p xi (x i ) p(y ′ |x) N j̸ =i p xj (x j )dx j /Z i (y ′ ) ≈ p xi (x i ) N (x i ; r i , τ ri ) (17) where p(y ′ |x) N j̸ =i p xj (x j )dx j /Z i (y ′
) is the extrinsic for x i and Z i (y ′ ) is a normalization factor. For a very large class of models for A ′ and x, it is clear that the CLT will allow to approximate the extrinsic by a Gaussian distribution, as indicated in (17). As explored in [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF], the Gaussian extrinsic approximation can be built from the Component-Wise Conditionally Unbiased MMSE estimate of x i . Now, from the non-Gaussian approximate posterior in (17), with Gaussian extrinsic but true prior, we can find the MMSE estimate x i and associated MMSE τ xi . These quantities allow us to build a Gaussian posterior approximation q xi (x i ) = N (x i ; x i , τ xi ) which minimizes the KLD to the non-Gaussian posterior. The (marginal) Gaussian posterior approximation q xi (x i ) in turn allows us to find a Gaussian prior approximation q i (x i ) by Gaussian division

q i (x i ) = N x i ; a i , σ 2 xi ∝ N (x i ; x i , τ xi ) /N (x i ; r i , τ ri ) , 1/σ 2 xi = 1/τ xi -1/τ ri , a i = σ 2 xi ( x i /τ xi -r i /τ ri ) .
(18) Note that in the Gaussian prior case, the σ 2

x are of course the variances of the true prior. The joint distribution N i=1 q i (x i ) equals N (x; a, σ 2

x ). Now, since also p(y|x) is Gaussian, the (independent) Gaussian priors in turn induce a multivariate Gaussian posterior approximation

q x (x) = N (x; m, C m ) ∝ p(y|x)N (x; a, σ 2 x ) (19) where C m = A ′T C -1 vv A ′ + D -1 x -1 m = C m A ′T C -1 vv y ′ + D -1 x a , (20)
D x is a short hand notation for D(σ 2 x ) and C vv = σ 2 v I. Furthermore, if we define τ m = diag(C m ), then we have at convergence of the reVAMP algorithm m = x and τ m = τ x , and q xi (x i ) is a marginal of q x (x). This reVAMP perspective on the (AMB)(U)AMP algorithm shows that asymptotically, the MMSE estimates can be associated to an underlying Gaussian linear model which will facilitate the large random matrix analysis for the variances.

HAAR LARGE SYSTEM ANALYSIS

Following [START_REF] Schniter | Vector Approximate Message Passing for the Generalized Linear Model[END_REF], we model A as a right rotationally invariant matrix. That means that in the economy SVD of A = UΣV T , U and Σ remain deterministic but V is considered drawn as M columns of a Haar distributed N × N random matrix (i.e. the columns of V form an orthonormal basis for a uniformly randomly oriented M -dimensional subspace of the N -dimensional space). The analysis of the large system primarily relies on the deterministic equivalent proposed in [START_REF] Dirk | Nonlinear MMSE using Linear MMSE Bricks and Application to Compressed Sensing and Adaptive Kalman Filtering[END_REF], which states Lemma 1. Let P be any Hermitian matrix with bounded spectral norm and let V ∈ R N ×M be M < N columns of a Haar distributed (unitary) random matrix. Let B be a nonnegative definite matrix with ||B|| < ∞ (||B|| represents the spectral norm) and D be any diagonal matrix with positive entries. Then the following convergence result holds almost surely,

1 N tr B VP V T + D -1 - 1 N tr B(e I + D) -1 a.s.
--→ 0. (

) 22
The MMSE solution for (1) is given by

x MMSE = 1 σ 2 v A T A + D 1./σ 2 x -1 1 σ 2 v A T y, C MMSE = 1 σ 2 v A T A + D 1./σ 2 x -1 . (23)
Thus, the MSE is

MSE = 1 N tr[C MMSE ]. (24)
By using Lemma 1, we obtain the large system approximation MSE a.s.

--→

1 N tr (e 0 I + D(1./σ 2 x )) -1 (25) with e 0 = 1 N tr 1 σ 2 v Σ 2 e 0 σ 2 v Σ 2 + (1 -e 0 e 0 )I -1 , e 0 = 1 N tr (e 0 I + D(1./σ 2 x )) -1 .
(26)

LARGE SYSTEM ANALYSIS OF AMBUAMP

We will first prove that under the assumption that V is Haardistributed, the averaged τ x , namely 1 N 1 T τ x , does not match the optimal MSE defined in (24). Then in the next section, we will propose a correction term such that τ x matches the optimal MSE. The steady state of variances in UAMP can be summarized as follows

1./τ ∞ s = σ 2 v + S ′ τ ∞ x 1./τ ∞ x = 1./σ 2 x + S ′T τ ∞ s . (27)
With the large system assumptions, as N tend to infinity, we approximate V T D N V and VD M V T to 1 N tr(D N )I and 1 N tr(D M )I respectively. Thus, we have

S ′ τ ∞ x = diag ΣV T D(τ ∞ x )VΣ = 1 N 1 T τ ∞ x λ, S ′T τ ∞ s = diag VΣD(τ ∞ s )ΣV T = 1 N λ T τ ∞ s 1. (28)
Now we show the following.

Lemma 2. In AMP with equivalent measurement matrix A ′ , the variance prediction 1 N 1 T τ ∞ x does not match the optimal MSE in (25).

Proof. If the noise is iid, the MSE remains unchanged under a unitary transformation. Therefore, equations ( 25) and ( 26) remain the same in this transformed system. We will prove by contradiction. Suppose that τ ∞

x matches the optimal MSE, we then have

1 N tr[D(τ ∞ x )] = 1 N tr D(S ′T τ ∞ s ) + D(1./σ 2 x ) -1 = e 0 = 1 N tr (e 0 I + D(1./σ 2 x )) -1 , (29)
which implies

e 0 = 1 N λ T τ ∞ s = 1 N tr Σ 2 D(τ ∞ s ) = 1 N tr Σ 2 1 N tr[D(τ ∞ x )]D(λ) + σ 2 v I -1 = 1 N tr Σ 2 e 0 Σ 2 + σ 2 v I -1 (30)
One can observe that e 0 in (30) only equals e 0 in (26) if e 0 e 0 = 0.

Correction Term For τ s

Considering the asymptotic MSE expression in (25), and the second equations in ( 27), (28), we can still write

1 N tr[D(τ ∞ x )] = 1 N tr D(S ′T τ ∞ s ) + D(1./σ 2 x ) -1 = 1 N tr 1 N λ T τ ∞ s I + D(1./σ 2 x ) -1 (31) Introduce e c = 1 N tr[D(τ ∞ x )] , e c = 1 N λ T τ ∞ s . (32)
Now compare (31),(32) with (25),(26), then we require e c to be of the form

e c = 1 N tr Σ 2 e c Σ 2 + (1 -e c e c )σ 2 v I -1 . (33)
From the definition of e c in (32), we have

e c = 1 N λ T τ ∞ s = 1 N tr Σ 2 D(τ ∞ s ) . (34)
Comparing (34) with (33), we want to design the update scheme of τ s such that at steady state,

D(τ ∞ s ) = e c Σ 2 + (1 -e c e c )σ 2 v I -1 . (35)
From the definition of e c in (32), we have

e c Σ 2 = 1 N tr[D(τ x )]Σ 2 = 1 N 1 T τ x Σ 2 = 1 N 1 T τ x D(λ).
(36) Under the large system approximation (28), we obtain

e c Σ 2 = D( 1 N 1 T τ x λ) = D(S ′ τ ∞ x ). (37)
Substituting ( 32) and ( 37) into (35), we get

D(τ ∞ s ) = D(S ′ τ ∞ x )+σ 2 v I - 1 N 2 1 T τ ∞ x λ T τ ∞ s σ 2 v I -1
.

Therefore, we propose a simple correction for the update of τ t s in the algorithm

τ t s = 1./ σ 2 v - 1 N 2 1 T τ t x λ T τ t-1 s σ 2 v + τ t p (38)
One can verify by Lemma 1 that with this correction, (e c , e c ) converge to a fixed point of (26) and hence τ x will converge to the optimal MSE.

CONVERGENCE ANALYSIS

Since the variance updates have been modified, the convergence of the variance subsystem needs to be reanalyzed. We provide here a sketch analysis following [START_REF] Kurisummoottil | Towards Convergent Approximate Message Passing by Alternating Constrained Minimization of Bethe Free Energy[END_REF], where we analyze the contractivity via the infinity norm of the Jacobian.

From the variance updates

1./τ t x = 1./σ 2 x + S T τ t-1 s , 1./τ t s = σ 2 v 1 + Sτ t x - σ 2 v N 2 (1 T τ t x )(λ T τ t-1 s ), (39)
we obtain the following Jacobians

J t xs = ∂τ t x ∂τ t-1,T s = -D t,2 x S T ; J t sx = ∂τ t s ∂τ t,T x = -D t,2 s (S - σ 2 v N 2 (λ T τ t-1 s )1 M 1 T N ); J t ss = ∂τ t s ∂τ t-1,T s = -D t,2 s ( σ 2 v N 2 (1 T τ t x )1 N λ T ) . (40)
Following the chain rule, we write with the gradient of τ t s

G t = -(D t s ) dτ t s dτ t-1 s D t-1 s = -(D t s ) -1 (J t sx J t xs + J t ss )D t-1 s = -D t s S - σ 2 v N (λ T τ t-1 s )1 M 1 T N D t,2 x S T D t-1 s -D t s σ 2 v N 2 (1 T τ t x )1 N λ T D t-1 s .
(41) The optimal test vector for the matrix infinity norm will again be 1. We then have

G t 1 = D t s S- σ 2 v N (λ T τ t-1 s )1 M 1 T N D t x D t x S T τ t-1 s +D t s σ 2 v N 2 (1 T τ t x )1 N λ T τ t-1 s . (42)
From (39), we see that

D t-1 x S T τ t-1 s ⪯ 1, leading to G t 1 ⪯ D t s S- σ 2 v N (λ T τ t-1 s )1 M 1 T N τ t x +D t s σ 2 v N 2 (1 T τ t x )1 N λ T τ t-1 s = D t s Sτ t x (43)
From (39), D t s Sτ t x ⪯ 1, and hence, ∥G t ∥ ∞ < 1.

GAUSSIAN MIXTURE MODEL PRIOR CASE

True MMSE Solution

Assume that the prior distribution of each x i is given by:

p xi (x i ) = L n=1 α n N (x i ; µ ni , σ 2 ni ); 3 n=L α n = 1.
(44) The true MMSE estimation mean and covariance matrix for a Gaussian mixture are derived in [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF] x

MMSE = E[x|y], C MMSE = E[xx T |y] -E[x|y] E[x|y] T . (45)

AMBUAMP Scalar MMSE Calculations

The computation of the posterior (13) in UAMP with a Gaussian mixture prior (44) can be derived analogously to (45), see [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF]. Define Z ni :

Z ni = α n N (x; µ ni , σ 2 ni )N (x; r i , τ ri )dx = αn √ 2π(σ 2 ni +τr i ) e - (µ ni -r i ) 2 2(σ 2 ni +τr i ) . (46)
Then we have

x i = x 3 n=1 αnN (x;µni,σ 2
τ xi = x 2 3 n=1 αnN (x;µni,σ 2 ni )N (x;ri,τr i )dx 3 n=1 αnN (x;µni,σ 2 ni )N (x;ri,τr i )dx -x 2 i = 3 n=1 ( µ ni τr i +r i σ 2 ni σ 2 ni +τr i ) 2 + τr i σ 2 ni σ 2 ni +τr i Zni 3 n=1 Zni -x 2 i . (47)

SIMULATION RESULTS

To verify our findings, we simulate the scenarios where the measurement matrices A ∈ R 5×10 are ill-conditioned. We set the condition number of the measurement matrix to 100 and the diagonal entries in the singular matrix of A form a geometry series. We fix the SNR to 20dB. Each element of the signal vector x is independently drawn from a Gaussian mixture model p xi (x i ) = 0.5N (x i ; 0, 4 × 0.25 i-1 ) + 0.5N (x i ; 0, 0.5 i-1 ) (49) To evaluate the averaged posterior variance prediction error, we compare the normalized square of the sum difference, which is defined as

(tr[C MMSE ] -1 T τ x ) 2 tr[C MMSE ] 2 . ( 50
)
Similarly, we use the normalized sum of squares to evaluate the individual prediction error of the posterior variances, which is defined as

(diag[C MMSE ] -τ x ) T (diag[C MMSE ] -τ x ) tr[C MMSE ] 2 . ( 51
)
The simulation results appear in Fig. 1. It is an average of 20 realizations of different measurement matrices A. The correction term is seen to lead to improved variance predictions even at these small dimensions.

CONCLUDING REMARKS

In this work, we have investigated the recovery of a sparse signal vector with a non-identically independently distributed arbitrary prior. This occurs typically in a compressed sensing problem but the approach is applicable to MMSE estimation in any GLM with Gaussian noise. The reVAMP algorithm we have introduced in [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF] also solves this problem without any convergence or variance approximation issues, only involving the asymptotic Gaussian extrinsic approximation. However, the AMBUAMP proposed here addresses the high complexity issues in reVAMP for high dimensions. It combines the convergent AMP in [START_REF] Kurisummoottil | Towards Convergent Approximate Message Passing by Alternating Constrained Minimization of Bethe Free Energy[END_REF] with the Haar LSA based variance correction in [START_REF] Zhao | Variance Predictions in VAMP/UAMP with Right Rotationally Invariant Measurement Matrices for niid Generalized Linear Models[END_REF]. The Haar LSA is based on the underlying approximate multivariate Gaussian posterior approximation elucidated in [START_REF] Zhao | Approximate Message Passing for Not So Large niid Generalized Linear Models[END_REF]. The Haar LSA is formulated for the sum MSE. However, the expressions in (28) indicate that asymptotically also the individual component MSEs become correct, based on the same sum MSE variance correction. We have illustrated AMBUAMP for a GMM prior, leading to a GMM posterior. The efficient computation of MMSE estimates with Gaussian extrinsics and arbitrary prior deserves further work however.

  The scalar e can be obtained as the unique solution (fixed point) of the following system of equations, eI + D) -1 .

  ni )N (x;ri,τr i )dx 3 n=1 αnN (x;µni,σ 2 ni )N (x;ri,τr i )dx

Fig. 1 .

 1 Fig. 1. MSE with or without correction.