SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Optimal Algorithms for the Inhomogeneous Spiked Wigner model
Justin Ko
justin.ko@ens-lyon.fr
Florent Krzakala
florent.krzakala@epfl.ch
Aleksandr Pak
ENS de Lyon, EPFL
aleksandr.pak@epfl.ch
We study a spiked Wigner problem with an inhomogeneous noise profile. Our aim in this problem is to recover the signal passed through an inhomogeneous low-rank matrix channel. While the information-theoretic performances are well-known, we focus on the algorithmic problem. First, we derive an approximate message-passing algorithm (AMP) for the inhomogeneous problem and show that its rigorous state evolution coincides with the fixed-point equations satisfied by the Bayes-optimal estimator. Second, we deduce a simple and efficient spectral method that outperforms PCA and is shown to match the information-theoretic transition.
Low-rank information extraction from a noisy data matrix is a crucial statistical challenge. The spiked random matrix models have recently gained extensive interest in the fields of statistics, probability, and machine learning, serving as a valuable platform for exploring this issue Donoho and Johnstone [1995], Péché [2014], Baik et al. [2005]. A prominent example is the spiked Wigner model, where a rank one matrix is observed through a component-wise homogeneous noise.
Here we consider the inhomogeneous version of the model which has been recently introduced in a series of papers Barbier and Reeves [2020], Behne and Reeves [2022], Alberici et al. [2021a, 2022] and naturally arises in different contexts such as community detection Behne and Reeves [2022], deep Boltzmann machines Alberici et al. [2021b] etc. In particular, it corresponds to the dense limit of the degree-corrected stochastic block model Karrer and Newman [2011] (see the explicit mapping in Guionnet et al. [2022]). In this model the signal is observed through an inhomogeneous Gaussian noise with block-constant variance profile. Such block-constant noise naturally arises in settings such as community detection and stochastic-block models. The assumption for the noise to be Gaussian is in no way restrictive as a universality result in Guionnet et al. [2022] shows that an entrywise transformation of data generated from a general inhomogeneous inference problem reduces any probabilistic noise to the spiked Wigner problem with Gaussian data. This generality is what makes the model appealing.
We now define the model: Consider a partition  $\{1, \ldots, N\} = [N]$  into  $q$  disjoint groups  $C_1^N \cup \dots \cup C_q^N = [N]$ . This partition is encoded by a function  $g : [N] \mapsto [q]$  which maps each index  $i \in [N]$  into its group  $g(i) \in [q]$ . Let  $\tilde{\Delta} \in \mathbb{R}^{q \times q}$  be a symmetric matrix encoding a block-constant symmetric matrix  $\Delta \in \mathbb{R}^{N \times N}$
$$
\boldsymbol {\Delta} _ {i j} = \tilde {\boldsymbol {\Delta}} _ {g (i) g (j)}. \tag {1}
$$
We assume both the noise profile  $\tilde{\Delta}$  and the partition function  $g$  to be known. This is often the case in practical applications, such as the degree-corrected block model, Karrer and Newman [2011]
Additionally, in practical applications where the noise profile is not known, one can often empirically estimate the noise profile and assign group membership according to the empirical variances. Such an approach is likely to yield good results but lies outside the scope of the present work.
We observe the signal  $\pmb{x}^{\star} \in \mathbb{R}^{N}$  which is assumed to have independent identically distributed coordinates generated from some prior distribution  $\mathbb{P}_0$  (i.e.  $\mathbb{P}(\pmb{x}^{\star} = \pmb{x}) = \prod_{i=1}^{N} \mathbb{P}_0(x_i^{\star} = x_i)$ ) through noisy measurements:
$$
\boldsymbol {Y} = \sqrt {\frac {1}{N}} \boldsymbol {x} ^ {\star} \left(\boldsymbol {x} ^ {\star}\right) ^ {T} + \boldsymbol {A} \odot \sqrt {\Delta}. \tag {2}
$$
Here and throughout the article  $\odot$  denotes the Hadamard product,  $\sqrt{\Delta}$  is the Hadamard square-root of  $\Delta$  and  $A$  is a real-valued symmetric GOE matrix with off-diagonal elements of unit variance. The Bayes-optimal performance of this model in the asymptotic limit  $N\to \infty$  was studied rigorously in Guionnet et al. [2022], Behne and Reeves [2022], Chen and Xia [2022], Chen et al. [2021] who characterized the fundamental information-theoretic limit of reconstruction in this model.
Here we focus instead on the algorithmic problem of reconstructing the (hidden) spike. Our contributions are many-fold:
- We show how one can construct an Approximate Message Passing (AMP) algorithm for the inhomogeneous Wigner problem, whose asymptotic performance can be tracked by a rigorous state evolution, generalizing the homogeneous version of the algorithm for low-rank factorization [Bayati and Montanari, 2011, Deshpande et al., 2015, Lesieur et al., 2017].
- We derive a fixed-point equation for AMP and show that it coincides with the fixed-point equation for the Bayes-optimal estimator obtained in Guionnet et al. [2022]
- Finally, we present a linear version of AMP [Maillard et al., 2022], that is equivalent to a spectral method. We conjecture it to be optimal in the sense that it can detect the presence of the spike in the same region as AMP. This is quite remarkable since, as shown in [Guionnet et al., 2022, Section 2.4], the standard spectral method (PCA) fails to do so.
Related work - The class of approximate message passing algorithms (AMP) has attracted a lot of attention in the high-dimensional statistics and machine learning community, see e.g. [Donoho et al., 2009, Bayati and Montanari, 2011, Rangan, 2011, Deshpande et al., 2015, Lesieur et al., 2017, Gerbelot and Berthier, 2021, Feng et al., 2022]. The ideas behind this algorithm have roots in physics of spin glasses Mészard et al. [1987], Bolthausen [2014], Zdeboróva and Krzakala [2016]. AMP algorithms are optimal among first order methods [Celentano et al., 2020], thus their reconstruction threshold provides a bound on the algorithmic complexity in our model. Our approach to the inhomogeneous version of AMP relies on several refinements of AMP methods to handle the full complexity of the problem, notably the spatial coupling technique Krzakala et al. [2012], Donoho et al. [2013], Javanmard and Montanari [2013], Gerbelot and Berthier [2021], Rossetti and Reeves [2023].
Figure 1: Performance of the inhomogeneous AMP algorithm against the information-theoretical optimal MMSE. The variance profile is proportional to  $\tilde{\Delta} = \begin{bmatrix} 1 & 3 \\ 3 & 2 \end{bmatrix}$  with two equally sized blocks with standard Gaussian prior when  $N = 500$  at various snr.
Factorizing low-rank matrices is a ubiquitous problem with many applications in machine learning and statistics, ranging from sparse PCA to community detection and sub-matrix localization. Many variants of the homogeneous problem have been studied in the high-dimensional limit [Deshpande et al., 2015, Lesieur et al., 2017, Barbier et al., 2018, Lesieur et al., 2017, Alaoui et al., 2020, Lelarge and Miolane, 2019, Louart and Couillet, 2018, Barbier and Reeves, 2020]. The inhomogeneous version was discussed in detail in Guionnet et al. [2022], Behne and Reeves [2022], Alberici et al. [2021a]. Spectral methods are a very popular tool to solve rank-factorization problems [Donoho and Johnstone, 1995, Péché, 2014, Baik et al., 2005]. Using AMP as an inspiration for deriving
new spectral methods was discussed, for instance, in Saade et al. [2014], Lesieur et al. [2017], Aubin et al. [2019], Mondelli and Montanari [2018], Mondelli et al. [2022], Maillard et al. [2022], Venkataramanan et al. [2022].
1 Main results
Message passing algorithm — For each  $t \geq 0$ , let  $(f_t^a)_{a \in [q]}$  be a collection of Lipschitz functions from  $\mathbb{R} \to \mathbb{R}$ , and define for  $t \in \mathbb{N}$ $f_t : \mathbb{R}^N \mapsto \mathbb{R}^N$  by
$$
f _ {t} (\pmb {x}) := (f _ {t} ^ {g (i)} (x _ {i})) _ {i \in [ N ]} \in \mathbb {R} ^ {N}.
$$
These functions are often called denoiser functions and can be chosen amongst several options, such as the Bayes optimal denoisers for practical applications (see Section 2), or even linear denoisers (see Section 3). We shall consider the following AMP 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 {3}
$$
with the so-called Onsager term  $\mathbf{b}_t = \frac{1}{N\Delta} f_t'(x^t) \in \mathbb{R}^N$  where  $\frac{1}{\Delta}$  is the Hadamard inverse of  $\Delta$  and  $f_t'$  is the vector of coordinate wise derivatives.
In practical implementations, we initialize the algorithm with some non-null  $x^0$  and let it run for a certain number of iterations. One efficient way to do this is the spectral initialization [Mondelli and Venkataramanan, 2021] with the method described in sec. 3. In Figure 1 we provide an example of the performance of the AMP together with the Bayes-optimal estimator predicted by the asymptotic theory. Even at very moderate sizes, the agreement between theory and simulation is clear.
State evolution — AMP was the basis of many works for the homogeneous case e.g. Deshpande et al. [2015], Lesieur et al. [2017]. Our first contribution is the introduction of an AMP and of its rigorous state evolution [Javanmard and Montanari, 2013, Theorem 1] in the inhomogeneous setting. To state a well-defined limit of the AMP, we have the following assumptions.
Assumption 1.1. To ensure that our inhomogeneous AMP has a well-defined limit, we assume that
1. For each  $a \in [q]$ , we have
$$
\lim  _ {N \to \infty} \frac {\left| C _ {a} ^ {N} \right|}{N} \to c _ {a} \in (0, 1).
$$
2. For each  $t \in [N]$  and  $a \in [q]$ , both  $(f_t^a)$  and  $(f_t^a)'$  are Lipschitz.
3. For each  $a \in [q]$ , there exists  $\left(\sigma_{a}^{0}\right)^{2} \in \mathbb{R}$  such that, in probability,
$$
\lim_{N\to \infty}\frac{1}{|C_{a}^{N}|}\sum_{i\in C_{a}^{N}}f_{0}^{a}(x_{i}^{0})f_{0}^{a}(x_{i}^{0}) = \left(\sigma_{a}^{0}\right)^{2}.
$$
Our first result describes the distribution of the iterates in the limit. Our mode of convergence will be with respect to  $L$ -pseudo-Lipschitz test functions  $\phi : \mathbb{R}^M \to \mathbb{R}$  satisfying
$$
\left| \phi (x) - \phi (y) \right| \leq L \left(1 + \| x \| + \| y \|\right) \| x - y \| \quad \text {f o r a l l} x, y \in \mathbb {R} ^ {M}. \tag {4}
$$
We define the following state evolution parameters  $\mu_b^t$  and  $\sigma_b^t$  for  $b\in [q]$  through the recursion
$$
\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] \text {w i t h} x _ {0} ^ {\star} \sim \mathbb {P} _ {0}, Z \sim \mathcal {N} (0, 1)
$$
$$
\left(\sigma_ {b} ^ {t + 1}\right) ^ {2} = \sum_ {a = 1} ^ {q} \frac {c _ {a}}{\tilde {\Delta} _ {a b}} \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] \text {w i t h} x _ {0} ^ {\star} \sim \mathbb {P} _ {0}, Z \sim \mathcal {N} (0, 1), \tag {5}
$$
where  $x_0^\star$  and  $Z$  are independent. We prove that the iterates  $\pmb{x}_i^t$  are asymptotically equal in distribution to  $\mu_{g(i)}^{t}x_{0}^{\star} + \sigma_{g(i)}^{t}Z$  where  $x_0^\star \sim \mathbb{P}_0$  and  $Z$  is an independent standard Gaussian.
Theorem 1.2 (State evolution of AMP iterates in the inhomogeneous setting). Suppose that Assumption 1.1 holds, and that  $\mathbb{P}_0$  has bounded second moment. Let  $\phi : \mathbb{R}^2 \to \mathbb{R}$  be a 2-pseudo-Lipschitz test functions satisfying (4). For any  $a \in [q]$ , 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 (\mu_ {a} ^ {t} x _ {0} ^ {\star} + \sigma_ {a} ^ {t} Z, x _ {0} ^ {\star})
$$
where  $Z$  is a standard Gaussian independent from all other variables.
Remark 1.3. The notion of convergence under the  $L$ -pseudo-Lipschitz test functions induces a topology that is equivalent to the one generated by the  $L$ -Wasserstein topology [Feng et al., 2022, Remark 7.18]. We can strengthen the second moment assumption on  $\mathbb{P}_0$  to finite  $k$ th moment, but the induced topology will then change to the  $k$ -Wasserstein topology, see [Javanmard and Montanari, 2013, Theorem 1].
Even though the theoretical result above applies in the high-dimensional limit, numerical simulations show that even for medium-sized  $N$  (around 500), the behaviour of the iterates is well described by the state evolution parameters. Through the state evolution equations (5) we are able to track the iterates of the AMP iteration with just two vectors of parameters obeying the state evolution recursion: the overlap with the true signal  $(\mu_a^t)_{a\in [q]}$  and its variance  $(\sigma_a^t)_{a\in [q]}$ . We next obtain the following necessary and sufficient condition for the overlaps of a fixed point of the iteration (2):
Theorem 1.4 (Bayes-Optimal fixed point). Assume  $AMP$  satisfies Assumption 1.1 and let the denoising functions be the Bayes ones (32). Then the overlaps  $\pmb{\mu} = (\mu_{a})_{a\in [q]}$  in (5) satisfy the following fixed point equation
$$
\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 {6}
$$
Remark 1.5. The state evolution fixed point equation above coincides with the fixed point equation satisfied by the Bayes optimal estimator in [Guionnet et al., 2022, Equation 2.14].
A spectral method — Given the matrix  $\mathbf{Y}$  defined in (2) we consider the transformed matrix
$$
\tilde {\boldsymbol {Y}} := \frac {\mathbb {E} _ {x _ {0} ^ {\star}} [ (x _ {0} ^ {\star}) ^ {2} ]}{\sqrt {N} \boldsymbol {\Delta}} \odot \boldsymbol {Y} - \mathbb {E} _ {x _ {0} ^ {\star}} [ (x _ {0} ^ {\star}) ^ {2} ] ^ {2} \operatorname {d i a g} \left(\frac {1}{N \boldsymbol {\Delta}} \left[ \begin{array}{l} 1 \\ \vdots \\ 1 \end{array} \right]\right). \tag {7}
$$
Let  $\pmb{c} = (c_{a})_{a\in [q]}$ . We define the inhomogeneous signal-to-noise (SNR) ratio of such a model by
$$
\operatorname {S N R} (\boldsymbol {\Delta}) := \lambda (\boldsymbol {\Delta}) = \mathbb {E} _ {x _ {0} ^ {\star}} [ (x _ {0} ^ {\star}) ^ {2} ] ^ {2} \left\| \operatorname {d i a g} (\sqrt {\boldsymbol {c}}) \frac {1}{\tilde {\boldsymbol {\Delta}}} \operatorname {d i a g} (\sqrt {\boldsymbol {c}}) \right\| _ {o p}. \tag {8}
$$
Conjecture 1.6. The top eigenvalue of  $\tilde{\mathbf{Y}}$  separates from the bulk if and only if the signal to noise ratio  $\lambda (\Delta) > 1$ . In particular, if  $\hat{\mathbf{x}}$  is the top eigenvector of  $\tilde{\mathbf{Y}}$ , then if and only if  $\lambda (\Delta) < 1$  we have:
$$
\lim  _ {N \rightarrow \infty} \frac {| \hat {\boldsymbol {x}} \cdot \boldsymbol {x} ^ {\star} |}{\| \hat {\boldsymbol {x}} \| \| \boldsymbol {x} ^ {\star} \|} = 0.
$$
The conjecture is based on heuristic arguments identifying the fixed point of AMP with a spectral method. This matches precisely the recovery transition in [Guionnet et al., 2022, Lemma 2.15 Part (b)]. In this paper, we rigorously show that with  $\mathrm{SNR}(\Delta) < 1$  our proposed spectral method fails to recover the signal. We illustrate the eigenvalue BBP-like transition in Fig.2.
2 The inhomogeneous AMP algorithm
In this section, we derive the formula for the inhomogeneous AMP iteration (25). We first recall the general matrix framework of AMP from [Javanmard and Montanari, 2013]:
Figure 2: Illustration of the spectrum of  $\tilde{Y} \in \mathbb{R}^{10^3 \times 10^3}$  evaluated at noise profiles with snr  $\lambda(\Delta) = 0.7$  (left, before the transition) and 1.8 (right, after the transition), with the outlying eigenvector correlated with the spike arises at eigenvalue one.
Matrix AMP - In the matrix setting an AMP algorithm operates on the vector space  $\mathcal{V}_{q,N} \equiv (\mathbb{R}^q)^N \simeq \mathbb{R}^{N \times q}$ . Each element of  $\boldsymbol{v} = (v_1, \ldots, v_N) \in \mathcal{V}_{q,N}$  will be regarded as  $N$ -vector with entries  $v_i \in \mathbb{R}^q$ .
Definition 2.1 (AMP). A matrix AMP acting on this space is represented by  $(\mathbf{A},\mathcal{F},\pmb{v}^0)$ , where:
1.  $\mathbf{A} = \mathbf{G} + \mathbf{G}^{\mathrm{T}}$ , where  $\mathbf{G} \in \mathbb{R}^{N \times N}$  has iid entries  $G_{ij} \sim N(0, \frac{1}{2})$ .
2.  $\mathcal{F}$  is a family of  $N$  Lipschitz functions  $f_{t}^{i}:\mathbb{R}^{q}\mapsto \mathbb{R}^{q},i\in [N]$  indexed by time  $t\in \mathbb{N}$ . The family  $\mathcal{F}$  encodes a function  $f_{t}:\mathcal{V}_{q,N}\to \mathcal{V}_{q,N}$  that acts separately on each coordinate  $v_{j}\in \mathbb{R}^{q}$ ,
$$
f _ {t} (\boldsymbol {v}) = \left(f _ {t} ^ {1} \left(v _ {1}\right), \dots , f _ {t} ^ {N} \left(v _ {N}\right) \in \mathcal {V} _ {q, N} \right. \tag {9}
$$
3.  $\pmb{v}^0\in \mathcal{V}_{q,N}$  is a starting condition.
The algorithm itself is a sequence of iterates generated by:
$$
\boldsymbol {v} ^ {t + 1} = \frac {\boldsymbol {A}}{\sqrt {N}} f _ {t} \left(\boldsymbol {v} ^ {t}\right) - f _ {t - 1} \left(\boldsymbol {v} ^ {t - 1}\right) \mathbf {B} _ {t} ^ {T} \tag {10}
$$
where  $\mathbf{B}_t$  is the  $q\times q$  Onsager matrix given by
$$
\mathbf {B} _ {t} = \frac {1}{N} \sum_ {j = 1} ^ {N} \partial f _ {t} ^ {j} \left(\boldsymbol {v} _ {j} ^ {t}\right). \tag {11}
$$
where  $\partial f_t^j$  denotes the Jacobian matrix of  $f_{t}^{j}$ . The limiting properties of the AMP sequences are well known and can be found in [Javanmard and Montanari, 2013, Theorem 1].
The inhomogeneous AMP - We now define an inhomogeneous AMP iteration which takes into account the block-constant structure of the noise:
Definition 2.2. An inhomogeneous AMP on  $\mathcal{V}_{1,N} = \mathbb{R}^N$  is represented by  $(A, \mathcal{F}, x^0, \Delta)$ , where the terms  $A$ ,  $\mathcal{F}$ ,  $x^0$  are defined in Definition 2.1 and  $\Delta$  is the  $N \times N$  variance profile encoded by  $\tilde{\Delta} \in \mathbb{R}^{q \times q}$  and grouping  $g: [N] \to [q]$  defined by (1). We further assume that the family of functions  $\mathcal{F}$  is encoded by functions  $f_t^a: \mathbb{R} \mapsto \mathbb{R}$  for  $a \in [q]$  which define the group dependent function
$$
f _ {t} (\boldsymbol {x}) = \left(f _ {t} ^ {g (1)} \left(x _ {1}\right), \dots , f _ {t} ^ {g (N)} \left(x _ {N}\right)\right) \in \mathbb {R} ^ {N}. \tag {12}
$$
The sequence of iterates  $\pmb{x}^t\in \mathbb{R}^N$  of the  $(A,\mathcal{F},\pmb{v}^{0},\pmb{\Delta})$  are defined as follows:
$$
\boldsymbol {x} ^ {t + 1} = \left(\frac {1}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} \odot \boldsymbol {A}\right) f _ {t} \left(\boldsymbol {x} ^ {t}\right) - \mathrm {b} _ {t} \odot f _ {t - 1} \left(\boldsymbol {x} ^ {t - 1}\right), \tag {13}
$$
where  $\frac{1}{\sqrt{\Delta}}$  is the Hadamard inverse square root of the noise, and the Onsager term  $\mathbf{b}_t$  has the following form
$$
\mathbf {b} _ {t} = \frac {1}{N} \left( \begin{array}{c c c c c} \frac {1}{\Delta_ {1 1}} \left(f _ {t} ^ {g (1)}\right) ^ {\prime} \left(x _ {1} ^ {t}\right) & + & \dots & + & \frac {1}{\Delta_ {1 N}} \left(f _ {t} ^ {g (N)}\right) ^ {\prime} \left(x _ {N} ^ {t}\right) \\ \vdots & & \vdots & & \vdots \\ \frac {1}{\Delta_ {N 1}} \left(f _ {t} ^ {g (1)}\right) ^ {\prime} \left(x _ {1} ^ {t}\right) & + & \dots & + & \frac {1}{\Delta_ {N N}} \left(f _ {t} ^ {g (N)}\right) ^ {\prime} \left(x _ {N} ^ {t}\right) \end{array} \right) = \frac {1}{N \boldsymbol {\Delta}} f _ {t} ^ {\prime} \left(x ^ {t}\right) \in \mathbb {R} ^ {N}. \tag {14}
$$
State evolution of the inhomogeneous AMP — Through a continuous embedding, we will reduce our inhomogeneous AMP to the matrix AMP framework, and recover the state evolution of the inhomogeneous AMP. We define the diagonal matrix operator blockdiag:  $\mathbb{R}^N\mapsto \mathcal{V}_{q,N}$  which outputs a block diagonal matrix according to the block structure of our discretization of  $[N]$ :
$$
\operatorname {b l o c k d i a g} (\boldsymbol {v}) = \boldsymbol {M} \quad \text {w h e r e} \quad M _ {i j} = \left\{ \begin{array}{l l} v _ {j} & g (j) = i \\ 0 & \text {o t h e r w i s e .} \end{array} \right.
$$
Likewise, we define the projection operator blockproj:  $\mathcal{V}_{q,N} \mapsto \mathbb{R}^N$  which extracts a vector of size  $N$  from a  $N \times q$  according to the block structure of  $[N]$  by
$$
\operatorname {b l o c k p r o j} (\boldsymbol {M}) = (M _ {i g (i)}) _ {i \leq N} \in \mathbb {R} ^ {N}.
$$
Under these changes of variables, we define
$$
\boldsymbol {r} ^ {t} = \operatorname {b l o c k d i a g} \left(\boldsymbol {x} ^ {t}\right) \in \mathcal {V} _ {q, N} \quad \text {f o r} t \geq 0
$$
and  $\tilde{f}_t\colon (\mathbb{R}^q)^N\mapsto (\mathbb{R}^q)^N$  by
$$
\left(\tilde {f} _ {t} \left(\boldsymbol {r} ^ {t}\right)\right) _ {i j} = \frac {1}{\sqrt {\tilde {\Delta} _ {g (i) j}}} f _ {t} ^ {g (i)} \left(x _ {i} ^ {t}\right) \quad \text {f o r} i, j \in [ N ] \times [ q ]. \tag {15}
$$
We encode the family of functions  $\tilde{f}_t$  by  $\tilde{\mathcal{F}} (\Delta)$
Lemma 2.3. Let  $\pmb{x}^t$  be iterates from the AMP  $(\mathbf{A},\mathcal{F},\pmb{v}^0,\pmb{\Delta})$ . Then the iterates  $\pmb{r}^t \coloneqq \mathrm{blockdiag}(\pmb{x}^t)$  follow the generalized matrix AMP  $(\mathbf{A},\tilde{\mathcal{F}} (\Delta),\pmb{r}^0)$ .
Proof. We will show that the projection of the iterates  $\boldsymbol{r}^t$  from  $(\mathbf{A}, \tilde{\mathcal{F}}(\Delta), r^0)$  are the iterates from  $(\mathbf{A}, \mathcal{F}, \boldsymbol{v}^0, \Delta)$ . It is easy to check that
$$
\frac {\boldsymbol {A}}{\sqrt {N} \sqrt {\boldsymbol {\Delta}}} f _ {t} \left(\boldsymbol {x} ^ {t}\right) = \operatorname {b l o c k p r o j} \left(\frac {\boldsymbol {A}}{\sqrt {N}} \tilde {f} _ {t} \left(\boldsymbol {r} ^ {t}\right)\right). \tag {16}
$$
Next, notice that Jacobian is given by
$$
\partial \tilde {f} _ {t} ^ {i} (\boldsymbol {r} _ {i} ^ {t}) = \left[ \begin{array}{c c c c c} 0 & \dots & \frac {1}{\sqrt {\Delta_ {g (i) 1}}} (f _ {t} ^ {g (i)}) ^ {\prime} (x _ {i} ^ {t}) & \dots & 0 \\ \vdots & \ddots & \vdots & \ddots & \vdots \\ 0 & \dots & \frac {1}{\sqrt {\Delta_ {g (i) q}}} (f _ {t} ^ {g (i)}) ^ {\prime} (x _ {i} ^ {t}) & \dots & 0 \end{array} \right],
$$
which is a matrix where only the column number  $g(i)$  has non-zero elements. Applying (11), we thus get for  $a, b \in [q] \times [q]$
$$
\left(\boldsymbol {B} _ {t}\right) _ {a b} = \frac {1}{N \sqrt {\tilde {\Delta} _ {a b}}} \sum_ {i: g (i) = a} \left(f _ {t} ^ {a}\right) ^ {\prime} \left(x _ {i} ^ {t}\right). \tag {17}
$$
$$
\mathrm {b} _ {t} \odot f _ {t - 1} (\boldsymbol {x} ^ {t - 1}) = \operatorname {b l o c k p r o j} \left(\tilde {f} _ {t - 1} (\boldsymbol {r} ^ {t - 1}) \boldsymbol {B} _ {t} ^ {T}\right).
$$
As a consequence, the inhomogeneous state evolution equations in Theorem 1.2 follow immediately from the state evolution equations of  $(\mathbf{A},\tilde{\mathcal{F}} (\Delta),\boldsymbol{r}^{0})$  discussed in [Javanmard and Montanari, 2013, Section 2.1]. This follows from the observation that given the law of  $\pmb{r}^t$  in the high dimensional limit, the law of  $\pmb{x}^t = \mathrm{blockproj}(\pmb{r}^t)$  is straightforward to compute. We define
$$
\left(\sigma_ {b} ^ {t + 1}\right) ^ {2} := \sum_ {a = 1} ^ {q} \frac {c _ {a}}{\tilde {\Delta} _ {a b}} \mathbb {E} \left[ \left(f _ {t} ^ {b} \left(Z _ {b} ^ {t}\right)\right) ^ {2} \right]. \tag {18}
$$
We will show that the distribution of the iterate  $\boldsymbol{x}_i^t$  is asymptotically normal with mean 0 and variance  $(\sigma_{g(i)}^t)^2$ .