SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma 2.4 (Behavior of the inhomogeneous AMP iterates). Suppose that Assumption 1.1 holds, and that  $\mathbb{P}_0$  has a bounded second moment. Let  $\phi : \mathbb{R}^2 \to \mathbb{R}$  be a  $L$ -pseudo-Lipschitz test function satisfying (4). For any  $a \in [q]$ , then the following limit holds almost surely
$$
\lim  _ {N \to \infty} \frac {1}{| C _ {a} ^ {N} |} \sum_ {i \in C _ {a} ^ {N}} \phi (x _ {i} ^ {t}, x _ {i} ^ {\star}) = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \phi (\sigma_ {a} ^ {t} Z, x _ {0} ^ {\star})
$$
where  $Z$  is an independent standard Gaussian.
Proof. In matrix-AMP [Javanmard and Montanari, 2013, Th. 1], the marginals of the iterates  $\boldsymbol{r}^t$  from  $(\mathbf{A}, \tilde{\mathcal{F}}(\Delta), r^0)$  are approximately Gaussian and encoded by the positive definite matrices
$$
\widehat {\boldsymbol {\Sigma}} _ {a} ^ {t} = \mathbb {E} \left[ \tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right) \tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right) ^ {\top} \right] \quad \text {f o r a l l} i \in C _ {a} ^ {N}, \tag {19}
$$
$$
\text {w i t h} \boldsymbol {Z} ^ {t} \sim N \left(0, \boldsymbol {\Sigma} ^ {t}\right) \text {a n d} \Sigma^ {\mathrm {t} + 1} = \sum_ {\mathrm {a} = 1} ^ {\mathrm {q}} \mathrm {c} _ {\mathrm {a}} \widehat {\Sigma} _ {\mathrm {a}} ^ {\mathrm {t}}. \tag {20}
$$
We now show that  $a \in [q]$  and  $i \in C_a^N$ ,  $\widehat{\Sigma}_a^t$  depends only on  $(\sigma_a^t)^2 = \Sigma_{aa}^t$ . Indeed, by the definition of  $\tilde{f}_t^i$  we have
$$
\widehat {\boldsymbol {\Sigma}} _ {a} ^ {t} (k, l) = \mathbb {E} \left[ \left(\tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right)\right) _ {k} \left(\tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right)\right) _ {l} \right] = \mathbb {E} \left[ \frac {1}{\sqrt {\tilde {\boldsymbol {\Delta}} _ {a k}}} \frac {1}{\sqrt {\tilde {\boldsymbol {\Delta}} _ {a l}}} \left(f _ {t} ^ {a} \left(Z _ {a} ^ {t}\right)\right) ^ {2} \right], \tag {21}
$$
where  $Z_{a}^{t}\sim N(0,(\sigma_{a}^{t})^{2})$  is the  $a$ th component of the Gaussian vector  $Z_{t}$ . The key observation here is that by construction our function  $\tilde{f}_t^i,\mathbb{R}^q\mapsto \mathbb{R}^q$  depends only on the  $i$ th component  $Z_{i}^{t}$  of the Gaussian vector  $Z^{t}$ . To characterize the limiting distribution of  $x^{t} = \mathrm{blockproj}(r^{t})$ , we only need to keep track of the variances  $(\sigma_j^t)^2,j\in [N]$ . Using (21), for a given  $a\in [q]$  and  $i\in C_a^N$  we write
$$
\widehat {\boldsymbol {\Sigma}} _ {a} ^ {t} (g (j), g (j)) = \mathbb {E} \left[ \left(\tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right)\right) _ {g (j)} \left(\tilde {f} _ {t} ^ {i} \left(\boldsymbol {Z} ^ {t}\right)\right) _ {g (j)} \right] = \frac {1}{\tilde {\Delta} _ {a g (j)}} \mathbb {E} \left[ \left(f _ {t} ^ {i} \left(Z _ {i} ^ {t}\right)\right) ^ {2} \right]. \tag {22}
$$
Finally, with (20) we get that for any  $b \in [q]$  and any  $j \in C_b^N$ , using  $Z_a^t \sim N(0, (\sigma_a^t)^2)$ ,
$$
\left(\sigma_ {j} ^ {t + 1}\right) ^ {2} = \left(\sigma_ {b} ^ {t + 1}\right) ^ {2} = \boldsymbol {\Sigma} _ {b b} ^ {t + 1} = \sum_ {a = 1} ^ {q} c _ {a} \widehat {\boldsymbol {\Sigma}} _ {a} ^ {t} (g (j), g (j)) = \sum_ {a = 1} ^ {q} \frac {c _ {a}}{\tilde {\boldsymbol {\Delta}} _ {a g (j)}} \mathbb {E} \left[ \left(f _ {t} ^ {a} \left(Z _ {a} ^ {t}\right)\right) ^ {2} \right] \tag {23}
$$
The inhomogeneous spiked Wigner model in the light of the AMP approach — We now generalize the state evolution equations from Lemma 2.4 to spiked matrices with an inhomogeneous noise profile as was stated in Theorem 1.2. This reduction via a change of variables is standard, see for example [Deshpande et al., 2015, Lemma 4.4]. Remember that in the inhomogeneous version of the spiked Wigner model we observe the signal  $x^{\star}$  through an inhomogeneous channel:
$$
\boldsymbol {Y} = \sqrt {\frac {1}{N}} \boldsymbol {x} ^ {\star} \left(\boldsymbol {x} ^ {\star}\right) ^ {T} + \boldsymbol {A} \odot \sqrt {\boldsymbol {\Delta}}. \tag {24}
$$
Our AMP algorithm is defined with the following recursion:
$$
\boldsymbol {x} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) f _ {t} \left(\boldsymbol {x} ^ {t}\right) - \mathbf {b} _ {t} \odot f _ {t - 1} \left(\boldsymbol {x} ^ {t - 1}\right) \tag {25}
$$
where  $\mathbf{b}_t = \frac{1}{N\Delta} f_t'(\boldsymbol{x}^t)$  and  $f$  is encoded by the family of functions in Definition 2.2. The main difference in contrast to the iteration (13) is that our data matrix  $\mathbf{Y}$  is no longer a centered matrix, while  $\frac{1}{\sqrt{\Delta}} \odot \mathbf{A}$  is. We would like to reduce (25) to an iteration of the form (13) with respect to a different parameter  $\boldsymbol{s}^t$  which is uniquely determined by  $\boldsymbol{x}^t$
$$
\boldsymbol {s} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) f _ {t} (\boldsymbol {s} ^ {t}) - \mathbf {b} _ {t} \odot f _ {t - 1} (\boldsymbol {s} ^ {t - 1}).
$$
Doing so will allow us to recover the limiting laws of the iterates from Lemma 2.4. This is done via a standard change of variables to recenter  $\mathbf{Y}$ . We will sketch the argument in this section, and defer the full proof of Theorem 1.2 to the Appendix A in the Supplementary Material.
To simplify notation, let us denote  $f_{t}(\boldsymbol{x}^{t}) \coloneqq \hat{\boldsymbol{x}}^{t}$ . We proceed following the approach of [Deshpande et al., 2015, Lemma 4.4]. We rewrite (25) using the definition of  $\mathbf{Y}$  to get
$$
\begin{array}{l} \boldsymbol {x} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) \hat {\boldsymbol {x}} ^ {t} - \mathbf {b} _ {t} \odot \hat {\boldsymbol {x}} ^ {t - 1} \tag {26} \\ = \left(\frac {1}{N \boldsymbol {\Delta}} \odot \boldsymbol {x} ^ {\star} (\boldsymbol {x} ^ {\star}) ^ {T}\right) \hat {\boldsymbol {x}} ^ {t} + \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) \hat {\boldsymbol {x}} ^ {t} - \mathbf {b} _ {t} \odot \hat {\boldsymbol {x}} ^ {t - 1}. \\ \end{array}
$$
If indices are independent, then by the strong law of large numbers one would expect that
$$
\left(\left(\frac {1}{N \boldsymbol {\Delta}} \odot \boldsymbol {x} ^ {\star} (\boldsymbol {x} ^ {\star}) ^ {\boldsymbol {T}}\right) \hat {\boldsymbol {x}} ^ {t}\right) _ {j} = x _ {j} ^ {\star} \sum_ {a \in [ q ]} \sum_ {i \in C _ {a} ^ {N}} \frac {1}{N} \frac {x _ {i} ^ {\star} \hat {x} _ {i} ^ {t}}{\boldsymbol {\Delta} _ {j i}} \rightarrow x _ {j} ^ {\star} \sum_ {a \in [ q ]} \frac {c _ {a}}{\boldsymbol {\Delta} _ {j i _ {a}}} \mathbb {E} \left[ x _ {0} ^ {\star} \hat {x} _ {i _ {a}} ^ {t} \right], \tag {27}
$$
where  $i_{a}$  is some index belonging to the group  $C_a^N$  and  $x_0^\star$  is a random variable distributed according to the prior distribution  $\mathbb{P}_0$ . For  $b \in [q]$  and  $i \in C_b^N$ , we define the block overlap  $\mu_b^t$  using the recursion
$$
\mu_ {i} ^ {t + 1} = \mu_ {b} ^ {t + 1} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\tilde {\Delta} _ {a b}} \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right], \tag {28}
$$
where  $Z$  is a standard Gaussian random variable independent from all others sources of randomness. Notice that (28) is precisely the asymptotic behavior of the summation appearing in (27).
We now make a change of variables and track the iterates
$$
\boldsymbol {s} ^ {0} = \boldsymbol {x} ^ {0} - \boldsymbol {\mu} ^ {0} \odot \boldsymbol {x} ^ {\star} \quad \boldsymbol {s} ^ {t} = \boldsymbol {x} ^ {t} - \boldsymbol {\mu} ^ {t} \odot \boldsymbol {x} ^ {\star}, \quad t \geq 1 \tag {29}
$$
where  $\mu^0$  is the vector of block overlaps of the initial condition  $x^0$  with the truth. We reduced the (25) iteration to the following iteration in which we easily recognize a version of (13):
$$
\boldsymbol {s} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) f _ {t} \left(\boldsymbol {s} ^ {t} + \boldsymbol {\mu} ^ {t} \odot \boldsymbol {x} ^ {\star}\right) - \mathbf {b} _ {t} \odot f _ {t - 1} \left(\boldsymbol {s} ^ {t - 1} + \boldsymbol {\mu} ^ {t - 1} \odot \boldsymbol {x} ^ {\star}\right) \tag {30}
$$
with the initial condition  $s^0 = x^0 - \mu^0 \odot x^*$  and the Onsager term taken from (14) is given by
$$
\mathbf {b} _ {t} = \frac {1}{N \boldsymbol {\Delta}} f _ {t} ^ {\prime} \left(\boldsymbol {s} ^ {t} + \boldsymbol {\mu} ^ {t} \odot \boldsymbol {x} ^ {\star}\right). \tag {31}
$$
Using Lemma 2.4, we can recover the asymptotic behavior of the iterates  $\boldsymbol{x}^t$  given in (25) by computing the iterates  $s^t + \mu_t \odot x^\star$  where  $s^t$  follows (30) and  $\mu_t$  satisfies (28). From this reduction we obtain the following state evolution equations describing the behaviour of (25):
1.  $x_{j}^{t}\cong \mu_{g(j)}^{t}x_{0}^{\star} + \sigma_{g(j)}^{t}Z$  for  $j\in [N]$ , where  $Z\sim \mathcal{N}(0,1)$
2.  $\mu_b^{t + 1} = \sum_{a\in [q]}\frac{c_a}{\tilde{\Delta}_{ab}}\mathbb{E}_{x_0^\star ,Z}[x_0^\star f_t^a (\mu_a^t x_0^\star +\sigma_a^t Z)]$  with  $x_0^\star \sim \mathbb{P}_0,Z\sim \mathcal{N}(0,1)$
3.  $(\sigma_{b}^{t + 1})^{2} = \sum_{a = 1}^{q}\frac{c_{a}}{\bar{\Delta}_{ab}}\mathbb{E}_{x_{0}^{\star},Z}\left[(f_{t}^{a}(\mu_{a}^{t}x_{0}^{\star} + \sigma_{a}^{t}Z))^{2}\right]$  with  $x_0^\star \sim \mathbb{P}_0$ $Z\sim \mathcal{N}(0,1)$
This (informally) characterizes the limiting distribution of the state evolution of the iterates from the inhomogeneous AMP stated in Theorem 1.2.
Fixed-point equation of state evolution in the Bayes-optimal setting — Suppose that we know the prior distribution  $\mathbb{P}_0$  of  $x_0^\star$ . The Bayes-optimal choice for the denoising functions  $f_t^j, j \in [N]$  is simply the expectation of  $x_0^\star$  with respect to the posterior distribution,
$$
f _ {t} ^ {j} (r) = f _ {t} ^ {g (j)} (r) = \mathbb {E} _ {\text {p o s t e r i o r}} \left[ x _ {0} ^ {\star} \mid \mu_ {g (j)} ^ {t} x _ {0} ^ {\star} + \sigma_ {g (j)} ^ {t} Z = r \right]. \tag {32}
$$
Under this Bayes-optimal setting, we can simplify the equations obtained in the previous section and see that AMP estimator is indeed an optimal one by studying its fixed point.
Proof of Theorem 1.4. For this choice of  $f_{t}^{j}$  the Nishimori identity (see for example [Lelarge and Miolane, 2019, Proposition 16]) states that for  $a \in [q]$  and  $j \in C_a^N$ ,
$$
\tilde {\mu} _ {a} ^ {t} := \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right] = \mathbb {E} \left[ \left(f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right)\right) ^ {2} \right]. \tag {33}
$$
In this setting, the state evolution equations from Theorem 1.2 reduce to
$$
\left\{ \begin{array}{l} \tilde {\mu} _ {a} ^ {t} = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right] \\ \mu_ {b} ^ {t + 1} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\bar {\Delta} _ {a b}} \tilde {\mu} _ {a} ^ {t}, b \in [ q ] \\ (\sigma_ {b} ^ {t + 1}) ^ {2} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\bar {\Delta} _ {a b}} \tilde {\mu} _ {a} ^ {t}, b \in [ q ]. \end{array} \right. \tag {34}
$$
Remarkably with the Bayes-optimal choice of the denoising functions we have that for  $t \geq 1$  for each block  $b \in [q]$ ,  $\mu_b^{t+1} = (\sigma_b^{t+1})^2$ . Therefore a necessary and sufficient condition for an estimator to be a fixed point of the state evolution is to simply have its overlaps  $\mu_b^t$  unchanged by an iteration of the state evolution. This translates into the following equation for the overlaps  $\mu_b, b \in [q]$
$$
\mu_ {b} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\tilde {\Delta} _ {a b}} \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} \mathbb {E} _ {\text {p o s t e r i o r}} \left[ x _ {0} ^ {\star} \mid \mu_ {a} x _ {0} ^ {\star} + \sqrt {\mu_ {a}} Z \right] \right]. \tag {35}
$$
The result of Theorem 1.4 now follows immediately.
3 A spectral method adapted to the inhomogeneous spiked Wigner model
From AMP to a spectral method — Remarkably, AMP and the state evolution machinery associated with it can help us design a simple spectral algorithm that matches the information-theoretic phase transition [Guionnet et al., 2022, Remark 2.16]. Recall that Theorem 1.2 does not require the denoising functions  $f_{t}$  to be Bayes-optimal, but can be applied to any Lipschitz family of functions. In this section, we analyze the state evolution for the family of identity functions,  $f_{t}(x) = x$ . We can gain some intuition motivating such a choice by considering simple priors, such as the Rademacher prior. In the case of the Rademacher prior a simple computation shows that the Bayes-optimal choice of the denoising functions yields  $f_{t}^{j}(\star) = \tanh(\star)$ . In the first order approximation we have  $\tanh(x) \approx x$ . Thus, at least for the Rademacher prior, the choice of identity functions as denoising functions corresponds to the first order approximation of the Bayes-optimal choice. By Remark 3.1, we can assume that the entries of  $x^{\star}$  have unit variance. With this choice of denoising functions the AMP iteration will simply become:
$$
\boldsymbol {x} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) \boldsymbol {x} ^ {t} - \mathbf {b} _ {t} \odot \boldsymbol {x} ^ {t - 1} \quad \text {w h e r e} \quad \mathbf {b} _ {t} = \frac {1}{N \boldsymbol {\Delta}} f _ {t} ^ {\prime} = \frac {1}{N \boldsymbol {\Delta}} \left[ \begin{array}{l} 1 \\ \vdots \\ 1 \end{array} \right]. \tag {36}
$$
If we denote  $B_{t} = \mathrm{diag}(\mathbf{b}_{t})$ , it is easy to see that the fixed point of this iteration yields
$$
\boldsymbol {x} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) \boldsymbol {x} - \boldsymbol {B} _ {t} \boldsymbol {x} \tag {37}
$$
so any  $x$  fixed by the AMP iteration (36) must be an eigenvector of the matrix
$$
\tilde {\boldsymbol {Y}} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) - \boldsymbol {B} _ {t} = \left(\frac {1}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y}\right) - \operatorname {d i a g} \left(\frac {1}{N \boldsymbol {\Delta}} \left[ \begin{array}{l} 1 \\ \vdots \\ 1 \end{array} \right]\right). \tag {38}
$$
A simple spectral method consists in taking the principal eigenvector (associated to the largest eigenvalue) of the matrix  $\tilde{\mathbf{Y}} = \frac{\mathbf{Y}}{\sqrt{N}\Delta} -\mathbf{B}_t$
Analysis of the spectral method using state evolution — The spectral algorithm described above behaves like the AMP iteration with identity denoising functions around its fixed point. Therefore we can analyze this spectral algorithm using state evolution machinery for the AMP iteration. In the case of identity functions we have  $f_{t}^{a}(x) = x$  for all  $a \in [q]$ , so
$$
\mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right] = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right] = \mu_ {a} ^ {t} \tag {39}
$$
$$
\mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ \left(f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right)\right) ^ {2} \right] = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) ^ {2} \right] = \left(\mu_ {a} ^ {t}\right) ^ {2} + \left(\sigma_ {a} ^ {t}\right) ^ {2} \tag {40}
$$
which transforms state evolution equations (5) into the following simple form:
$$
\mu_ {b} ^ {t + 1} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\Delta_ {b a}} \mu_ {a} ^ {t} \quad \text {a n d} \quad \left(\sigma_ {b} ^ {t + 1}\right) ^ {2} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\Delta_ {b a}} \left(\left(\mu_ {a} ^ {t}\right) ^ {2} + \left(\sigma_ {a} ^ {t}\right) ^ {2}\right). \tag {41}
$$
Rewriting the overlap state evolution in a matrix form we get for  $\pmb{c} = (c_{a})_{a\in [q]}$  that
$$
\operatorname {d i a g} (\sqrt {c}) \boldsymbol {\mu} ^ {t + 1} = \operatorname {d i a g} (\sqrt {c}) \frac {1}{\Delta} \operatorname {d i a g} (\sqrt {c}) \left(\operatorname {d i a g} (\sqrt {c}) \boldsymbol {\mu} ^ {t}\right). \tag {42}
$$
The special form of state evolution above is rather informative and we can interpret it as follows. First of all, in the regime where  $\lambda (\Delta) = \left\| \mathrm{diag}(\sqrt{c})\frac{1}{\Delta}\mathrm{diag}(\sqrt{c})\right\|_{op} < 1$ , we can see from (42) that any iteration of the AMP recursion (36) contracts the vector overlap (multiplying it by a matrix with the operator norm smaller than 1). Thus the only possible overlap of a fixed point of (36) is 0. Moreover, we have defined the matrix  $\tilde{\pmb{Y}}$  in a way that any eigenvector of this matrix is a fixed point of the AMP recursion (36). Thus any eigenvector of  $\tilde{\pmb{Y}}$  must have zero expected overlap with the signal, meaning that the spectral method is uninformative in the regime  $\lambda (\Delta) < 1$ .
Second of all, even though we cannot say anything rigorous in the regime  $\lambda (\Delta) > 1$ , looking at state evolution equations of the AMP recursion (36), we see instability (overlap blows up but so does the variance). We conjecture that the principal eigenvector of  $\tilde{\mathbf{Y}}$  correlates with the signal exactly in the regime  $\lambda (\Delta) > 1$  which would imply that the BBP transition happens precisely at  $\lambda (\Delta) = 1$ .
Remark 3.1. In general, we can let  $\gamma = \mathbb{E}_{x_0^*}[(x_0^\star)^2]$  and consider the normalized matrix
$$
\bar {\boldsymbol {Y}} = \frac {\boldsymbol {Y}}{\gamma} = \sqrt {\frac {1}{N}} \frac {\boldsymbol {x} ^ {\star} (\boldsymbol {x} ^ {\star}) ^ {T}}{\gamma} + \boldsymbol {A} \odot \frac {\sqrt {\boldsymbol {\Delta}}}{\gamma} = \sqrt {\frac {1}{N}} \bar {\boldsymbol {x}} ^ {\star} (\bar {\boldsymbol {x}} ^ {\star}) ^ {T} + \boldsymbol {A} \odot \sqrt {\bar {\boldsymbol {\Delta}}}
$$
for  $\bar{\pmb{x}} = \frac{\pmb{x}}{\sqrt{\gamma}}$  and  $\bar{\Delta} = \frac{\Delta}{\gamma^2}$ . Notice that the entries of  $\bar{\pmb{x}}$  now have unit variance. Under this setting, the transition of the transformation in (38) applied to  $\bar{Y}$ , which appears in (7), has transition at
$$
\lambda (\bar {\Delta}) = \left\| \mathrm {d i a g} (\sqrt {\boldsymbol {c}}) \frac {1}{\bar {\Delta}} \mathrm {d i a g} (\sqrt {\boldsymbol {c}}) \right\| _ {o p} = \mathbb {E} _ {x _ {0} ^ {\star}} [ (x _ {0} ^ {\star}) ^ {2} ] ^ {2} \left\| \mathrm {d i a g} (\sqrt {\boldsymbol {c}}) \frac {1}{\Delta} \mathrm {d i a g} (\sqrt {\boldsymbol {c}}) \right\| _ {o p} <   1
$$
which is the generalized SNR defined in (8).
We thank Alice Guionnet and Lenka Zdeborova for valuable discussions. We acknowledge funding from the ERC Project LDRAM: ERC-2019-ADG Project 884584, and by the Swiss National Science Foundation grant SNFS OperaGOST, 200021 200390
Marco Mondelli, Christos Thrampoulidis, and Ramji Venkataramanan. Optimal combination of linear and spectral estimators for generalized linear models. Foundations of Computational Mathematics, 22(5):1513-1566, 2022.
Ramji Venkataramanan, Kevin Kögler, and Marco Mondelli. Estimation in rotationally invariant generalized linear models via approximate message passing. In International Conference on Machine Learning, pages 22120-22144. PMLR, 2022.
Marco Mondelli and Ramji Venkataramanan. Approximate message passing with spectral initialization for generalized linear models. In International Conference on Artificial Intelligence and Statistics, pages 397-405. PMLR, 2021.
Greg W. Anderson, Alice Guionnet, and Ofer Zeitouni. An introduction to random matrices, volume 118 of Cambridge Studies in Advanced Mathematics. Cambridge University Press, Cambridge, 2010. ISBN 978-0-521-19452-5.
A Proof of Theorem 1.2
We now provide a rigorous proof of the result that was sketched in Section 2. This proof is essentially identical to [Deshpande et al., 2015, Lemma 4.4]. Recall the iterates (30) given by  $\boldsymbol{s}^0 = \boldsymbol{x}^0 - \boldsymbol{\mu}^0 \odot \boldsymbol{x}^\star$  and
$$
\boldsymbol {s} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) f _ {t} \left(\boldsymbol {s} ^ {t} + \boldsymbol {\mu} ^ {t} \odot \boldsymbol {x} ^ {\star}\right) - \mathbf {b} _ {t} \odot f _ {t - 1} \left(\boldsymbol {s} ^ {t - 1} + \boldsymbol {\mu} ^ {t - 1} \odot \boldsymbol {x} ^ {\star}\right), \tag {43}
$$
where  $\pmb{\mu}^t = (\mu_i^t)_{i\leq N}$  is given by the recursion
$$
\mu_ {i} ^ {t + 1} = \mu_ {g (i)} ^ {t + 1} = \sum_ {a \in [ q ]} \frac {c _ {a}}{\tilde {\Delta} _ {a g (i)}} \mathbb {E} _ {x _ {0} ^ {\star}, Z} \left[ x _ {0} ^ {\star} f _ {t} ^ {a} \left(\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z\right) \right]
$$
and  $\pmb{x}^{\star} = (x_{i}^{\star})_{i\in [N]}$  is a vector with independent coordinates distributed according to  $\mathbb{P}_0$ . By Lemma 2.4, for each  $a\in [q]$ , and any pseudo-Lipschitz function  $\phi :\mathbb{R}\to \mathbb{R}\mapsto \mathbb{R}$  we have that almost surely
$$
\lim  _ {N \rightarrow \infty} \frac {1}{\left| C _ {a} ^ {N} \right|} \sum_ {i \in C _ {a} ^ {N}} \phi \left(s _ {i} ^ {t}, x _ {i} ^ {\star}\right) = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \phi \left(\sigma_ {a} ^ {t} Z, x _ {0} ^ {\star}\right) \tag {44}
$$
$$
\left(\sigma_ {b} ^ {t + 1}\right) ^ {2} := \sum_ {a = 1} ^ {q} \frac {c _ {a}}{\tilde {\Delta} _ {a b}} \mathbb {E} _ {Z} \left[ \left(f _ {t} ^ {b} \left(Z _ {b} ^ {t}\right)\right) ^ {2} \right].
$$
as was defined in (18). For any pseudo-Lipschitz function  $\psi : \mathbb{R} \to \mathbb{R}$ , we have  $\phi(x, y) = \phi(x - \mu_a^t y)$  is also pseudo-Lipschitz, so (44) implies that
$$
\lim  _ {N \rightarrow \infty} \frac {1}{\left| C _ {a} ^ {N} \right|} \sum_ {i \in C _ {a} ^ {N}} \psi \left(s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star}\right) = \mathbb {E} _ {x _ {0} ^ {\star}, Z} \psi \left(\sigma_ {a} ^ {t} Z + \mu_ {a} ^ {t} x _ {0} ^ {\star}\right) \tag {45}
$$
Now let  $\boldsymbol{x}^t$  be the iterates from the spiked AMP iteration for the inhomogeneous Wigner matrix (24) we derived in (26)
$$
\boldsymbol {x} ^ {t + 1} = \left(\frac {1}{N \boldsymbol {\Delta}} \odot \boldsymbol {x} ^ {\star} \left(\boldsymbol {x} ^ {\star}\right) ^ {T}\right) \hat {\boldsymbol {x}} ^ {t} + \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) \hat {\boldsymbol {x}} ^ {t} - \mathbf {b} _ {t} \odot \hat {\boldsymbol {x}} ^ {t - 1}. \tag {46}
$$
It now suffices to show that for fixed  $t$  and all  $a \in [q]$  that
$$
\lim  _ {N \rightarrow \infty} \frac {1}{\left| C _ {a} ^ {N} \right|} \sum_ {i \in C _ {a} ^ {N}} \left(\psi \left(s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star}\right) - \psi \left(x _ {i} ^ {t}\right)\right) = 0 \tag {47}
$$
almost surely. This will imply that  $s_i^t + \mu_a^t x_i^\star$  and  $x_i^t$  have the same asymptotic distribution which finish the proof of Theorem 1.2 by (45).
We now prove (47). Since  $\psi$  is  $L$ -pseudo-Lipschitz we have
$$
\begin{array}{l} | \psi (s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star}) - \psi (x _ {i} ^ {t}) | \leq L (1 + | s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star} | + | x _ {i} ^ {t} |) | s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star} - x _ {i} ^ {t} | \\ \leq 2 L \left| s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star} - x _ {i} ^ {t} \right| \left(1 + \left| s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star} \right| + \left| s _ {i} ^ {t} + \mu_ {a} ^ {t} x _ {i} ^ {\star} - x _ {i} ^ {t} \right|\right). \\ \end{array}
$$
The Cauchy-Schwarz inequality implies that