## Quantitative Group Testing and Pooled Data in the Linear Regime with Sublinear Tests

Nelvin Tan, Pablo Pascual Cobo, and Ramji Venkataramanan

## Abstract

In the pooled data problem, the goal is to identify the categories associated with a large collection of items via a sequence of pooled tests. Each pooled test reveals the number of items in the pool belonging to each category. A prominent special case is quantitative group testing (QGT), which is the case of pooled data with two categories. We consider these problems in the non-adaptive and linear regime, where the fraction of items in each category is of constant order. We propose a scheme with a spatially coupled Bernoulli test matrix and an efficient approximate message passing (AMP) algorithm for recovery. We rigorously characterize its asymptotic performance in both the noiseless and noisy settings, and prove that in the noiseless case, the AMP algorithm achieves almost-exact recovery with a number of tests sublinear in the number of items. For both QGT and pooled data, this is the first efficient scheme that provably achieves recovery in the linear regime with a sublinear number of tests, with performance degrading gracefully in the presence of noise. Numerical simulations illustrate the benefits of the spatially coupled scheme at finite dimensions, showing that it outperforms i.i.d. test designs as well as other recovery algorithms based on convex programming.

## 1 Introduction

Consider a large collection of items, each of which is either defective or non-defective. In group testing [1], the goal is to identify the defective set via pooled tests, where groups of items are tested together, with as few tests as possible. In the original Boolean group testing model, which has been studied extensively [1-5], each test returns a positive outcome if it includes at least one defective item and a negative outcome otherwise. Its variant, the quantitative group testing (QGT) model [6], is useful when tests are more informative: each test reveals the number of defective items in that pool. QGT is of interest in a range of modern applications, including genomics [7], multi-access communication [8], and network traffic monitoring [9]. A more general version of QGT, where each item belongs to one of L &gt; 2 categories, is known as the pooled data problem [10]. The goal is to identify the categories via a sequence of pooled tests, where each pooled test reveals the number of items of each category within the pool.

In this paper, we consider non-adaptive QGT and pooled data, where the tests are all designed in advance, making them amenable to being implemented in parallel. We also consider the linear regime, which for QGT, means that the number of defective items is proportional to the total number

N. Tan was supported by a Cambridge Trust Scholarship and the Harding Distinguished Postgraduate Scholars Programme Leverage Scheme. P. Pascual Cobo was supported by a Engineering and Physical Sciences Research Council Doctoral Training Award. This work was presented in part at the 2024 International Zurich Seminar on Information and Communication (IZS). The authors are with the Department of Engineering, University of Cambridge, UK. Emails: tcnt2@cam.ac.uk, pp423@cam.ac.uk, rv285@cam.ac.uk .

of items. For pooled data, the linear regime implies that the proportion of items in each category is non-vanishing as the number of items increases, a realistic assumption in practical applications.

## 1.1 Problem Setup

Quantitative group testing. There are p items, whose status is denoted by the binary vector β ∈ { 0 1 , } p , where 1 represents a defective item and 0 a non-defective item. Items are allocated to tests using a binary design (or test) matrix X ∈ { 0 1 , } n × p , where n is the number of tests and p is the number of items. The i th row X i, : determines the pooling design of the i th test, where X ij = 1 indicates that the j th item will be included in the i th test, and X ij = 0 indicates otherwise. Let d be the number of defective items with d &lt; p . We consider the linear regime, where the fraction of defective items d/p converges to π ∈ (0 1) , . Mathematically, we define the QGT model as

$$y _ { i } = \left ( X _ { i, \colon } \right ) ^ { \top } \beta + \Psi _ { i } \quad \text{for $i\in\{1,\dots,n\}$},$$

where y i is the i th element of y ∈ R n , X i, : is the i th row of X ∈ R n × p represented as a column vector, and Ψ i is the i th element of the additive noise Ψ ∈ R n . Under the noiseless setting (i.e., all entries of Ψ are zero), the output y i is the number of defective items in the i th test. The goal of QGT is to recover β with as few tests as possible.

We will use the almost-exact recovery criterion, which is achieved by an estimator ˜ β if

̸

$$\frac { 1 } { p } \sum _ { j = 1 } ^ { p } 1 \{ \tilde { \beta } _ { j } \neq \beta _ { j } \} \to 0 \text{ as } p \to \infty.$$

̸

This is a weaker notion of recovery as compared to the exact recovery criterion [11] where we want the probability of error P [ ˜ = β β ] → 0 as p →∞ . We note that an almost-exact recovery criterion is meaningful in the linear regime, but not in the sublinear regime where the number of defectives d = ( ) o p , since setting ˜ β to be the zero vector would trivially satisfy (2).

Pooled data. The signal to be estimated is a matrix B ∈ { 0 1 , } p × L , where each row is a onehot vector. For example B j, : = [0 1 0 , , , . . . , 0] represents the j th item belonging to category 2 (the position of one in B j, : ). We consider the linear regime, where the fraction of items in each category l converges to π l ∈ (0 1) , , where ∑ L l =1 π l = 1 . The model is

$$Y _ { i, \colon } = B ^ { \top } X _ { i, \colon } + \Psi _ { i, \colon } \in \mathbb { R } ^ { L } \quad \text{for $i\in\{1,\dots,n\}$},$$

where Y i, : is the i th row of Y ∈ R n × L represented as a column vector, and Ψ i, : is the i th row of the additive noise Ψ ∈ R n × L represented as a column vector. Under the noiseless setting (i.e., all entries of Ψ are zero), the output of each test Y i, : tells us the number of items from each category present in the test, which can be viewed as a histogram. Similar to QGT, an estimator ˜ B achieves almost-exact recovery if

̸

$$\frac { 1 } { p L } \sum _ { j = 1 } ^ { p } \sum _ { l = 1 } ^ { L } { \mathbf 1 } \{ \widetilde { B } _ { j l } \neq B _ { j l } \} \to 0 \text{ as } p \to \infty.$$

The number of categories L does not grow with p .

Information-theoretic limits. For noiseless pooled data in the linear regime, the informationtheoretic limit on the number of tests required was established by Scarlett and Cevher [11], closing the gap between previously derived upper and lower bounds [10,12,13]. It was shown in [11] that the minimum number of tests needed for exact recovery is n ∗ = γ ∗ p log p (1 + o (1)) , where

$$\gamma ^ { * } = \max _ { r \in \{ 1, \dots, L - 1 \} } \frac { 2 [ H ( \pi ) - H ( \pi ^ { ( r ) } ) ] } { L - r },$$

while H π ( ) = -∑ L l =1 π l log π l is the Shannon entropy function, and π ( r ) = ( π ( r ) 1 , . . . , π ( r ) r ) is a vector whose first entry sums the largest ( L -r + 1) entries of π , and whose remaining entries coincide with the remaining ( r -1) entries of π . Setting L = 2 above gives the information-theoretic limit for noiseless QGT.

For noisy pooled data where the entries of the noise matrix are independent with zero mean, it was shown in [11] that we require n = Ω( p log p ) tests for exact recovery, in contrast with the sublinear Θ ( p log p ) required in the noiseless case.

Efficient Algorithms. For pooled data, Wang et al. [10] proposed a deterministic design matrix and a polynomial-time algorithm that achieves exact recovery with n = Ω ( p log p ) tests, matching the optimal sample complexity above. However, both the test design and the recovery algorithm (based on Gaussian elimination) are tailored to the noiseless setting, and do not extend to the noisy case. In this paper, we focus on random test designs, which are more robust with respect to the items included in each test, and on recovery algorithms whose performance degrades gracefully with the noise level. For random i.i.d. designs (where X ij i . i . d . ∼ Bernoulli ( α ) for some α ∈ (0 1) , ), efficient recovery using Approximate Message Passing (AMP) algorithms was studied in [13, 14]. Rigorous guarantees on the recovery performance of AMP, established in [14], imply that with the i.i.d. Bernoulli design, the AMP algorithm needs n = Θ( ) p tests for almost-exact recovery, even in the noiseless QGT/pooled data setting. Thus, for random designs there is an orderlog p gap between the information-theoretic limit and the best known algorithms.

## 1.2 Approximate Message Passing and Spatial Coupling

Approximate message passing (AMP) is a family of iterative algorithms that can be tailored to take advantage of structural information about the signals and the model, e.g., a known prior on the signal vector or on the proportion of observations that come from each signal. AMP algorithms were first proposed for the standard linear model [15-18], but have since been applied to a range of statistical problems, including estimation in generalized linear models and their variants [19-23], as well as low-rank matrix and tensor estimation [24-30]. In all these settings, under suitable model assumptions the performance of AMP in the high-dimensional limit is characterized by a succinct deterministic recursion called state evolution . The literature on AMP is vast, and we refer the interested reader to [31] for a survey.

In this paper, we will use a spatially coupled design matrix X and a suitable AMP algorithm for recovery. The spatially coupled matrix has a block-wise structure, with blocks along a band-diagonal having i.i.d. Bernoulli entries and the remaining blocks being all zeros (see Figure 1). Our scheme is inspired by a line of work on compressed sensing with spatially coupled designs [18, 32-34]. For a noiseless linear model defined via a spatially coupled Gaussian sensing matrix, Donoho et al. [34] showed that AMP recovers the signal with high probability when the sampling ratio δ = n/p

exceeds the Rényi information dimension of the signal prior. The Rényi information dimension is zero for priors supported on a finite set, which implies that AMP can recover the signal with n = o p ( ) measurements. Recently, it was also shown that using a spatially coupled sensing matrix in a generalized linear model allows AMP to achieve the Bayes-optimal error (corresponding to an i.i.d. Gaussian matrix) [35].

## 1.3 Main Contributions

In Section 3, we describe the spatially coupled random test design and an AMP algorithm (SCAMP) for signal recovery. In Theorem 3.2 we give a precise characterization of the performance of SC-AMP in the asymptotic regime where the number of tests n grows proportionally with the number of items p (with n/p → δ , a constant). Using this characterization, we bound the MSE of the algorithm in the low noise regime (Theorem 3.7) and show that for noiseless QGT, the SCAMP algorithm achieves almost-exact recovery with probability one, for any sampling ratio δ &gt; 0 (Corollary 3.8). This implies that it achieves almost-exact recovery with n = ( ) o p tests. In Section 4, we generalize the SC-AMP algorithm to the pooled data setting and again establish almost-exact recovery for any δ &gt; 0 (Theorem 4.1).

At the heart of our theoretical guarantees is a rigorous analysis of an AMP algorithm for a generalized linear model (GLM) with a generic spatially coupled design matrix. (The matrix consists of blocks of independent entries drawn from an arbitrary zero-mean distribution satisfying certain moment conditions.) GLMs include many important nonlinear estimation problems such as phase retrieval and logistic regression. Theorem 5.2 shows that the AMP algorithm and its performance characterization originally developed for GLMs with spatially coupled Gaussian designs [35] remain valid for a much broader class of designs.

To our knowledge, our scheme is the first one based on a random design to efficiently achieve recovery for QGT and pooled data in the linear regime with a sublinear number of tests. Moreover, the scheme is robust to noise. Numerical simulations show that the spatially coupled scheme outperforms the i.i.d. Bernoulli test design with AMP, as well as recovery algorithms based on convex programming.

Key technical ideas. Although the QGT model (1) is an instance of a linear model, an important constraint is that the test design matrix X can only contain binary entries. Therefore, we cannot apply the analysis from [34], which assumes a spatially coupled Gaussian design matrix. To prove the theoretical guarantees for our scheme, we reduce the SC-AMP algorithm to an abstract AMP iteration defined for any generalized white noise matrix. A state evolution result for this abstract AMP iteration was established by Wang et al. in [36], using which we obtain a rigorous asymptotic characterization of the SC-AMP algorithm (Theorem 3.2). To establish conditions for almost-exact recovery, we then need to analyze the fixed points of the SC-AMP state evolution. We do this in Theorem 3.5 via the potential function method [37], a powerful tool for characterizing the fixed points of coupled recursions. This characterization then yields the noise robustness and exactrecovery results (Theorem 3.7 and Corollary 3.8).

## 1.4 Other Related Work

Spatial coupling. Spatial coupling was introduced in coding theory as a means to construct LDPC codes that achieve capacity with an efficient belief propagation decoder [38, 39]. Spatial

coupling has since been applied in many estimation problems to improve on the performance of 'regular' (or i.i.d.) designs. For Boolean group testing in the sublinear regime (the number of defectives is p θ for θ ∈ (0 , 1) ), spatially coupled test designs enable efficient recovery with the asymptotically optimal number of tests, in both the noiseless [40] and noisy settings [41]. For QGT, Mashauri et al. [42, 43] investigated efficient schemes based on spatial coupled LDPC codes, and showed that they outperform previous constructions based on generalized LDPC codes [44,45].

Sublinear category regime. A few recent works have studied pooled data in the sublinear category regime, where one category is dominant with p -o p ( ) items, and the remaining ( L -1) categories have d = o p ( ) items. (In contrast, we consider the linear category regime, where the proportion of items in each category is Θ(1) , i.e., π l = Θ(1) for l ∈ [ L ] .) For the sublinear category regime, the information-theoretic lower bound for exact recovery is n = Ω( ) d tests [46, 47]. An efficient algorithm proposed in [46] achieves the lower bound when d = Θ( p κ ) , for a constant κ ∈ (0 , 1) . A lower complexity algorithm for the special case of QGT with d = Θ( p κ ) was recently proposed in [48]. For QGT in the sublinear regime, a number of recent works have proposed algorithms based on ideas from coding theory [42] and thresholding [49], which require Ω( d log p ) tests for exact recovery. Noisy versions of QGT were recently studied in [50] and [51], and QGT in the adaptive setting has been studied in [6,52].

## 2 Preliminaries

Notation. We write [ n ] := { 1 , . . . , n } and [ n : m ] for { n, n + 1 , . . . , m } where n &lt; m . All vectors (including those corresponding to rows of matrices) are assumed to be column vectors unless otherwise stated. For a, b ∈ R n , ⟨ a, b ⟩ = a ⊤ b ∈ R is the inner product, a ⊙ b ∈ R n is the entry-wise product, and ⟨ a ⟩ = 1 n ∑ n i =1 a i denotes the empirical average of the entries of a . Matrices are denoted by upper case letters, and given a matrix A , we write A i, : for its i th row and A : ,j for its j th column. The operator norm is denoted by ∥ A ∥ op . For r ∈ [1 , ∞ ) and a vector a = ( a , . . . , a 1 n ) ∈ R n , we write ∥ a ∥ r for the ℓ r -norm, so that ∥ a ∥ r = ( ∑ n i =1 | a i | r ) 1 /r . We use 1 p to denote the vector of p ones, 0 p for the vector of p zeros, and I p for the p × p identity matrix. Given random variables U, V , we write U d = V to denote equality in distribution. We write ∂ f i ( ) · to denote the partial derivative of f with respect to (w.r.t.) the i th argument. Throughout, the function log( ) · has base e , and we use Bachmann-Landau asymptotic notation (i.e., O o , , Ω , ω , Θ ).

Almost-sure and Wasserstein convergence. Let { A n } be a sequence of random elements taking values in a Euclidean space E . We say that A n converges almost surely to a deterministic limit a ∈ E , and write A n a.s. → a , if P [lim n →∞ A n = a ] = 1 .

For a vector a ∈ R n and a random variable A ∈ R , we write a W r → A as n → ∞ , for the Wassersteinr convergence of the empirical distribution of the entries of a to the law of A . More generally, for vectors a , . . . , a 1 k ∈ R n and a random vector ( A , . . . , A 1 k ) ∈ R k , we write

$$a ^ { 1 }, \dots, a ^ { k } \stackrel { W _ { \bar { \gamma } } } { \rightarrow } ( A ^ { 1 }, \dots, A ^ { k } ) \text{ as } n \rightarrow \infty,$$

for the Wassersteinr convergence of the empirical distribution of rows of ( a , . . . , a 1 k ) ∈ R n × k to the joint law of ( A , . . . , A 1 k ) . This means that, for any continuous function ϕ : R k → R and input

Figure 1: The entries of X sc are independent with X sc ij ∼ Bernoulli ( αW r ( ) i , c ( j ) ) . Here W is an ( ω, Λ) base matrix with ω = 3 and Λ = 7 (see Definition 3.1). The white parts of X sc and W correspond to zeros.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-bfb3c42b-acc7-44f0-8ad5-c1c12fd6e2e2.jpg)

vector ( a , . . . , a 1 i k i ) ∈ R k satisfying the polynomial growth condition [36]

$$| \phi ( a _ { i } ^ { 1 }, \dots, a _ { i } ^ { k } ) | \leq C ( 1 + \| ( a _ { i } ^ { 1 }, \dots, a _ { i } ^ { k } ) \| _ { 2 } ^ { r } ),$$

for a constant C &gt; 0 , we have

$$\frac { 1 } { n } \sum _ { i = 1 } ^ { n } \phi ( a _ { i } ^ { 1 }, \dots, a _ { i } ^ { k } ) \to \mathbb { E } \left [ \phi ( A ^ { 1 }, \dots, A ^ { k } ) \right ] \text{ as } n \to \infty.$$

We write

$$a \stackrel { W } { \to } A ; \ \ ( a _ { 1 }, \dots, a _ { k } ) \stackrel { W } { \to } ( A _ { 1 }, \dots, A _ { k } ) \text{ as } n \to \infty,$$

to mean that the above Wassersteinr convergences hold for every order r ≥ 1 .

Model assumptions for QGT. The signal β ∈ { 0 1 , } p is independent of the design matrix. As n, p →∞ , we have n/p → δ &gt; 0 , and the empirical distribution of the entries of the signal converges in Wasserstein distance to well-defined limits. More precisely, β W → ¯ β where ¯ β ∼ Bernoulli ( π ) . We note that the entries of β are not assumed to be independent or identically distributed.

## 3 Spatially Coupled Design for Quantitative Group Testing

The spatially coupled (SC) design matrix consists of independent Bernoulli entries whose parameters are specified by a base matrix W of dimension R × C . The SC design matrix is obtained by replacing each entry of the base matrix W rc by an n R × p C matrix with entries drawn independently from Bernoulli ( αW rc ) , where α, αW rc ∈ (0 , 1) . An example of a SC design matrix is shown in Figure 1. In this paper, we will use the following base matrix.

Definition 3.1. An ( ω, Λ) base matrix W is described by two parameters: the coupling width ω ≥ 1 and the coupling length Λ ≥ 2 ω -1 . The matrix has R = Λ+ ω -1 rows and C = Λ columns, with each entry indexed by ( r , c ) , for r ∈ [ R ] and c ∈ [ C ] . For α ≤ 0 5 . , the entries are given by

$$W _ { r c } = \begin{cases} \frac { 1 } { 2 \alpha } \left ( 1 - \sqrt { 1 - \frac { 4 \alpha ( 1 - \alpha ) } { \omega } } \right ) & \text{if $c\leq r\leq c+\omega-1$} \\ 0 & \text{otherwise.} \end{cases}$$

$$\text{For } \alpha > 0. 5, \text{ the non-zero entries } ( \text{when } c \leq r \leq c + \omega - 1 ) \text{ are given by } \frac { 1 } { 2 \alpha } \left ( 1 + \sqrt { 1 - \frac { 4 \alpha ( 1 - \alpha ) } { \omega } } \right ).$$

Figure 1 shows an ( ω, Λ) base matrix with ω = 3 and Λ = 7 . The spatially coupled (SC) design matrix, denoted by X sc , has entries

$$X _ { i j } ^ { s c } \underset { \sim } { \text{emodii} } \text{ Bernoulli} ( \alpha W _ { r ( i ) \epsilon ( j ) } ), \ \ i \in [ n ], \, j \in [ p ].$$

for some fixed constant α ∈ (0 , 1) . Here the operators r ( ) · : [ n ] → [ R ] and c ( ) · : [ p ] → [ C ] map a particular row or column index to its corresponding row block or column block index in W . The band-diagonal structure of ( ω, Λ) base matrix here is similar to the ones used for SC sparse regression codes [53] and for SC generalized linear models [35], but the values of the non-zero entries are different. Here the base matrix specifies the Bernoulli parameters for each block of the design, whereas in [35,53] it specifies the variances for the Gaussian entries in each block.

With the i.i.d. design, the number of defectives per test has expected value απp and standard deviation √ απ (1 -απ p ) . Similarly, for the spatially coupled design, it can be verified that the number of defective items per test has mean that is linear in p and standard deviation of order √ p . Since the fluctuation around the mean contains the useful information in each test, and because AMP requires a design matrix with zero-mean entries, we recenter and rescale the data before applying the AMP algorithm. We now describe this preprocessing of the data, which was also done in [14] for i.i.d. designs.

The i.i.d. design matrix, denoted by X iid , has each entry sampled i.i.d. ∼ Bernoulli ( α ) , for some fixed constant α ∈ (0 , 1) . Note that the i.i.d. matrix is a special case of the SC matrix, with R = C = 1 and W = 1 . A key difference between the i.i.d. design and the SC design (with an ( ω, Λ) base matrix) is that the latter includes many fewer items in each test. Indeed, with the i.i.d. design each item is included in a test with probability α , whereas in the SC design, each test includes items from at most ω adjacent column blocks, each with p/ C items (see Figure 1). In the SC design, since the tests corresponding to the first and last row blocks involve the fewest items, the corresponding entries of β are easier to recover than the others. A good estimate for these entries helps the algorithm recover the entries in the adjacent blocks, creating a decoding wave that propagates from the ends towards the center.

The rescaled i.i.d. matrix, denoted by ˜ X iid , is defined as

$$\widetilde { X } ^ { \text{iid} } = \frac { X ^ { \text{iid} } - \alpha 1 _ { n } 1 _ { p } ^ { \top } } { \sqrt { n \alpha ( 1 - \alpha ) } }.$$

We note that ˜ X iid has independent entries with E [ ˜ X ij ] = 0 and Var[ ˜ X ij ] = 1 /n .

The rescaled spatially coupled (SC) matrix ˜ X sc is defined as follows. For i ∈ [ n , j ] ∈ [ p ] , using the shorthand r ≡ r ( ) i , c ≡ c ( j ) , its entries are given by

$$\widetilde { X } _ { i j } ^ { s c } = \frac { X _ { i j } ^ { s c } - \alpha W _ { r c } } { \sqrt { n \alpha ( 1 - \alpha ) / R } } = \begin{cases} \frac { 1 - \alpha W _ { r c } } { \sqrt { n \alpha ( 1 - \alpha ) / R } } & \text{with probability $\alpha W_{rc}$,} \\ \frac { - \alpha W _ { r c } } { \sqrt { n \alpha ( 1 - \alpha ) / R } } & \text{with probability $1-\alpha W_{rc}$.} \end{cases}$$

It is straightforward to verify that E [ ˜ X sc ij ] = 0 , Var[ ˜ X sc ij ] = R W rc (1 -αW rc ) n (1 -α ) . In particular, for the ( ω, Λ) base matrix in Definition 3.1, we have

$$\text{Var} [ \widetilde { X } _ { i j } ^ { s c } ] = \begin{cases} \frac { R } { n \omega } & \text{if $c\leq r\leq c+\omega-1$} \\ 0 & \text{otherwise.} \end{cases}$$

Rewriting QGT. The AMP algorithm and its analysis require the design matrix to have independent zero-mean entries, so we recenter and rescale the QGT model in (1) to express it in terms of the rescaled design ( ˜ X sc or ˜ X iid ). For the SC design, using (10), we have for i ∈ [ n ] :

$$\overset { \cdot } { \text{of the rescaled design } ( \widetilde { X } ^ { s } \text{$c$} \text{ or $\widetilde{X}^{iid}$} ). \text{ For the SC design, using } (10), we have for $i\in [n] \colon \\ y _ { i } = \overset { p } { \underset { j = 1 } { \underset { k = 1 } { \sum } } } X ^ { s c } _ { i j } \beta _ { j } + \Psi _ { i } = \overset { p } { \underset { j = 1 } { \sum } } \left ( \alpha W _ { r ( i ) c ( j ) } + \sqrt { \frac { n \alpha ( 1 - \alpha ) } { R } } \widetilde { X } ^ { s c } _ { i j } \right ) \beta _ { j } + \Psi _ { i } \\ \Longrightarrow & y _ { i } - \overset { p } { \underset { j = 1 } { \sum } } \alpha W _ { r ( i ) c ( j ) } \beta _ { j } = \sqrt { \frac { n \alpha ( 1 - \alpha ) } { R } } \overset { p } { \underset { j = 1 } { \sum } } \widetilde { X } ^ { s c } _ { i j } \beta _ { j } + \Psi _ { i } \\ \Longrightarrow & \frac { 1 } { \sqrt { n \alpha ( 1 - \alpha ) / R } } \left ( y _ { i } - \alpha \left ( W _ { r ( i ) 1 } \underset { j \in \mathcal { J } _ { 1 } } { \underset { k = 1 } { \sum } } \beta _ { j } + \cdots + W _ { r ( i ) c } \sum _ { j \in \mathcal { J } _ { c } } \beta _ { j } \right ) \right ) = ( \widetilde { X } ^ { s c } _ { i, i } ) ^ { \top } \beta + \frac { \Psi _ { i } } { \sqrt { n \alpha ( 1 - \alpha ) / R } }, \\ \text{where $\mathcal{J}_{c}=[(c-1)p/C+1\coloncp/C]$ for $c\in[C]$. Denoting the left-hand side above by}$$

where J c = ( [ c -1) p/ C +1 : c p/ C ] for c ∈ [ C ] . Denoting the left-hand side above by

$$\dot { \vec { y } } _ { i } \coloneqq \frac { 1 } { \sqrt { n \alpha ( 1 - \alpha ) / R } } \Big ( y _ { i } - \alpha \left ( W _ { r ( i ) 1 } \sum _ { j \in \mathcal { J } _ { 1 } } \beta _ { j } + \cdots + W _ { r ( i ) C } \sum _ { j \in \mathcal { J } _ { c } } \beta _ { j } \right ) \right ),$$

gives us the rescaled QGT model:

$$\tilde { y } _ { i } = ( \widetilde { X } _ { i, ; } ^ { \text{sc} } ) ^ { \text{T} } \beta + \widetilde { \Psi } _ { i }, \quad \text{with} \quad \widetilde { \Psi } _ { i } \coloneqq \frac { \Psi _ { i } } { \sqrt { n \alpha ( 1 - \alpha ) / R } } \,, \quad i \in [ n ].$$

The term ∑ j ∈J c β j in (12) is the number of defective items in the sub-vector of β indexed by J c , for c ∈ [ C ] . In the noiseless setting, the terms ∑ j ∈J c β j can be obtained with an extra C tests, where we only include items from J c in the c th test, for c ∈ [ C ] . The extra C = O (1) tests does not affect our results since the limiting sampling ratio lim n →∞ n/p = δ remains the same. Since 1 p ∑ j ∈J c β j a.s. → π C for p →∞ via the strong law of large numbers, so we can also estimate ∑ j ∈J c β j using p C π . (However the error in this estimate would be of order √ p/ C .)

For an i.i.d. design X iid , we similarly recenter and rescale the QGT model to express it in terms of the rescaled i.i.d. matrix ˜ X iid in (9). For i ∈ [ n ] , we have

$$y _ { i } = \left ( X _ { i, i } ^ { i \text{id} } \right ) ^ { \top } \beta + \Psi _ { i } \implies & \tilde { y } _ { i } = \left ( \widetilde { X } _ { i, i } ^ { \text{id} } \right ) ^ { \top } \beta + \widetilde { \Psi } _ { i }, \\ & \quad \text{where } \tilde { y } _ { i } \coloneqq \frac { y _ { i } - \alpha d } { \sqrt { n \alpha ( 1 - \alpha ) } }, \quad \widetilde { \Psi } _ { i } = \frac { \Psi _ { i } } { \sqrt { n \alpha ( 1 - \alpha ) } }, \quad \quad ( 1 4 )$$

where d is the number of defective items.

$$\dots$$

Choice of α . For a spatially coupled design X sc constructed from an ( ω, Λ) base matrix, recall from (10) that the rescaled matrix ˜ X sc has independent zero-mean Bernoulli entries with variances given by (11). Notice that the distribution of ˜ X sc does not depend on α . Hence, in the rescaled model, α only affects the variance of the noise ˜ Ψ , which is minimized when α = 0 5 . . We therefore use α = 0 5 . for all our experiments.

Taking α = 0 5 . , it is useful to compare the SC design with the i.i.d. Bernoulli( 0 5 . ) design. (For the i.i.d. design, taking the Bernoulli parameter to be 0 5 . is optimal with respect to both the information-theoretic limits [11] and efficient recovery via AMP [14].) Taking α = 0 5 . in Definition 3.1, we have that the non-zero blocks of the SC design are drawn independently from an Bernoulli distribution with parameter 1 2 ( 1 -√ 1 -1 /ω ) . From structure of X sc (see Figure 1), it follows that the expected number of items included in each test is at most p ω C 2 ( 1 -√ 1 -1 /ω ) . For large ω , this is approximately p 4 C items per test. In contrast, for the Bernoulli( 0 5 . ) design, the expected number of items per test is p 2 . Thus, for large C , the expected number of items per test is much smaller for the SC design than the i.i.d. design.

Noise scaling assumption. In the rescaled QGT models (13) and (14), it can be verified that the terms ( ˜ X sc i, : ) ⊤ β and ( ˜ X iid i, : ) ⊤ β have zero mean and variance of constant order, for each i ∈ [ n ] . Therefore, for the rescaled model to be meaningful, the noise ˜ Ψ i should also have a mean and variance of constant order. This is guaranteed by the following assumption. The empirical distribution of rescaled noise vector ˜ Ψ in (13) converges to a well defined limit. More precisely, there exists, ¯ Ψ ∼ P ¯ Ψ with E [ ¯ Ψ 2 ] =: σ 2 &lt; ∞ , such that ˜ Ψ W → ¯ Ψ as n →∞ . We emphasize that the base matrix parameter R is fixed as n →∞ . A similar distributional assumption holds for the rescaled noise vector with the i.i.d. design in (14).

## 3.1 SC-AMP Algorithm

Consider the rescaled SC model (13). Given ( ˜ X ,y sc ˜) , the SC-AMP algorithm iteratively produces estimates ˆ β k for k ≥ 1 . For iteration k ≥ 0 , the algorithm computes:

$$\widetilde { \Theta } ^ { k } & = \tilde { y } \, - \, \widetilde { X } ^ { s c } \hat { \beta } ^ { k } \, + \, b ^ { k } \odot Q ^ { k - 1 } \odot \widetilde { \Theta } ^ { k - 1 } \, \in \, \mathbb { R } ^ { n } \\ \beta ^ { k + 1 } & = ( \widetilde { X } ^ { s c } ) ^ { \top } ( Q ^ { k } \odot \widetilde { \Theta } ^ { k } ) \, - \, c ^ { k } \odot \hat { \beta } ^ { k } \, \in \, \mathbb { R } ^ { p }, \quad \hat { \beta } ^ { k + 1 } = f _ { k + 1 } ( \beta ^ { k + 1 }, \mathcal { C } ) \, \in \, \mathbb { R } ^ { p },$$

where f k : R 2 → R acts row-wise on the inputs ( β , k C ) and

$$\mathcal { C } \coloneqq ( \underbrace { 1, \dots, 1 } _ { p / C \text{ entries } p / C \text{ entries } } \,, \, \underbrace { 2, \dots, 2 } _ { p / C \text{ entries } } \,, \, \dots, \, \underbrace { C, \dots, C } _ { p / C \text{ entries } } ) ^ { \top } \in \mathbb { R } ^ { p }.$$

The algorithm is initialized with ˆ β 0 = E [ ¯ ]1 β p and ˜ Θ = ˜ 0 y -˜ X β sc ˆ 0 . We recall that ⊙ is the Hadamard (entry-wise) product. To define b k , we use the partitions [ p ] = ⋃ C c =1 J c and [ n ] = ⋃ R r =1 I r where

$$\mathcal { J } _ { c } = \left \{ ( c - 1 ) \frac { p } { C } + 1, \dots, c \frac { p } { C } \right \} \text{ for } c \in [ C ], \quad \mathcal { I } _ { \tau } = \left \{ ( r - 1 ) \frac { n } { R } + 1, \dots, r \frac { n } { R } \right \} \text{ for } r \in [ R ]. \quad ( 1 7 )$$

Then, letting

$$\widetilde { W } _ { \text{rc} } \coloneqq \frac { W _ { \text{rc} } ( 1 - \alpha W _ { \text{rc} } ) } { 1 - \alpha } = \begin{cases} 1 / \omega & \text{if $c\leq r\leq c+\omega-1$} \\ 0 & \text{otherwise.} \end{cases},$$

the entries of b k ∈ R n are

$$b _ { i } ^ { k } = \sum _ { c = 1 } ^ { c } \frac { \widetilde { W } _ { r c } } { n / R } \sum _ { j \in \mathcal { J } _ { c } } \partial _ { 1 } f _ { k } ( \beta _ { j } ^ { k }, \mathbf c ), \ \ i \in [ n ].$$

The second equality in (18) follows from Definition 3.1. The function f k and the vectors Q ,c k k in (15) will be defined later in (19) and (23), respectively, in terms of state evolution parameters. We note that the time complexity of each iteration in (15) is O np ( ) .

We introduce some additional notation to handle the block-wise structure of the SC-AMP iterates. For c ∈ [ C ] and r ∈ [ R ] , define β c := β J c ∈ R p/ C and ˆ β k c := ˆ β k J c ∈ R p/ C to be the c th blocks of β ∈ R p and ˆ β k ∈ R p respectively, and Θ := Θ r I r ∈ R n/ R and ˜ Θ k r := Θ ˜ k I r ∈ R n/ R to be the r th blocks of Θ := ˜ X β sc ∈ R n and ˜ Θ k ∈ R n respectively. Similar notation simplifications will be used for other vectors where c and r will replace J c and I r in the subscripts of the vectors.

State evolution. The 'memory' terms b k ⊙ Q k -1 ⊙ ˜ Θ k -1 and -c k ⊙ ˆ β k in (15) debias the iterates ˜ Θ k and β k +1 , ensuring that their empirical distributions are accurately captured by state evolution in the high-dimensional limit. These iterates have a block-wise distributional structure. Recall from the model assumptions that the empirical distribution of the signal β converges to the law of ¯ β ∼ Bernoulli ( π ) . Theorem 3.2 below shows that for each k ≥ 1 and c ∈ [ C ] , the empirical distribution of β k c converges to the distribution of ( χ k c ) 2 ¯ + β χ G k c , where G ∼ N (0 1) , is a standard Gaussian independent of ¯ β , and the deterministic parameter χ k c ∈ R is defined below via the state evolution recursion. Thus, f k in (15) can be viewed as a denoising function that can be tailored to take advantage of the prior on ¯ β . Indeed, we choose f k in this manner, and for k ≥ 1 , define:

$$f _ { k } ( s, c ) = \mathbb { E } [ \bar { \beta } \, | \, ( \chi _ { \epsilon } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { \epsilon } ^ { k } G = s ] \,, \quad \text{for } c \in [ C ].$$

State evolution iteratively computes the parameter χ k c ∈ R as follows, for k ≥ 1 . Letting

$$\delta _ { \text{in} } = \frac { n / R } { p / C } = \frac { C } { R } \delta \,,$$

given χ k c for c ∈ [ C ] , we compute χ k +1 c as:

$$( \chi _ { c } ^ { k + 1 } ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \left ( \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \sum _ { c ^ { \prime } = 1 } ^ { C } \widetilde { W } _ { r c ^ { \prime } } \mathbb { E } \left [ \left ( \bar { \beta } - f _ { k } \left ( \chi _ { c ^ { \prime } } ^ { k } \right ) ^ { 2 } \bar { \beta } + \chi _ { c ^ { \prime } } ^ { k } G, \, c ^ { \prime } \right ) \right ) ^ { 2 } \right ] \right ) ^ { - 1 },$$

where G ∼ N (0 1) , is independent of ¯ β , and σ 2 is the variance defined in the noise scaling assumption on p.9. The recursion is initialized with

$$( \chi ^ { 1 } _ { c } ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \left ( \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \sum _ { c ^ { \prime } = 1 } ^ { C } \widetilde { W } _ { r c ^ { \prime } } \text{Var} ( \bar { \beta } ) \right ) ^ { - 1 }, \ \ c \in [ C ].$$

The vectors Q k ∈ R n and c k ∈ R p in (15) have a block-wise structure, and are defined as follows. For i ∈ [ n ] , j ∈ [ p ] , recalling that r ( ) i and c ( j ) denote the respective row-block and column-block indices, we have:

$$Q _ { i } ^ { k } = \left ( \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { c } \widetilde { W } _ { r ( i ) c } \mathbb { E } \left [ \left ( \bar { \beta } - f _ { k } ( \left ( \chi _ { c } ^ { k } \right ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G, \, c \right ) ) ^ { 2 } \right ] \right ) ^ { - 1 }, \quad c _ { j } ^ { k } = - ( \chi _ { c ( j ) } ^ { k + 1 } ) ^ { 2 }.$$

The SC-AMP algorithm in in (15) and its state evolution are equivalent to the one proposed for a spatially coupled Gaussian design [34]. The key difference is that our algorithm uses a rescaled spatially coupled Bernoulli design. The theorem below shows that the state evolution guarantees remain valid for this setting. For an i.i.d. design (where R = C = 1 ), SC-AMP reduces to the standard AMP algorithm [16] for an i.i.d. Gaussian design.

Theorem 3.2 (State evolution result for SC-AMP) . Consider the QGT model (1) with a spatially coupled design defined via the ( ω, Λ) base matrix in Definition 3.1. Let the model assumptions in Section 2 and the noise scaling assumption (p. 9) be satisfied. Then, for the SC-AMP algorithm in (15) , run on the recentered and rescaled QGT model (13) with the denoising functions f k in (19) , we have the following convergence guarantee. For each k ≥ 0 and c ∈ [ C ] :

$$( \beta _ { \mathrm c }, \beta _ { \mathrm c } ^ { k + 1 } ) \stackrel { W _ { 2 } } { \to } ( \bar { \beta }, ( \chi _ { \mathrm c } ^ { k + 1 } ) ^ { 2 } \bar { \beta } + \chi _ { \mathrm c } ^ { k + 1 } G )$$

almost surely as n, p →∞ with n/p → δ .

The theorem is proved in Section 5, where the SC-AMP algorithm is shown to be a special case of an AMP algorithm for a generalized linear model with a generic spatially coupled design. We prove Theorem 3.2 by establishing a state evolution result for this general AMP algorithm (Theorem 5.2).

Performance measures. Theorem 3.2 allows us to compute the limiting values of performance measures such as the mean-squared error (MSE) and the normalized squared correlation, via the convergence property in (6). The MSE of the AMP estimates ˆ β k satisfy the following almost surely, for k ≥ 1 :

$$\lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \mathbb { E } \left [ \left ( \bar { \beta } - f _ { k } ( ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G, c ) \right ) ^ { 2 } \right ].$$

And the normalized squared correlation of the AMP estimates ˆ β k satisfy:

$$\lim _ { p \to \infty } \frac { \langle \hat { \beta } ^ { k }, \beta \rangle } { \| \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } \cdot \| \beta \| _ { 2 } ^ { 2 } } = \frac { ( \frac { 1 } { \bar { c } } \sum _ { c = 1 } ^ { C } \mathbb { E } [ f _ { k } ( ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G, c ) \cdot \bar { \beta } ] ) ^ { 2 } } { ( \frac { 1 } { \bar { c } } \sum _ { c = 1 } ^ { C } \mathbb { E } [ f _ { k } ( ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G, c ) ^ { 2 } ] ) \cdot ( \mathbb { E } [ \bar { \beta } ^ { 2 } ] ) }, \quad k \geq 1.$$

We can also obtain formulas for the limiting values of the false positive rate (FPR) and false negative rate (FNR). The choice of f k in (19) outputs a vector in R p , but we can obtain an estimate in { 0 1 , } p by thresholding the AMP iterate ˆ β K in the final iteration K to output a hard decision. For some chosen constant ζ , let us define the hard decision to be

$$1 \left \{ \hat { \beta } _ { j } ^ { K } > \zeta \right \} = 1 \left \{ f _ { K } ( \beta _ { j } ^ { K }, \mathfrak { c } ) > \zeta \right \}, \quad \text{for } j \in \mathcal { J } _ { \mathfrak { c } }, \mathfrak { c } \in [ \mathbb { C } ].$$

where the indicator function is applied component-wise to ˆ β k c . That is, we declare large entries of ˆ β K to be one (i.e., defective) and small entries of ˆ β K to be zero (i.e., non-defective). Based on the above function, let us denote the estimated defective set as ̂ S = { j : ˆ β K j &gt; ζ } .

The false positive rate (FPR) and the false negative rate (FNR) are defined as:

$$\ F P R = \frac { \sum _ { j = 1 } ^ { p } 1 \{ \beta _ { j } = 0 \cap j \in \widehat { \mathcal { S } } \} } { p - \sum _ { j = 1 } ^ { p } \beta _ { j } } \quad \text{and} \quad \ F N R = \frac { \sum _ { j = 1 } ^ { p } 1 \{ \beta _ { j } = 1 \cap j \notin \widehat { \mathcal { S } } \} } { \sum _ { j = 1 } ^ { p } \beta _ { j } }.$$

Corollary 3.3. Under the same assumptions as for Theorem 3.2, with a threshold ζ ∈ [0 , 1] for the final iteration K , as p →∞ , we have

$$\ F P R \stackrel { a. s. } { \rightarrow } \frac { 1 } { \overline { c } } \sum _ { c = 1 } ^ { c } \mathbb { P } [ f _ { K } ( \chi _ { \epsilon } ^ { K } G, c ) > \zeta ] \ \ a n d \ F N R \stackrel { a. s. } { \rightarrow } \frac { 1 } { \overline { c } } \sum _ { c = 1 } ^ { c } \mathbb { P } [ f _ { K } ( ( \chi _ { \epsilon } ^ { K } ) ^ { 2 } + \chi _ { \epsilon } ^ { K } G, c ) \leq \zeta ].$$

The result follows from Theorem 3.2 by applying the convergence property to suitable indicator functions. The proof uses the same steps as the analogous result for i.i.d. Bernoulli designs [14, Corollary 5.2], and is omitted.

## 3.2 Almost-Exact Recovery

Given ˆ β k , the SC-AMP estimate after k iterations, let us define the quantized estimate to be

$$\tilde { \beta } _ { j } ^ { k } = \begin{cases} 1 & \text{if $\hat{\beta}_{j}^{k}>0.5$} \\ 0 & \text{otherwise,} \end{cases}$$

̸

In this section, we show that the SC-AMP algorithm can attain almost-exact recovery with n = o p ( ) tests, by proving that it attains almost-exact recovery for any δ &gt; 0 (recall that δ = lim n →∞ n p ). To this end, we introduce the potential function to analyze the asymptotic MSE of SC-AMP as k → ∞ . Potential functions are widely used to characterize the limiting MMSE and mutual information in high-dimensional estimation problems (see, e.g., [54, 55]). Here we will use it only to characterize the asymptotic MSE of the AMP algorithm, both with and without spatial coupling (see Theorem 3.5).

Then, recalling the almost-exact recovery criterion in (2), the SC-AMP algorithm achieves almostexact recovery if lim k →∞ lim p →∞ 1 p ∑ p j =1 1 { ˜ β k j = β j } = 0 .

Definition 3.4. For b ∈ [0 , Var( ¯ )] β , δ &gt; 0 , the scalar potential function for the rescaled QGT model is defined as

$$U ( b ; \delta ) \coloneqq - \delta \Big ( 1 - \frac { \sigma ^ { 2 } } { ( b / \delta ) + \sigma ^ { 2 } } \Big ) + \delta \log \Big ( 1 + \frac { b } { \delta \sigma ^ { 2 } } \Big ) + 2 I \Big ( \bar { \beta } ; \sqrt { \frac { 1 } { ( b / \delta ) + \sigma ^ { 2 } } } \bar { \beta } + G \Big ). \quad \ ( 3 1 )$$

Here the mutual information I ( ; · · ) is computed with ¯ β ∼ P ¯ β independent of G ∼ N (0 1) , , and σ 2 = E [ ¯ Ψ 2 ] is the second moment of the rescaled noise (see p. 9).

Figure 2 plots the potential function for various values of δ . The next theorem characterizes the limiting MSE of the AMP algorithm via the minimizers and stationary points of the potential function. For clarity, we refer to the AMP algorithm under the rescaled i.i.d. QGT model in (14) as the iid-AMP algorithm, and the AMP algorithm under the rescaled spatially coupled QGT model in (13) as the SC-AMP algorithm. The SC-AMP algorithm in (15) reduces to iid-AMP with the trivial base matrix ( R = C = 1 ).

Theorem 3.5 (MSE of SC-AMP and iid-AMP) . Consider the QGT model (1) , and let the model assumptions in Section 2 and the noise scaling assumption (p. 9) be satisfied.

Figure 2: U b δ ( ; ) vs. b , for different δ . π = 0 1 . and σ = 1 × 10 -30 .

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-bf695aa1-e5cf-4299-a546-3d0cffcdd82b.jpg)

- 1. Consider a spatially coupled design defined via an ( ω, Λ) base matrix (Definition 3.1). For any γ &gt; 0 , there exist ω 0 &lt; ∞ and k 0 &lt; ∞ such that for all ω &gt; ω 0 and k &gt; k 0 , the asymptotic MSE of the SC-AMP algorithm almost surely satisfies:

$$\lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } < \left ( \max \left \{ \underset { b \in [ 0, V \arrow ( \bar { \beta } ) ] } { \arg \min } U ( b ; \delta _ { \text{in} } ) \right \} + \gamma \right ) \frac { \Lambda + \omega } { \Lambda }.$$

- 2. With an i.i.d. design (i.e., 1 × 1 base matrix with W 11 = 1 ), the asymptotic MSE of the iid-AMP algorithm almost surely satisfies:

$$\lim _ { k \to \infty } \lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } = \max \left \{ b \in [ 0, \text{Var} ( \bar { \beta } ) ] \, \colon \partial _ { 1 } U ( b ; \delta ) = 0 \right \},$$

where ∂ 1 denotes the partial derivative w.r.t. the first argument.

The proof is given in Section 6. Part 1 of the theorem says that for sufficiently large base matrix parameters (with ω ≪ Λ ), the MSE of the SC-AMP algorithm is bounded by the largest minimizer of the potential function. (The max {·} indicates that if there are multiple minimizers, the largest one is chosen.) Part 2 of the theorem says that the MSE of iid-AMP algorithm is given by the largest stationary point of the potential function. In Figure 2, we observe that for δ = 0 05 . and δ = 0 02 . , the unique minimizer is b = 0 , but the largest stationary point is strictly larger than zero. This implies that the limiting MSE of SC-AMP algorithm is 0, but that of iid-AMP algorithm is strictly larger than 0. The next lemma quantifies this observation, showing that for any δ &gt; 0 , the largest minimizer of the potential function tends to zero as the noise variance σ → 0 .

Lemma 3.6. Consider the scalar potential function U b δ ( ; ) in (31) with δ &gt; 0 . For any ∆ ∈ (0 , δ ) , there exists σ 0 (∆) &gt; 0 such that for all σ &lt; σ 0 (∆) , we have the rate

$$\max \left \{ \underset { b \in ( 0, \text{Var} ( \bar { \beta } ) ] } { \arg \min } U ( b ; \delta ) \right \} < \frac { 7 } { 2 } \delta \left ( \sigma ^ { 2 - 2 \Delta / \delta } \right ).$$

The proof is given in Appendix B. Using Lemma 3.6 in Part 1 of Theorem 3.5 yields the following bound on the MSE of the SC-AMP algorithm in the low noise regime.

Theorem 3.7 (MSE of SC-AMP in the low-noise regime) . Consider the setup of part 1 of Theorem 3.5, for any δ &gt; 0 . Then for any ϵ &gt; 0 and ∆ ∈ (0 , δ/ (1 + ϵ )) , there exists σ 0 (∆) &gt; 0 such that following holds for any noise variance σ &lt; σ 0 (∆) . There exist finite ω 0 and k 0 such that for all ω &gt; ω 0 , k &gt; k 0 , and Λ sufficiently large, the asymptotic MSE of the SC-AMP algorithm almost surely satisfies:

$$\lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } < 4 \delta \sigma ^ { 2 - 2 \Delta ( 1 + \epsilon ) / \delta } + \epsilon.$$

Proof. Let ϵ, γ &gt; 0 . Recalling that δ in = δ Λ Λ+ ω -1 , for sufficiently large Λ /ω we have δ in &lt; δ 1+ ϵ . Then, from Part 1 of Theorem 3.5, for k &gt; k 0 , ω &gt; ω 0 and sufficiently large Λ /ω , we have

$$\lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } < \left ( \max \left \{ \underset { b \in [ 0, V \bar { \epsilon } ] } { \arg \min } U ( b ; \delta _ { \text{in} } ) \right \} + \gamma \right ) ( 1 + \epsilon ),$$

Taking γ to be small enough and using Lemma 3.6 in (35) gives the result.

Corollary 3.8 (SC-AMP achieves almost exact recovery for any δ &gt; 0 ) . Consider the setting and assumptions of Theorem 3.7. Let ˜ β k be the quantized estimate produced from the SC-AMP iterate ˆ β k , according to (30) . Then, almost surely we have:

̸

$$\lim _ { p \to \infty } \frac { 1 } { p } \sum _ { j = 1 } ^ { p } \mathbf 1 \{ \tilde { \beta } _ { j } ^ { k } \neq \beta _ { j } \} \leq 4 \lim _ { p \to \infty } \frac { 1 } { p } \| \beta - \hat { \beta } ^ { k } \| _ { 2 } ^ { 2 } < 4 \left ( 4 \delta \sigma ^ { 2 - 2 \Delta ( 1 + \epsilon ) / \delta } + \epsilon \right ).$$

̸

In particular, for noiseless QGT we have lim k →∞ lim p →∞ 1 p ∑ p j =1 1 { ˜ β k j = β j } = 0 almost surely for any δ &gt; 0 .

̸

̸

̸

Proof. From the definition of the quantized estimate in (30), it follows that for j ∈ [ p ] , we have | β j -ˆ β k j | ≥ (0 5) . 1 { ˜ β k j = β j } , which implies 1 { ˜ β k j = β j } ≤ 4 | β j -ˆ β k j | 2 . This gives the first inequality in (36). The second equality follows from Theorem 3.7. The result for the noiseless case follows by setting σ = 0 , and taking a sequence ( ϵ k ) such that ϵ k → 0 as k → ∞ . We note that lim p →∞ 1 p ∑ p j =1 1 { ˜ β k j = β j } exists for each k by the state evolution result in Theorem 3.2.

The guarantees in Corollary 3.8 are analogous to those in [34, Theorem 1.7 and Corollary 1.8] for a linear model with a spatially coupled Gaussian design matrix. Specifically, [34, Theorem 1.7] shows that when the sampling ratio δ is larger than the Rényi information dimension of the signal prior, the MSE of the SC-AMP algorithm satisfies lim k →∞ lim p →∞ 1 p ∥ β -ˆ β k ∥ 2 2 &lt; Cσ 2 , for sufficiently small noise variance σ 2 . Here the constant C depends on δ and on the prior. The Rényi information dimension for a Bernoulli prior is 0, so Corollary 3.8 is consistent with the result in [34]. The key difference is that we use a binary-valued SC design for the QGT model rather than the Gaussian one in [34]. Our analysis of the fixed point of the SC state evolution to establish Theorem 3.5 is also simpler than in [34], where the authors use a continuum version of the state evolution along with a perturbation argument. In contrast, we use a straightforward potential function analysis based on the recipe provided in [37] for analyzing coupled recursions.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-28603977-21f1-4a0c-9fbd-35f3be700470.jpg)

(a) SC-AMP vs. iid-AMP

Figure 3: Normalized squared correlation for Noiseless QGT. π = 0 3 . , p = 20000 . With spatial coupling parameters ω = 6 Λ = 40 , , the inner block size p/ Λ = 500 . In (b), we use p = 2000 for both iid-LP and SC-LP due to computational constraints. Error bars indicate one standard deviation.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-858ada76-efab-4c81-a4de-2eff40d71d40.jpg)

## 3.3 Numerical Simulations

We present simulation results for finite length SC-AMP and compare its performance against alternative algorithms and the information-theoretic limit. The performance in all the plots is either measured via the normalized squared correlation between the SC-AMP estimate and the signal (see (26)) or via the FPR and FNR (see (28)). In the plots, curves labeled 'SC-AMP' show the empirical performance of the SC-AMP algorithm, while the 'SC-SE' curves refer to its theoretical performance predicted via state evolution. The corresponding curves for an i.i.d. design are labeled 'iid-AMP' and 'iid-SE'. For empirical performance curves, each point is obtained from 10 independent runs, where in each run, the SC-AMP algorithm is executed for 300 iterations. Other implementation details are described in Appendix C. Python code for all the simulations is available at [56].

Figure 3b shows how the AMP algorithm compares to linear programming (LP) estimator,

Figure 3a shows how normalized squared correlation varies with the sampling ratio δ for noiseless QGT, for both spatially coupled and i.i.d. designs. We observe that SC-AMP outperforms iid-AMP, justifying the use of the SC design. The orange vertical line show the information-theoretic lower bound on the ratio n/p obtained from (4). Specializing (4) to the case of L = 2 , we get the information-theoretic lower bound on the number of tests for noiseless QGT: n ∗ = 2 H π ( ) p log p . We observe that the performance of SC-SE improves and approaches n ∗ as we increase the size of the spatial coupling parameters ( ω, Λ) from (6,40) to (40,400). We did not implement the SC-AMP for ( ω, Λ) = (40 400) , as it requires a large amount of computational memory. The difference between the SC-SE plot and the SC-AMP plot for (6 , 40) is due to finite length effects, since the inner block size p/ Λ is only 500.

defined as the solution of the following linear program:

$$\minimize & \ \| \beta \| _ { 1 } \\ \text{subject to} & \ y = X \beta, \text{ and } 0 \leq \beta _ { j } \leq 1, \ \ j \in [ p ].$$

Similar reconstruction algorithms are commonly used for compressed sensing [57]. LP based estimators have also been used in Boolean in group testing [1]. We observe that the AMP algorithm outperforms LP for both i.i.d. and SC designs, and that the performance of LP is similar with both designs. This is because the LP algorithm is not equipped to take advantage of the spatially coupled design. LP is also more computationally intensive than the SC-AMP algorithm and challenging to implement for large values of p . Therefore, we use a smaller p for all our LP experiments.

Figure 4b shows the tradeoff between the FPR and the FNR for noisy QGT with δ = 0 46 . and σ 2 = 0 0016 . . Following the model in (1), for the i.i.d. design we consider Gaussian noise with Ψ i i . i . d . ∼ N (0 , pσ 2 ) , as previously investigated in [11, 14]. For the SC design, we consider Ψ i i . i . d . ∼ N (0 , pσ 2 / (2 C ) ) . As described in Section 3, for α = 0 5 . , the expected number of items in each test is approximately p 4 C for the SC design, compared to p/ 2 for the i.i.d. design. This choice of noise variance for the SC model ensures that the signal-to-noise ratio E [ ∥ Xβ ∥ 2 ] / E [ ∥ Ψ ∥ 2 ] is similar for both designs.

Figure 4a shows the tradeoff between the FPR and the FNR for noiseless QGT with δ = 0 38 . . The tradeoff curve is obtained by thresholding the AMP or LP estimate with different thresholds ζ , as described in (27). SC-AMP achieves perfect recovery at this value of δ , so its FPR and FNR are both 0, for all threshold values. As expected, SC-AMP does significantly better than iid-AMP and LP.

In the noisy setting, the AMP algorithm is compared to the following convex programming (CVX) estimator:

$$\minimize & \quad \frac { 1 } { 2 \text{Var} ( \Psi ) } \| y - X \beta \| _ { 2 } ^ { 2 } + \| \beta \| _ { 1 } \log \frac { 1 - \pi } { \pi } \\ \text{subject to} & \quad 0 \leq \beta _ { j } \leq 1, \quad j \in [ p ].$$

This estimator is obtained via a convex relaxation of the MAP estimator for QGT. Figure 4b shows that in the presence of a small amount of noise, SC-AMP continues to achieve perfect recovery, outperforming both CVX and iid-AMP. Surprisingly, in the presence of noise, the performance of CVX is worse with the SC design than with the i.i.d. one, possibly because it does not take advantage of the band-diagonal structure in the SC design matrix.

## 4 SC-AMP for Pooled Data

In this section we extend the SC-AMP algorithm to the pooled data model in (3) with the spatially coupled design X sc defined in (8). We apply SC-AMP to a centered and rescaled version of the pooled model, as we did for QGT in (13). Recalling the decomposition X sc ij = αW rc + √ nα (1 -α ) R ˜ X sc ij ,

Figure 4: FPR vs. FNR tradeoff. In both (a) and (b), we use π = 0 3 . , and thresholds ζ ∈ { 0 1 . , 0 2 . , . . . , 0 9 . } . We take p = 20000 for AMP, and p = 2000 for LP and CVX.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-d59c94d5-a9a2-4975-94b8-e127f4d907ed.jpg)

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-ebe0aaa8-e35e-4b25-bac1-3fa6d6d58207.jpg)

we have

$$Y _ { i ; } & = \sum _ { j = 1 } ^ { p } X _ { i j } ^ { s c } B _ { j, } + \Psi _ { i, ; } = \sum _ { j = 1 } ^ { p } \left ( \alpha W _ { r ( i ) \epsilon ( j ) } + \sqrt { \frac { n \alpha ( 1 - \alpha ) } { R } } \widetilde { X } _ { i j } ^ { s c } \right ) B _ { j, ; } + \Psi _ { i, ; } \\ & \implies Y _ { i, ; } - \sum _ { j = 1 } ^ { p } \alpha W _ { r ( i ) \epsilon ( j ) } B _ { j, ; } = \sqrt { \frac { n \alpha ( 1 - \alpha ) } { R } } \sum _ { j = 1 } ^ { p } \widetilde { X } _ { i j } ^ { s c } B _ { j, ; } + \Psi _ { i, ; } \,. \\ - \cdots & \cdots$$

Defining

$$\widetilde { Y } _ { i, \colon } \coloneqq \frac { 1 } { \sqrt { n \alpha ( 1 - \alpha ) / R } } \left ( Y _ { i, \colon } - \alpha \left ( W _ { r ( i ) 1 } \sum _ { j \in \mathcal { J } _ { 1 } } B _ { j, \colon } + \cdots + W _ { r ( i ) C } \sum _ { j \in \mathcal { J } _ { C } } B _ { j, \colon } \right ) \right ),$$

and ˜ Ψ i, : := 1 √ nα (1 -α / ) R Ψ i, : gives us the rescaled pooled data model:

$$\widetilde { Y } _ { i, ; } = B ^ { \top } \widetilde { X } _ { i, ; } ^ { s c } + \widetilde { \Psi } _ { i, ; } \in \mathbb { R } ^ { L }, \quad \text{for } i \in [ n ].$$

The sets ( J c ) c ∈ [ C ] are defined in (17). In the noiseless setting, the terms ∑ j ∈J c B j, : can be obtained with an extra C = O (1) tests, where the c th test only includes items from J c .

Model and noise scaling assumptions. The signal matrix B ∈ R p × L and the rescaled noise matrix ˜ Ψ ∈ R n × L Ψ are both independent of the design matrix. As p → ∞ , we assume that n/p → δ &gt; 0 . As p, n → ∞ , the empirical distributions of the rows of B and ˜ Ψ each converge to well-defined limits. More precisely, B W → ¯ B and ˜ Ψ W → ¯ Ψ , for L -dimensional random vectors ¯ B ∼ Categorical ( π ) and ¯ Ψ ∼ P ¯ Ψ .

## 4.1 Matrix SC-AMP Algorithm

The goal is to recover B from ˜ Y generated according to the rescaled model (39). The matrix SCAMP algorithm is initialized with ̂ B 0 j, : = E [ ¯ ] B for j ∈ [ p ] , and ˜ Θ = 0 ˜ Y -˜ X B sc ̂ 0 . For iteration k ≥ 1 , we compute:

$$\widetilde { \Theta } ^ { k } & = \widetilde { Y } - \widetilde { X } ^ { s c } \widehat { B } ^ { k } + U ^ { k } \\ B ^ { k + 1 } & = V ^ { k } + \widehat { B } ^ { k } ; \quad \widehat { B } ^ { k + 1 } = f _ { k + 1 } ( B ^ { k + 1 }, \mathcal { C } ),$$

where U k ∈ R n × L , V k ∈ R p × L , and f k : R L +1 → R , which acts row-wise on its input, will be defined after the state evolution parameters are introduced. Note that the vector C ∈ R p is defined in (16).

State evolution. For k ≥ 0 , we iteratively compute the following L × L covariance matrices, starting from initialization ψ 0 c = Cov( ¯ ) B , c ∈ [ C ] :

$$\phi _ { r } ^ { k } = \text{Cov} [ \bar { \Psi } ] + \frac { 1 } { \delta _ { \text{in} } } \sum _ { r = 1 } ^ { C } \widetilde { W } _ { \text{rc} } \psi _ { \text{c} } ^ { k } \,, \quad r \in [ \text{R} ],$$

$$\phi _ { r } ^ { k } & = C o v [ \bar { \Psi } ] + \frac { 1 } { \delta _ { \text{in} } } \sum _ { r = 1 } ^ { \text{c} } \widetilde { W } _ { \text{rc} } \psi _ { \text{c} } ^ { k }, \quad \ \ r \in [ \text{R} ], \\ \text{T} _ { \text{c} } ^ { k } & = \left [ \sum _ { r = 1 } ^ { \text{R} } \widetilde { W } _ { \text{rc} } [ \phi _ { r } ^ { k } ] ^ { - 1 } \right ] ^ { - 1 }, \quad \ \psi _ { \text{c} } ^ { k + 1 } = \mathbb { E } \left [ \left ( f _ { k } ( \bar { B } + G _ { \text{c} } ^ { k }, \text{c} ) - \bar { B } \right ) \left ( f _ { k } ( \bar { B } + G _ { \text{c} } ^ { k }, \text{c} ) - \bar { B } \right ) ^ { \text{T} } \right ], \quad \text{c} \in [ \text{C} ]. \\ & \quad. \quad. \quad - \quad \ \sim$$

where G k c ∼ N (0 T ) , k c is independent of ¯ B , and we recall that ˜ W is defined in (18).

We define a matrix Q k ∈ R L R × L C in terms of sub-matrices Q k r c , ∈ R L × L , which for r ∈ [ R ] , c ∈ [ C ] , are given by

$$Q _ { r, \mathfrak { c } } ^ { k } = ( \phi _ { r } ^ { k } ) ^ { - 1 } \left ( \sum _ { r ^ { \prime } = 1 } ^ { \mathrm R } \widetilde { W } _ { r ^ { \prime } \mathfrak { c } } ( \phi _ { r ^ { \prime } } ^ { k } ) ^ { - 1 } \right ) ^ { - 1 }.$$

Then, the rows of the matrix U k ∈ R n × L in (40) are defined as

$$U _ { i, \colon } ^ { k } = \frac { 1 } { \delta _ { \text{in} } } \widetilde { \Theta } _ { i, \colon } ^ { k - 1 } \sum _ { c = 1 } ^ { c } \widetilde { W } _ { r ( i ), c } Q _ { r ( i ), c } ^ { k - 1 } \left \langle f _ { k } ^ { \prime } ( B _ { j, \colon } ^ { k }, c ) \right \rangle _ { c } ^ { \top }, \quad \text{for } i \in [ n ]$$

where ⟨·⟩ c is the empirical average over rows j ∈ J c , i.e., 〈 f ′ k ( B k j, : , c ) 〉 c = 1 p/ C ∑ j ∈J c f ′ k ( B k j, : , c ) . Next, the rows of the matrix V k ∈∈ R p × L are given by

$$V _ { j, \colon } ^ { k } = \sum _ { i = 1 } ^ { n } \widetilde { X } _ { i j } ^ { s c } \, \widetilde { \Theta } _ { i, \colon } ^ { k } Q _ { r ( i ), \mathfrak { c } ( j ) } ^ { k }, \quad \text{for } j \in [ p ].$$

Finally, the Bayes-optimal denoising function f k is given by the following, for j ∈ J c , c ∈ [ C ] :

$$f _ { k } ( B _ { j, \cdot } ^ { k }, \mathbf c ) = \mathbb { E } [ \bar { B } | \bar { B } + G _ { \mathbf c } ^ { k } = B _ { j, \cdot } ^ { k } ], \ \ G _ { \mathbf c } ^ { k } \sim \mathcal { N } \left ( 0, \mathbf T _ { \mathbf c } ^ { k } \right ) \text{ independent of } \bar { B }.$$

The matrix SC-AMP algorithm was proposed and analyzed in [58] for a model with a spatially coupled Gaussian design matrix. Similar to our analysis of SC-AMP for QGT, we could apply a

reduction technique along with the universality result of [36] to establish a state evolution characterization for the matrix SC-AMP applied to pooled data. Such a result would be analogous to Theorem 3.2, and show that for each iteration k ≥ 1 , the joint empirical distribution of the rows of ( B,B k ) converges as:

$$( B, B ^ { k } ) _ { \mathcal { J } _ { c } ; } \stackrel { W _ { 2 } } { \to } ( \bar { B }, \bar { B } + G _ { c } ^ { k } ), \ \ c \in [ C ].$$

To analyze the limiting MSE and error rate of the matrix SC-AMP algorithm, we need to characterize the fixed point of the state evolution recursion in (41) (as k →∞ ). In QGT, the state evolution fixed point was characterized via the minimizer of a scalar potential function (Theorem 3.5). Extending this approach to the pooled data setting is challenging as the state evolution parameters are now L × L matrices rather than scalars. In the following section, we circumvent this issue by showing that a suboptimal AMP algorithm still achieves almost-exact recovery for any δ &gt; 0 . The suboptimal algorithm applies the SC-AMP algorithm column-wise to ˜ Y ∈ R n × L , ignoring the correlation between the columns of the signal matrix B .

## 4.2 Almost-Exact Recovery via Column-wise SC-AMP

Given ˜ X ,Y sc ˜ from the rescaled model (39), we run the SC-AMP algorithm column-wise on ˜ Y . Specifically, for l ∈ [ L ] , run the SC-AMP algorithm (15) with inputs Y : ,l and ˜ X sc to produce the estimate ̂ B k : ,l after k ≥ 1 iterations. For the SC-AMP algorithm applied to column l ∈ [ L ] , the denoiser f k in (19) is computed with ¯ β ∼ Bernoulli ( π l ) . (We recall that the rows of the signal follow the prior ¯ B ∼ Categorical ( π ) where π = ( π , . . . , π 1 L ) .)

The column-wise SC-AMP algorithm can be viewed as an instance of the matrix SC-AMP algorithm with a suboptimal denoiser, obtained by replacing the conditional expectation E [ ¯ B | ¯ + B G k c ] in (43) with the marginal conditional expectations E [ ¯ B l | ¯ B l +( G k c ) ] l , for l ∈ [ L ] . Let us define the quantized estimate after k iterations of the column-wise SC-AMP algorithm to be

$$\widetilde { B } _ { j l } ^ { k } = \begin{cases} 1 & \text{if } \widehat { B } _ { j l } ^ { k } > 0. 5 \\ 0 & \text{otherwise}, \end{cases}$$

where ̂ B k is the estimate obtained from the SC-AMP algorithm. In practice, we can quantize the estimate in a better manner, by setting the largest entry in the row of ̂ B to one and the remaining entries in the row to zero. We do not use this form of quantization for our almost-exact recovery result since we want to directly apply the SC-AMP results for QGT to the pooled data setting.

Theorem 4.1. Consider the noiseless pooled data problem with the assumptions stated on p.17, for any δ &gt; 0 . There exist finite ω 0 and k 0 such that for all ω &gt; ω 0 , k &gt; k 0 , and sufficiently large Λ , the quantized estimate ˜ B k ∈ { 0 1 , } p × L produced by the column-wise SC-AMP algorithm almost surely satisfies:

̸

$$\lim _ { k \to \infty } \lim _ { p \to \infty } \frac { 1 } { p } \sum _ { j = 1 } ^ { p } 1 \left \{ \widetilde { B } _ { j, \colon } ^ { k } \neq B _ { j, \colon } \right \} = 0.$$

̸

̸

Proof. The model assumptions imply that for each l ∈ [ L ] , the empirical distribution of column B : ,l converges in Wasserstein distance to Bernoulli ( π l ) . By Corollary 3.8, we have that the SC-AMP algorithm applied to ˜ Y : ,l satisfies lim k →∞ lim p →∞ 1 p ∑ p j =1 1 { ˜ B k j,l = B j,l } = 0 almost surely, for each l ∈ [ L ] . The result follows by noting that 1 p ∑ p j =1 1 { ˜ B k j, : = B j, : } ≤ 1 p ∑ L l =1 ∑ p j =1 1 { ˜ B k jl = B jl } .

̸

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-9d387348-9863-402d-9903-b19ee81b1a7c.jpg)

(a) SC-AMP vs. iid-AMP

Figure 5: Noiseless Pooled Data with π = [1 3 1 / , / , 3 1 / 3] with p = 20000 and spatial coupling parameters ω = 6 Λ = 40 , . In (b), we set p = 1000 for LP due to the high computational cost. Error bars indicate one standard deviation.

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-a86ca6cb-711e-43b2-8d29-547c7afef588.jpg)

We can also obtain error guarantees in the low-noise regime for the column-wise SC-AMP algorithm, similar to Theorem 3.7 and Corollary 3.8. We remark that although column-wise SC-AMP is convenient for theoretical analysis, at finite dimensions it is inferior to the matrix SC-AMP algorithm that takes advantage of the correlation in the columns of B via the denoiser in (43). This is illustrated in the numerical experiments below.

## 4.3 Numerical Simulations

We present simulation results for noiseless pooled data using matrix SC-AMP, measuring the performance via the normalized correlation 1 p ∑ p j =1 ⟨ ̂ B k j, : , B j, : ⟩ after k iterations of the algorithm. By (44), we this normalized correlation converges as

$$\frac { 1 } { p } \sum _ { j = 1 } ^ { p } \langle \widehat { B } _ { j, \cdot } ^ { k }, B _ { j, \cdot } \rangle \stackrel { a. s. } { \to } \frac { 1 } { \text{C} } \sum _ { \text{c=1} } ^ { \text{c} } \mathbb { E } \left [ \langle f _ { k } ( \bar { B } + G _ { \text{c} } ^ { k }, \text{c} ), \bar { B } \rangle \right ].$$

Each point on the AMP performance curves is obtained from 10 independent runs, where in each run, the algorithm is executed for 300 iterations. Our benchmark will be the linear programming (LP) estimator tailored to the pooled data problem [14, Section 4.1]. Recall from (4) that the information theoretic lower bound on the number of tests for the noiseless pooled data problem is

$$n ^ { * } = \frac { p } { \log p } \left ( \max _ { r \in \{ 1, \dots, L - 1 \} } \frac { 2 [ H ( \pi ) - H ( \pi ^ { ( r ) } ) ] } { L - r } \right ).$$

Figure 5 shows how the normalized correlation varies with the sampling ratio δ for pooled data with L = 3 equally likely categories. Figure 5a shows that the state evolution prediction for performance of the matrix SC-AMP algorithm (SC-SE curves) improves as the spatial coupling

Figure 6: Matrix-SC AMP vs column-wise SC-AMP for noiseless pooled data with π = [1 / , 3 1 / , 3 1 / 3] , p = 20000 and spatial coupling parameters ω = 6 Λ = 40 , .

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-ff33ac9c-54f5-471d-aabe-327396dbbffa.jpg)

parameters ( ω, Λ) increase from (6 40) , to (20 200) , . As in QGT, the gap between the empirical performance of matrix SC-AMP and the state evolution prediction for (6 , 40) is due to finite length effects. We did not implement the matrix SC-AMP algorithm for (20 200) , as it requires a large amount of computational memory. In Figure 5b, we observe that iid-AMP, iid-LP, and SC-LP have similar performance, and SC-AMP outperforms all of these, justifying the use of a SC design matrix with matrix SC-AMP for recovery. We also implemented the iterative hard thresholding algorithm [14, Section 4.1] but found that it performed significantly worse AMP and LP, and so omitted it from our comparisons.

Figure 6 compares the performance of matrix SC-AMP with the column-wise SC-AMP algorithm. To make the algorithms comparable, the estimates from each algorithm (and the corresponding SE) were quantized in the same way after their final iteration, using the rule in (45). This leads to a slight difference in the AMP performance curves and the theoretical SE estimates compared to Figure 5, where no quantization was used. As expected, the matrix SC-AMP algorithm outperforms the column-wise SC-AMP algorithm since the former takes advantage of the correlation within each row of the matrix signal. Nevertheless, the column-wise SC-AMP algorithm performs slightly better than the matrix AMP algorithm with an i.i.d. matrix.

## 5 Proof of Theorem 3.2

We start by defining generalized white noise matrices , which will be used in the proof of the theorem.

Definition 5.1. [36, Definition 2.15] A generalized white noise matrix ˜ X ∈ R n × p with a (deterministic) variance profile S ∈ R n × p is one satisfying the following conditions, for i ∈ [ n , j ] ∈ [ p ] :

- 1. All entries ˜ X ij are independent.
- 2. Each entry ˜ X ij has mean 0, variance n -1 S ij , and higher moments satisfying, for each integer

m ≥ 3 ,

$$\lim _ { n, p \to \infty } p \cdot \max _ { i \in [ n ] } \max _ { j \in [ p ] } \mathbb { E } \left [ | \widetilde { X } _ { i j } | ^ { m } \right ] = 0.$$

## 3. For a constant C &gt; 0 ,

$$\max _ { i \in [ n ] } \max _ { j \in [ p ] } S _ { i j } \leq C, \quad \lim _ { n, p \to \infty } \max _ { i \in [ n ] } \left | \frac { 1 } { p } \sum _ { j = 1 } ^ { p } S _ { i j } - 1 \right | = 0, \quad \lim _ { n, p \to \infty } \max _ { j \in [ p ] } \left | \frac { 1 } { n } \sum _ { i = 1 } ^ { n } S _ { i j } - 1 \right | = 0. \quad ( 4 7 )$$

Definition 5.1 simplifies for the case of S ij = 1 for all ( i, j ) ∈ [ n ] × [ p ] . In this case, the entries are all i.i.d. with variance 1 /n , the third condition in the definition is trivially satisfied, and the second condition requires moments of order 3 and higher to decay faster than 1 /p . The rescaled i.i.d. design matrix ˜ X iid in (9) is a generalized white noise matrix, but the rescaled spatially coupled matrix ˜ X sc in (10) is not. Indeed, ˜ X sc satisfies the first two requirements in Definition 5.1 and from (11), its variance profile of satisfies the first and last conditions in (47), but not the second: for i ∈ [ n ] , we have 1 p ∑ p j =1 S ij = R C , which is close to, but not equal to 1 for large Λ /ω .

We prove Theorem 3.2 via a more general result, for a generalized linear model with a spatial coupled design, where the observations y i ∈ R are generated as:

$$\tilde { y } _ { i } = q \left ( \left ( \widetilde { X } _ { i, \cdot } ^ { \text{sc} } \right ) ^ { \top } \beta, \, \widetilde { \Psi } _ { i } \right ) = q \left ( \Theta _ { i }, \widetilde { \Psi } _ { i } \right ), \quad \text{for } i \in [ n ],$$

Here β ∈ R p is the signal to be estimated, ˜ Ψ ∈ R n is a noise vector, and q : R 2 → R is a known output function. We also allow ˜ X sc to be more general than the one in Definition 3.1. The generalized spatially coupled matrix ˜ X sc ∈ R n × p consists of independent zero-mean entries whose variances are specified by a generic base matrix ˜ W ∈ R R × C , which satisfies the following conditions:

$$\sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } = 1 \quad \text{for $c\in [C]$,} \quad \text{and} \quad \kappa _ { 1 } \leq \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \leq \kappa _ { 2 }, \quad \text{for $r\in [R]$,}$$

for some κ , κ 1 2 &gt; 0 . Given a base matrix ˜ W satisfying (49), we construct the spatially coupled matrix ˜ X sc with independent entries drawn from a distribution with zero-mean and variance E [ | ˜ X sc ij | 2 ] = ˜ W r ( ) i c ( j ) n/R , for i ∈ [ n , j ] ∈ [ p ] . We also assume that the higher moments E [ | ˜ X sc ij | m ] for m ≥ 3 , satisfy (46).

The first condition in (49) ensures that the expected squared norm of each column of ˜ X sc is 1 , and the second condition in (49) bounds the variance of each entry of ˜ X β sc from above and below.

## High-level sketch of proof of Theorem 3.2. The proof consists of three reductions.

- 1. In Section 5.1, we introduce the spatially coupled generalized approximate message passing algorithm (SC-GAMP) for the generalized linear model in (48), and characterize its performance via state evolution (Theorem 5.2). We then reduce the SC-AMP algorithm in (15) to SC-GAMP, and use the state evolution result of the latter to prove Theorem 3.2.

- 2. To prove the state evolution result for SC-GAMP (Theorem 5.2), we show that the algorithm can be written as an instance of an abstract matrix-AMP iteration defined via a generalized white noise matrix. This reduction, shown in Appendix A.1, is similar to the one used in [35, Appendix A] for reducing the SC-GAMP algorithm for a Gaussian design to an abstract matrix-AMP iteration.
- 3. To prove the state evolution result for the abstract matrix-AMP (Theorem A.1), we show that it is a special case of an AMP iteration for generalized white noise matrices, for which a rigorous state evolution result was established in [36]. We refer to the latter iteration as UAMP, where the 'U' stands for universal. The technique for reducing the abstract matrix-AMP to U-AMP is similar to the one presented in [14]. This is shown in Appendix A.2.

As before, to simplify notation, for vectors a ∈ R p and b ∈ R n , we will use a c := a J c and b r := b I r , where J c and I r are defined in (17). There will be no notation simplification for matrices.

## 5.1 The SC-GAMP Algorithm and its State Evolution

The SC-GAMP algorithm aims to estimate the signal β ∈ R p from observations ˜ y ∈ R n generated according to the generalized linear model (48). For iteration k ≥ 0 , the algorithm computes:

$$& \Theta ^ { k } = \widetilde { X } ^ { \text{sc} } \hat { \beta } ^ { k } - b ^ { k } \odot \widehat { R } ^ { k - 1 } ; \quad \widehat { R } ^ { k } = g _ { k } ( \Theta ^ { k }, \tilde { y }, \mathcal { R } ) \\ & \beta ^ { k + 1 } = ( \widetilde { X } ^ { \text{sc} } ) ^ { \top } \widehat { R } ^ { k } - \epsilon ^ { k } \odot \hat { \beta } ^ { k } ; \quad \hat { \beta } ^ { k + 1 } = f _ { k + 1 } ( \beta ^ { k + 1 }, \mathcal { C } ),$$

where ⊙ denotes element-wise product. The algorithm is initialized with some ˆ β 0 ∈ R p and Θ = 0 ˜ X β sc ˆ 0 . The functions g k and f k +1 act row-wise on their input, and

$$\mathcal { C } & = ( \underbrace { 1, \dots, 1 } _ { p / C \text{ entries } } _ { p / C \text{ entries } }, \underbrace { 2, \dots, 2 } _ { p / C \text{ entries } }, \dots, \underbrace { C, \dots, C } _ { p / C \text{ entries } } ) ^ { \top } \in \mathbb { R } ^ { p }, \\ \mathcal { R } & = ( \underbrace { 1, \dots, 1 } _ { n / R \text{ entries } }, \underbrace { 2, \dots, 2 } _ { n / R \text{ entries } }, \dots, \underbrace { R, \dots, R } _ { n / R \text{ entries } } ) ^ { \top } \in \mathbb { R } ^ { n }.$$

The entries of c k ∈ R p and b k ∈ R n are defined as follows, for j ∈ [ p ] , i ∈ [ n ] :

$$c _ { j } ^ { k } = \sum _ { r = 1 } ^ { \mathrm R } \frac { \widetilde { W } _ { \mathrm c } } { n / \mathrm R } \sum _ { i \in \mathcal { I } _ { r } } \partial _ { 1 } g _ { k } ( \Theta _ { i } ^ { k }, \tilde { y } _ { i }, r ), \ \ b _ { i } ^ { k } = \sum _ { c = 1 } ^ { \mathrm c } \frac { \widetilde { W } _ { \mathrm c } } { n / \mathrm R } \sum _ { j \in \mathcal { J } _ { c } } \partial _ { 1 } f _ { k } ( \beta _ { j } ^ { k }, \mathrm c ).$$

State evolution. The 'memory' terms -b k ⊙ ̂ R k -1 and -c k ⊙ ˆ β k in (50) debias the iterates Θ k and β k +1 , ensuring that their empirical distributions are accurately captured by state evolution in the high-dimensional limit. Theorem 5.2 below shows that for each k ≥ 1 and c ∈ [ C ] , the empirical distribution of β k c converges to the distribution of µ k β, c ¯ + β G k β, c where ¯ β is the random variable representing the limiting distribution of the entries of the signal β c , and G k β, c ∼ N ( 0 ( , σ β, c ) 2 ) is independent of ¯ β . The deterministic parameters µ k β, c ∈ R and σ k β, c ∈ R are defined below. The result implies that the empirical distribution of the estimate ˆ β k c converges to the distribution of f k ( µ k β. c ¯ + β G k β, c ) . Thus, f k can be viewed as a denoising function that can be tailored to take

advantage of the prior in ¯ β . Theorem 5.2 also shows that the joint empirical distribution of the rows of (Θ r , Θ ) k r converges to N (0 Σ , k, r ) , where Σ k, r ∈ R 2 × 2 is defined below.

We now describe the state evolution recursion defining µ k β, c , σ k β, c ∈ R and Σ k, r ∈ R 2 × 2 . Define ¯ g k : R 3 × [ R ] → R such that

$$g _ { k } ( \Theta _ { i } ^ { k }, \tilde { y } _ { i }, r ) = \bar { g } _ { k } ( \Theta _ { i }, \Theta _ { i } ^ { k }, \widetilde { \Psi } _ { i }, r ) \text{ for } i \in \mathcal { I } _ { r },$$

since ˜ y i = q (Θ i , ˜ Ψ ) i . Starting with an initializer Σ 0 , r ∈ R 2 × 2 for r ∈ [ R ] (defined later in (56)), the state evolution parameters are iteratively computed as follows for k ≥ 0 , and r ∈ [ R ] , c ∈ [ C ] :

$$\mu _ { \beta, c } ^ { k + 1 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r e } \mathbb { E } [ \partial _ { 1 } \bar { g } _ { k } ( Z _ { r }, Z _ { r } ^ { k }, \bar { \Psi }, r ) ], \quad ( \sigma _ { \beta, c } ^ { k + 1 } ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r e } \mathbb { E } [ \bar { g } _ { k } ( Z _ { r }, Z _ { r } ^ { k }, \bar { \Psi }, r ) ^ { 2 } ], \\ \Sigma ^ { k + 1, r } = \begin{bmatrix} \Sigma _ { 1 1 } ^ { k + 1, r } & \Sigma _ { 1 2 } ^ { k + 1, r } \\ \Sigma _ { 2 1 } ^ { k + 1, r } & \Sigma _ { 2 2 } ^ { k + 1, r } \end{bmatrix}, \\. \quad \, \quad \, \, \,$$

where ( Z , Z r k r ) ∼ N (0 Σ , k, r ) are independent of ¯ Ψ , and

$$Z _ { r }, Z _ { r } ^ { \kappa } ) & \sim \mathcal { N } ( 0, \Sigma ^ { \kappa, r } ) \text{ are independent of } \Psi, \text{ and} \\ \Sigma _ { 1 1 } ^ { k + 1, r } & = \mathbb { E } [ ( Z _ { r } ) ^ { 2 } ] = \frac { \mathbb { E } [ \bar { \beta } ^ { 2 } ] } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \\ \Sigma _ { 1 2 } ^ { k + 1, r } & = \Sigma _ { 2 1 } ^ { k + 1, r } = \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \mathbb { E } [ \bar { \beta } f _ { k + 1 } ( \mu _ { \beta, c } ^ { k + 1 } \bar { \beta } + G _ { \beta, c } ^ { k + 1 }, c ) ] \\ \Sigma _ { 2 2 } ^ { k + 1, r } & = \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \mathbb { E } [ f _ { k + 1 } ( \mu _ { \beta, c } ^ { k + 1 } \bar { \beta } + G _ { \beta, c } ^ { k + 1 }, c ) ^ { 2 } ]. \\ ^ { k + 1 } _ { a _ { \kappa } ^ { \kappa } } & \sim \mathcal { N } ( 0, ( \sigma _ { a _ { \kappa } ^ { k + 1 } } ) ^ { 2 } ) \text{ is independent of } \bar { \beta }.$$

Here G k +1 β, c ∼ N (0 ( , σ k +1 2 β, c ) ) is independent of ¯ β .

The SC-GAMP algorithm and its state evolution equations are similar to those introduced in [35], the only difference being that the SC design matrix ˜ X sc ∈ R n × p is now a generalized spatially coupled matrix instead of the spatially coupled Gaussian one used in [35]. We note that ˜ X sc is not a generalized white noise matrix since it has variance profile S sc ij := R ˜ W r ( ) i , c ( j ) , which is not guaranteed to satisfy the condition max i ∈ [ n ] ∣ ∣ 1 p ∑ p j =1 S sc ij -1 ∣ ∣ → 0 in Definition 5.1. Nevertheless, ˜ X sc can be related to a generalized white noise matrix ˜ X ∈ R n × p defined as follows. For i ∈ [ n , j ] ∈ [ p ] , let:

̸

$$\widetilde { X } _ { i j } \coloneqq \begin{cases} \frac { \widetilde { X } _ { i j } ^ { s c } } { \sqrt { R \cdot \widetilde { W } _ { r ( i ), c ( j ) } } }, & \text{ if } \widetilde { W } _ { r ( i ), c ( j ) } \neq 0, \\ \underset { \text{indep.} } { \sim } \mathcal { N } ( 0, \frac { 1 } { n } ), & \text{ otherwise.} \end{cases}$$

(In the second line of the definition, we could use any sub-Gaussian distribution with mean zero and variance 1 /n instead of N (0 1 , /n ) .) From the construction of ˜ X sc (see below (49)), it follows that ˜ X is a generalized white noise matrix with variance profile S ij = 1 for all ( i, j ) ∈ [ n ] × [ p ] .

The state evolution result for SC-GAMP requires the following assumptions on the model and the algorithm:

(A1) As n, p → ∞ , we have n p → δ . The signal β , initializer ˆ β 0 , and the noise vector ˜ Ψ are independent of X sc , and their empirical distributions have well-defined limits. There exist random variables ¯ β ∼ P ¯ β and ¯ Ψ ∼ P ¯ Ψ with β W → ¯ β and ˜ Ψ W → ¯ Ψ , respectively.

- (A2) As p → ∞ , ( β , β c ˆ ) 0 c W → ( ¯ β, β ¯ ) 0 c almost surely, for joint laws ( ¯ β, β ¯ ) 0 c having finite moments of all orders, for c ∈ [ C ] . Furthermore, multivariate polynomials are dense in the real L 2 -spaces of functions f : R → R and g : R 2 → R with the inner-products

$$\langle f, \tilde { f } \rangle \coloneqq \mathbb { E } [ f ( \bar { \Psi } ) \tilde { f } ( \bar { \Psi } ) ] \ \text{ and } \ \langle g, \tilde { g } \rangle \coloneqq \mathbb { E } [ g ( \bar { \beta }, \bar { \beta } _ { \mathfrak { c } } ^ { 0 } ) \tilde { g } ( \bar { \beta }, \bar { \beta } _ { \mathfrak { c } } ^ { 0 } ) ].$$

- (A3) For k ≥ 0 and r ∈ [ R ] , c ∈ [ C ] , the functions f k ( · , c ) and ¯ ( g k · , · , · , r ) are each continuous, Lipschitz w.r.t. the first argument, and satisfy the polynomial growth condition in (5) for some order r ≥ 1 .
- (A4) The matrix ˜ X defined as in (55) satisfies ∥ ˜ X ∥ op &lt; C for some constant C , and for any fixed polynomial function f † : R 2 R +2 → R , as n, p →∞ ,

$$\max _ { i \in [ n ] } \Big | & \Big < f ^ { \dagger } \Big ( \beta _ { \epsilon }, \beta _ { \epsilon } \sqrt { \widetilde { R W } _ { 1 } }, \dots, \beta _ { \epsilon } \sqrt { \widetilde { R W } _ { \epsilon } }, \hat { \beta } _ { \epsilon } ^ { 0 } \sqrt { \widetilde { R W } _ { 1 } }, \dots, \hat { \beta } _ { \epsilon } ^ { 0 } \sqrt { \widetilde { R W } _ { \epsilon } }, c \Big ) \odot S _ { i, \mathcal { J } } \Big > \\ & - \Big < f ^ { \dagger } \Big ( \beta _ { \epsilon }, \beta _ { \epsilon } \sqrt { \widetilde { R W } _ { 1 } }, \dots, \beta _ { \epsilon } \sqrt { \widetilde { R W } _ { \epsilon } }, \hat { \beta } _ { \epsilon } ^ { 0 } \sqrt { \widetilde { R W } _ { 1 } }, \dots, \hat { \beta } _ { \epsilon } ^ { 0 } \sqrt { \widetilde { R W } _ { \epsilon } }, c \Big ) \Big > \cdot \langle S _ { i, \mathcal { J } } \rangle \Big | \stackrel { a. s } { \to } 0,$$

for all c ∈ [ C ] , where S is the variance profile of ˜ X (see Definition 5.1) and f † acts element-wise on β c .

- (A5) For any fixed polynomial function f ‡ : R 2 → R , as n, p →∞ ,

$$\max _ { j \in [ p ] } \left | \left \langle f ^ { \dagger } \left ( \widetilde { \Psi } _ { r }, r \right ) \odot S _ { \mathcal { I } _ { r }, j } \right \rangle - \left \langle f ^ { \dagger } \left ( \widetilde { \Psi } _ { r }, r \right ) \right \rangle \cdot \left \langle S _ { \mathcal { I } _ { r }, j } \right \rangle \right | \stackrel { a. s. } { \to } 0,$$

for all r ∈ [ R ] , where f ‡ acts element-wise on ˜ Ψ r .

The state evolution recursion in (53)-(54) is initialized with

$$\Sigma ^ { 0, r } = \left [ \begin{smallmatrix} \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \mathbb { E } [ \bar { \beta } ^ { 2 } ] \\ \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \mathbb { E } [ \bar { \beta } ^ { 0 } \bar { \beta } ] \\ \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \mathbb { E } [ ( \bar { \beta } ^ { 0 } _ { c } ) ^ { 2 } ] \end{smallmatrix} \right ], \quad r \in [ R ].$$

Theorem 5.2 (State evolution for SC-GAMP) . Consider the GLM in (48) with spatially coupled design ˜ X sc defined via a base matrix satisfying (49) , and signal estimation using the SC-GAMP recursion in (50) . Let Assumptions (A1)-(A5) be satisfied, and assume σ 1 β, c &gt; 0 for c ∈ [ C ] . Then for each k ≥ 0 , we have

$$\left ( \beta _ { \epsilon }, \beta _ { \epsilon } ^ { k + 1 } \right ) \stackrel { W _ { 2 } } { \to } \left ( \bar { \beta }, \mu _ { \beta, \epsilon } ^ { k + 1 } \bar { \beta } + G _ { \beta, \epsilon } ^ { k + 1 } \right ), \ \left ( \Theta _ { r }, \Theta _ { r } ^ { k } \right ) \stackrel { W _ { 2 } } { \to } \left ( Z _ { r }, Z _ { r } ^ { k } \right ),$$

almost surely as n, p →∞ with n/p → δ .

The proof is given in Appendix A.1. We now use Theorem 5.2 to prove Theorem 3.2.

## 5.2 Proof of Theorem 3.2 using Theorem 5.2

We first verify that rescaled QGT model

$$\tilde { y } _ { i } = q \left ( \left ( \widetilde { X } _ { i, \colon } ^ { s c } \right ) ^ { \top } \beta, \, \widetilde { \Psi } _ { i } \right ) = \left ( \widetilde { X } _ { i, \colon } ^ { s c } \right ) ^ { \top } \beta + \widetilde { \Psi } _ { i }$$

is a special case of the GLM (48) with a generalized spatially coupled design ˜ X sc constructed as described below (49). The rescaled QGT design ˜ X sc in (10) has independent zero mean entries with variances ˜ W r ( ) i c ( j ) n/R for i ∈ [ n , j ] ∈ [ p ] , where from (18) we have ˜ W rc = 1 /ω for c ≤ r ≤ c + ω -1 , and 0 otherwise. Moreover, this ˜ W satisfies the conditions in (49).

Next, we show that the SC-AMP in (15) is a special case of the SC-GAMP algorithm by choosing

$$f _ { k } ( \beta _ { j } ^ { k } \,, c ) = \mathbb { E } [ \bar { \beta } \, | \, \mu _ { \beta, c } ^ { k } \bar { \beta } + G _ { \beta, c } ^ { k } = \beta _ { j } ^ { k } ], \ \ j \in \mathcal { J } _ { c }.$$

$$g _ { k } ( \Theta _ { i } ^ { k }, \tilde { y } _ { i }, r ) = \frac { \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \cdot \colon } { \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } }, \quad \text{for $i\in \mathcal{I}_{r}$.}$$

The choices in (58) and (59) are based on the Bayes-optimal denoisers for an i.i.d. design and Gaussian noise ¯ Ψ ∼ N (0 , σ 2 ) (see [31, Section 4.2]), and take into account the block-wise dependence of the state evolution parameters. With this choice of f k , in (54) we have E [ ¯ βf k ] = E [ f 2 k ] which implies that Σ k, r 12 = Σ k, r 21 = Σ k, r 22 , for k ≥ 1 .

With our choice of denoisers, the iterate ̂ R k ∈ R n in (50) can be written as

$$\widehat { R } ^ { k } = Q ^ { k } \odot ( \tilde { y } - \Theta ^ { k } ),$$

where the entries of Q k ∈ R n are Q k i = Σ ( k, r ( ) i 11 -Σ k, r ( ) i 12 + σ 2 ) -1 , for i ∈ [ n ] . For j ∈ [ p ] , we have:

$$c _ { j } ^ { k } & = \sum _ { r = 1 } ^ { R } \frac { \widetilde { W } _ { \text{rc} } } { n / R } \sum _ { i \in \mathcal { I } _ { r } } \partial _ { 1 } g _ { k } ( \Theta _ { i } ^ { k }, \tilde { \jmath } _ { y }, r ) = \frac { R } { n } \sum _ { r = 1 } ^ { R } \widetilde { W } _ { \text{rc} } \sum _ { i \in \mathcal { I } _ { r } } \frac { - 1 } { \Sigma _ { 1 1 } ^ { k, r ( i ) } - \Sigma _ { 1 2 } ^ { k, r ( i ) } + \sigma ^ { 2 } } \\ & = \frac { R } { n } \sum _ { r = 1 } ^ { R } \widetilde { W } _ { \text{rc} } \cdot \frac { n } { R } \cdot \frac { - 1 } { \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } } = - \sum _ { r = 1 } ^ { R } \widetilde { W } _ { \text{rc} } \left ( \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } \right ) ^ { - 1 }.$$

Next, we have

$$( \sigma _ { \beta, c } ^ { k + 1 } ) ^ { 2 } & = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \mathbb { E } [ \bar { g } _ { k } ( Z _ { r }, Z _ { r } ^ { k }, \bar { \Psi }, r ) ^ { 2 } ] \stackrel { ( a ) } { = } \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \mathbb { E } \left [ \frac { ( Z _ { r } + \bar { \Psi } - Z _ { r } ^ { k } ) ^ { 2 } } { ( \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } ) ^ { 2 } } \right ] \\ & \stackrel { ( b ) } { = } \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \frac { \sigma ^ { 2 } + \mathbb { E } [ ( Z _ { r } - Z _ { r } ^ { k } ) ^ { 2 } ] } { ( \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } ) ^ { 2 } } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } ( \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } ) ^ { - 1 } \stackrel { ( c ) } { = } - c _ { j } ^ { k }, \\. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad. \quad.$$

where (a) applies (59), (b) uses the independence between ( Z , Z r k r ) and ¯ Ψ , and (c) uses (61). Substituting the definitions of Σ k, r 11 and Σ k, r 12 in (54) into (62), we get

$$( \sigma _ { \beta, \epsilon } ^ { k + 1 } ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r \epsilon } \left ( \sigma ^ { 2 } + \frac { 1 } { \delta _ { \ln } } \sum _ { \epsilon ^ { \prime } = 1 } ^ { C } \widetilde { W } _ { r \epsilon ^ { \prime } } \mathbb { E } \left [ ( \bar { \beta } - f _ { k } ( \mu _ { \beta, \epsilon } ^ { k } \bar { \beta } + G _ { \beta, \epsilon } ^ { k }, \epsilon ) ) ^ { 2 } \right ] \right ) ^ { - 1 }.$$

We also have the identity

$$\mu _ { \beta, \epsilon } ^ { k + 1 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { \text{rc} } \mathbb { E } [ \partial _ { 1 } \bar { g } _ { k } ( Z, Z _ { r } ^ { k }, \bar { \Psi }, r ) ] \stackrel { ( a ) } { = } \sum _ { r = 1 } ^ { R } \widetilde { W } _ { \text{rc} } ( \Sigma _ { 1 1 } ^ { k, r } - \Sigma _ { 1 2 } ^ { k, r } + \sigma ^ { 2 } ) ^ { - 1 } \stackrel { ( b ) } { = } ( \sigma _ { \beta, \epsilon } ^ { k + 1 } ) ^ { 2 },$$

where (a) uses (59) and ¯ y r = Z r + ¯ Ψ , and (b) uses the last equality in (62).

Letting χ k c = σ k β, c , we observe that the update equations in (63)-(64) match the state evolution recursion of the SC-AMP algorithm in (21). Then, substituting (58)-(63) into SC-GAMP in (50), followed by a change of variables from ˜ Θ k to ˜ Θ := ˜ k y -Θ k , gives us the SC-AMP algorithm in (15). Finally, we check that the assumptions of Theorem 5.2 are satisfied:

- · Assumptions (A1) and (A2) hold due to the model assumptions in Section 2, the noise scaling assumption (p. 9), and the SC-AMP initialization ˆ β 0 = E [ ¯ ]1 β p . Recalling that ¯ β ∼ Bernoulli ( π ) for the QGT model, the state evolution initialization in (56) becomes

$$\Sigma ^ { 0, r } = \frac { 1 } { \delta _ { \text{in} } } \left [ \begin{matrix} \sum _ { c = 1 } ^ { \text{C} } \widetilde { W } _ { r c } \pi & \sum _ { c = 1 } ^ { \text{C} } \widetilde { W } _ { r c } \pi ^ { 2 } \\ \sum _ { c = 1 } ^ { \text{C} } \widetilde { W } _ { r c } \pi ^ { 2 } & \sum _ { c = 1 } ^ { \text{C} } \widetilde { W } _ { r c } \pi ^ { 2 } \end{matrix} \right ], \quad r \in [ R ].$$

Using this in (63), we obtain that σ 1 β, c = χ 1 c , where the latter is defined in (22).

- · (A3) . With ¯ β ∼ Bernoulli ( π ) , the denoiser f k ( · , c ) in (58) can be explicitly computed (see (116)), and the choice for ¯ ( g k · , · , · , r ) is given by (59) and (52). From these expressions, it can be verified that both functions are continuous, Lipschitz w.r.t. the first argument, and satisfy the polynomial growth condition with r = 2 .
- · (A4) and (A5) . Recalling the definition of ˜ X in (55) and of ˜ X sc in (10), we note the matrix √ nX ˜ has independent sub-Gaussian entries of variance 1. Using a concentration inequality for the operator norm of sub-Gaussian matrices [59, Theorem 4.4.5] together with the BorelCantelli lemma, we obtain that ∥ ˜ X ∥ op &lt; C almost surely for sufficiently large p . Since the variance profile S ij = 1 for all ( i, j ) , the second condition in (A3) is trivially satisfied. Assumption (A5) is similarly satisfied.

This completes the proof.

## 6 Proof of Theorem 3.5

## 6.1 Proof of (32)

The idea is to rewrite the SC-AMP state evolution in (21) in terms of a general coupled recursion analyzed by Yedla et al. in [37]. We then apply the fixed point characterization of [37, Theorem 1] to the SC-AMP state evolution to obtain (32).

General coupled recursion [37]. Let X = [0 , x max ] , Y = [0 , y max ] with x max , y max ∈ (0 , ∞ ) . Let f : Y → X be a non-decreasing C 1 function, and let g : X → Y be a strictly increasing C 2 function with y max = g x ( max ) . (We say a function f : Z → R is C d if its d th derivative exists and

is continuous on Z .) Consider a matrix A ∈ R C × R with R = C + ω -1 , whose entries are as defined as follows, for r ∈ [ R ] , c ∈ [ C ] :

$$A _ { \text{cr} } = \begin{cases} \frac { 1 } { \omega } & \text{if $c\leq r\leq c+\omega-1$} \\ 0 & \text{otherwise.} \end{cases}$$

Using A , we define the following coupled recursion. For r ∈ [ R ] :

$$y _ { r } ^ { k + 1 } = g ( x _ { r } ^ { k } ), \ \ x _ { r } ^ { k + 1 } = \sum _ { c = 1 } ^ { C } A _ { c r } f \left ( \sum _ { r ^ { \prime } = 1 } ^ { R } A _ { c r ^ { \prime } } y _ { r ^ { \prime } } ^ { k + 1 } \right ),$$

The recursion is initialized with x 0 r = x max for r ∈ [ R ] . This initialization, along with the monotonicity of f and g , ensures that the coupled recursion converges to a fixed point [37]. The fixed point { lim k →∞ x k r } r ∈ [ R ] is characterized by the lemma below in terms of the following potential function:

$$V ( x ) \coloneqq x g ( x ) - \int _ { 0 } ^ { x } g ( z ) d z - \int _ { 0 } ^ { g ( x ) } f ( z ) d z.$$

Lemma 6.1. [37, Theorem 1] For any γ &gt; 0 , there exists ω 0 &lt; ∞ such that for all ω &gt; ω 0 and C ∈ [1 , ∞ ] , the fixed point x ∞ r := lim k →∞ x k r , for r ∈ [ R ] , of the coupled recursion in (65) satisfies the upper bound

$$\max _ { r \in [ R ] } x _ { r } ^ { \infty } \leq \max \left \{ \underset { x \in \mathcal { X } } { \arg \min V ( x ) } \right \} + \gamma.$$

Analyzing state evolution using Lemma 6.1. Let us define the function

$$\ m m s e ( s ) = \mathbb { E } \left \lceil \left ( \bar { \beta } - \mathbb { E } [ \bar { \beta } \, | \, \sqrt { s } \, \bar { \beta } + G ] \right ) ^ { 2 } \right \rceil,$$

where G ∼ N (0 1) , is independent of ¯ β . Then, recalling the definition of f k from (19), the state evolution recursion in (21) is:

$$( \chi _ { c } ^ { k + 1 } ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \underbrace { \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r c } \underbrace { \ m m s e ( ( \chi _ { c } ^ { k } ) ^ { 2 } ) } _ { = \colon \phi _ { r } ^ { k } } } _ { = \colon \phi _ { r } ^ { k } } ) ^ { - 1 }. \\ \cdot \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }. \hat { \ }.$$

Using the definitions above, the state evolution recursion can be rewritten as:

$$\phi _ { r } ^ { k } = \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \sum _ { c = 1 } ^ { C } \widetilde { W } _ { r \epsilon } \psi _ { c } ^ { k }, \quad \psi _ { c } ^ { k + 1 } = \text{mmse} \left ( \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r \epsilon } ( \phi _ { r } ^ { k } ) ^ { - 1 } \right ),$$

which combined into one equation gives:

$$\phi _ { r } ^ { k + 1 } = \sigma ^ { 2 } + \frac { 1 } { \delta _ { \text{in} } } \underbrace { \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r } \text{mmse} \left ( \sum _ { r ^ { \prime } = 1 } ^ { R } \widetilde { W } _ { r ^ { \prime } \text{c} } \left ( \phi _ { r ^ { \prime } } ^ { k } \right ) ^ { - 1 } \right ) } _ { = \colon x _ { r } ^ { k + 1 } }.$$

Rewriting the recursion in terms of x k +1 r defined above, we get:

$$x _ { r } ^ { k + 1 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r \epsilon } \text{mmse} \left ( \sum _ { r ^ { \prime } = 1 } ^ { R } \widetilde { W } _ { r ^ { \prime } \epsilon } \left ( \sigma ^ { 2 } + \frac { x _ { r } ^ { k } } { \delta _ { \text{in} } } \right ) ^ { - 1 } \right ), \ \ r \in [ R ].$$

The modified recursion in (70) is an instance of the coupled recursion in (65), which can be seen by taking A = ˜ W ⊤ and

$$f ( y ) = \text{mmse} \left ( \frac { 1 } { \sigma ^ { 2 } } - y \right ), \ \ g ( x ) = \frac { 1 } { \sigma ^ { 2 } } - \frac { 1 } { \sigma ^ { 2 } + x / \delta }.$$

It is shown in [37, Section VI.E] that with these functions, which satisfy the assumptions stated at the start of this section, the potential function V ( x ) in (66) equals U b δ ( ; in ) defined in (31) (upto an additive constant). Invoking Lemma 6.1, we have that the fixed points of (70), denoted by ( x ∞ r ) r ∈ [ R ] satisfy:

$$\max _ { r \in [ R ] } x _ { r } ^ { \infty } \leq \max \left \{ \underset { x \in [ 0, \text{mmse} ( 0 ) ] } { \arg \min } V ( x ) \right \} + \gamma = \max \left \{ \underset { b \in [ 0, V \text{ar} ( \bar { \beta } ) ] } { \arg \min } U ( b ; \delta _ { \text{in} } ) \right \} + \gamma,$$

where the last equality uses the fact that mmse (0) = Var[ ¯ ] β .

We now use the bound on x ∞ r to upper bound the asymptotic MSE. Using (25), the asymptotic MSE (as k → ∞ ) can be written as 1 C ∑ C c =1 mmse ( ( χ ∞ c ) 2 ) which can be further written as 1 C ∑ C c =1 ψ ∞ c using (69). From (69) and (70), we can write x ∞ r = ∑ C c =1 ˜ W ψ rc ∞ c , which can be written more explicitly as

![Image](https://ac.marcodsn.me/arxiv-markdown-images/2408.00385v1-figure-e3519839-64b0-481f-8723-a5dd5fcddee8.jpg)

$$\nmid \geq c = 1 + c \, \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { m } { \dots } { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } { \stackrel { \dots } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } ) } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } { } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } \ } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } _ } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } \\ } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } }. } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } & } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } | } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } ( } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } }, } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } [ } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } ; } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } } = \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \$$

For notational convenience, let us denote x ∗ = max { argmin b ∈ [0 Var( ¯ ) , β U b δ ( ; in ) } + γ . From (71), we have that x ∞ r ≤ x ∗ for all r ∈ [ R ] where R = Λ + ω -1 . In the rightmost vector in (72), we observe that each entry contains the sum of at most ω consecutive terms. This implies that

$$\sum _ { c ^ { \prime } = c } ^ { c + \omega - 1 } \psi _ { c ^ { \prime } } ^ { \infty } \leq x ^ { * } \omega, \ \ c \in [ C ].$$

Recalling that C = Λ and dividing the elements of [ ψ ∞ 1 , . . . , ψ ∞ Λ ] into groups of non-intersecting consecutive terms - with index groups [1 : ω , ] [ ω +1 : 2 ω , . . . , ] [Λ -ω +1 : Λ] - gives us at most

⌈ Λ ω ⌉ disjoint groups, with the sum of each group having an upper bound of x ω ∗ by (73). Hence, the asymptotic MSE can be bounded as

$$\frac { 1 } { C } \sum _ { c = 1 } ^ { C } \psi _ { c } ^ { \infty } \leq \frac { 1 } { \Lambda } \left \lceil \frac { \Lambda } { \omega } \right \rceil x ^ { * } \omega < \frac { 1 } { \Lambda } \left ( \frac { \Lambda } { \omega } + 1 \right ) x ^ { * } \omega = \frac { \Lambda + \omega } { \Lambda } x ^ { * },$$

This completes the proof of the first part of Theorem 3.5.

## 6.2 Proof of (33)

For the i.i.d. design, we have R = C = 1 and ˜ W 11 = 1 , so the state evolution reduces to

$$x ^ { k + 1 } = m m s e \left ( \left ( \sigma ^ { 2 } + \frac { x ^ { k } } { \delta } \right ) ^ { - 1 } \right ),$$

with the initialization x 0 = mmse (0) = Var( ¯ ) β . Since mmse ( s ) is strictly decreasing in s ∈ [0 , ∞ ) , the sequence ( x k ) is monotonically decreasing in k , and since it is bounded below, it converges to a fixed point. Since the recursion is initialized at x 0 = Var( ¯ ) β , the fixed point is given by the largest solution of x = mmse ( ( σ 2 + x δ ) -1 ) . Finally, we observe that the same equation is obtained by setting ∂ U b δ 1 ( ; ) = 0 . This completes the proof of (33).

## 7 Discussion and Future Directions

We have shown that for noiseless QGT and pooled data, a spatially coupled Bernoulli test design with an AMP recovery algorithm achieves almost-exact recovery with n = o p ( ) tests. A key open question is to determine how n scales with p for almost-exact recovery with SC-AMP. Deriving this scaling is beyond the reach of our asymptotic analysis, which requires that n/p → δ &gt; 0 , but recent nonasymptotic analyses of AMP [29,60] might provide tools to address this question, and allow us to compare with information-theoretic bound of n ∗ = γ ∗ p log p (see (4)).

̸

In this paper, the only assumption on the QGT signal vector β is that its empirical distribution converges to a Bernoulli distribution. The items are not required to be independent, and in some applications there may be known correlations between the items. Although the current SC-AMP algorithm does not exploit correlations between the signal entries, it can be adapted to do so, using nonseparable denoising functions [61]. Recent work in Boolean group testing has shown that exploiting correlations or community structure among the items can significantly improve testing efficiency [62, 63]. Exploring how community structure can be harnessed in quantitative group testing is an interesting direction for future work.

Another open question is to determine the number of tests required for exact recovery in the linear regime for an efficient scheme with a random design. We recall that exact recovery requires P [ ˜ = β β ] → 0 as p →∞ , in contrast to the almost-exact recovery criterion in (2).

## A Proof of Theorem 5.2

## A.1 Proof of Theorem 5.2 via Reduction to Abstract Matrix-AMP

We describe an abstract matrix-AMP iteration for which a state evolution result can be established, and then prove Theorem 5.2 by reducing the SC-GAMP algorithm to the abstract matrix-AMP. For k ≥ 0 , the abstract matrix-AMP produces iterates H k +1 ∈ R p × l H and E k +1 ∈ R n × l E as follows:

$$H ^ { k + 1 } = \widetilde { X } ^ { \top } \widehat { R } ^ { k } - \widehat { H } ^ { k } \cdot ( { \mathbf D } ^ { k } ) ^ { \top }, \ \widehat { R } ^ { k } = \tilde { g } _ { k } ( E ^ { k }, \gamma, \mathcal { R } ), \ \ \mathbf D ^ { k } = \frac { 1 } { \mathbf R } \sum _ { r = 1 } ^ { \mathbf R } \mathbb { E } [ \tilde { g } ^ { \prime } _ { k } ( \tilde { E } ^ { k } _ { r }, \tilde { \gamma }, r ) ]$$

$$E ^ { k + 1 } = \widetilde { X } \widehat { H } ^ { k + 1 } - \widehat { R } ^ { k } \cdot ( B ^ { k + 1 } ) ^ { T }, \ \widehat { H } ^ { k + 1 } = \tilde { f } _ { k + 1 } ( H ^ { k + 1 }, \beta, C ), \ \ B ^ { k + 1 } = \frac { 1 } { \delta C } \sum _ { c = 1 } ^ { c } \mathbb { E } [ \tilde { f } ^ { \prime } _ { k + 1 } ( \tilde { H } ^ { k + 1 } _ { \epsilon }, \bar { \beta }, c ) ]$$

$$H ^ { k + 1 } & = \widetilde { X } ^ { \top } \widehat { R } ^ { k } - \widehat { H } ^ { k } \cdot ( { \mathbf D } ^ { k } ) ^ { \top }, \quad \widehat { R } ^ { k } = \tilde { g } _ { k } ( E ^ { k }, \gamma, \mathcal { R } ), \quad { \mathbf D } ^ { k } = \frac { 1 } { R } \sum _ { r = 1 } ^ { \bullet } \mathbb { E } [ \tilde { g } _ { k } ^ { \prime } ( \bar { E } _ { r } ^ { k }, \bar { \gamma }, r ) ] \\ E ^ { k + 1 } & = \widetilde { X } \widehat { H } ^ { k + 1 } - \widehat { R } ^ { k } \cdot ( { \mathbf B } ^ { k + 1 } ) ^ { \top }, \quad \widehat { H } ^ { k + 1 } = \tilde { f } _ { k + 1 } ( H ^ { k + 1 }, \beta, \mathcal { C } ), \quad { \mathbf B } ^ { k + 1 } = \frac { 1 } { \delta \mathbb { C } } \sum _ { \varepsilon - 1 } ^ { \mathbb { C } } \mathbb { E } [ \tilde { f } _ { k + 1 } ^ { \prime } ( \bar { H } _ { \varepsilon } ^ { k + 1 }, \bar { \zeta } }$$

where β ∈ R p , γ ∈ R n , and C and R are defined in (51). The functions ˜ f k +1 : R l H × R × [ C ] → R l E and ˜ g k : R l E × R × [ R ] → R l H act row-wise on their inputs. The joint laws of ( ¯ E ,γ k r ¯) and ( ¯ H k +1 c , β ¯ ) are described later (below (75)). The algorithm is initialized with ̂ H 0 ∈ R p × l H and E 0 = ˜ ̂ XH 0 ∈ R n × l E . We have the following assumptions for the abstract matrix-AMP algorithm.

- (C1) As dimensions p, n →∞ , the ratio n/p → δ &gt; 0 . Furthermore, l E , l H , R , and C are positive integers that do not scale with p as n, p →∞ .
- (C2) Almost surely for all c ∈ [ C ] , as n, p → ∞ , ( β , H c ̂ 0 J c , : ) W → ( ¯ β, H ¯ 0 c ) and γ r W → ¯ γ , with the joint law of ( ¯ β, H ¯ 0 c ) ∈ R × R l E having finite moments of all orders, where J c is defined in (17). Multivariate polynomials are dense in the real L 2 -spaces of functions f : R 2 → R and g : R l E +2 → R with the inner products

$$\langle f, \tilde { f } \rangle \coloneqq \mathbb { E } [ f ( \bar { \gamma }, r ) \tilde { f } ( \bar { \gamma }, r ) ] \ \text{ and } \ \langle g, \tilde { g } \rangle \coloneqq \mathbb { E } [ g ( \bar { \beta }, \bar { H } _ { \mathfrak { c } } ^ { 0 }, \mathfrak { c } ) \tilde { g } ( \bar { \beta }, \bar { H } _ { \mathfrak { c } } ^ { 0 }, \mathfrak { c } ) ],$$

for all r ∈ [ R ] and c ∈ [ C ] .

- (C3) For k ≥ 0 , the functions ˜ f k +1 and ˜ g k are continuous, Lipschitz w.r.t. their first argument, and satisfy the polynomial growth condition in (5) for some order r ≥ 1 .
- (C4) ˜ X is a generalized white noise matrix where ∥ ˜ X ∥ op &lt; C almost surely for sufficiently large n, p for some constant C . For any fixed polynomial functions f † : R l E +2 → R and f ‡ : R 2 → R , as n, p →∞ ,

$$\max _ { i \in [ n ] } \left | \left \langle f ^ { \dagger } ( \beta _ { \mathrm c }, \widehat { H } _ { \mathcal { I } _ { \mathrm c } } ^ { 0 }, \mathrm c ) \odot S _ { i, \mathcal { I } _ { \mathrm c } } \right \rangle - \left \langle f ^ { \dagger } ( \beta _ { \mathrm c }, \widehat { H } _ { \mathcal { I } _ { \mathrm c } } ^ { 0 }, \mathrm c ) \right \rangle \cdot \langle S _ { i, \mathcal { I } _ { \mathrm c } } \rangle \right | \stackrel { a. s. } { \to } 0 \\ \max _ { j \in [ p ] } \left | \left \langle f ^ { \dagger } ( \gamma _ { r }, r ) \odot S _ { \mathcal { I } _ { r }, j } \right \rangle - \left \langle f ^ { \dagger } ( \gamma _ { r }, r ) \right \rangle \cdot \langle S _ { \mathcal { I } _ { r }, j } \right \rangle \right | \stackrel { a. s. } { \to } 0,$$

for all c ∈ [ C ] and r ∈ [ R ] , where S is the variance profile of ˜ X (see Definition 5.1).

State evolution. The state evolution parameters for k ≥ 0 are

$$\Omega ^ { k + 1 } & = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \widehat { \Omega } ^ { k + 1, r } ; \quad \widehat { \Omega } ^ { k + 1, r } = \mathbb { E } [ \tilde { g } _ { k } ( \bar { E } _ { r } ^ { k }, \bar { \gamma }, r ) \tilde { g } _ { k } ( \bar { E } _ { r } ^ { k }, \bar { \gamma }, r ) ^ { \top } ] \in \mathbb { R } ^ { l _ { H } \times l _ { H } } \\ \Pi ^ { k + 1 } & = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \widehat { \Pi } ^ { k + 1, c } ; \quad \widehat { \Pi } ^ { k + 1, c } = \frac { 1 } { \delta } \mathbb { E } [ \tilde { f } _ { k + 1 } ( \bar { H } _ { c } ^ { k + 1 }, \bar { \beta }, c ) \tilde { f } _ { k + 1 } ( \bar { H } _ { c } ^ { k + 1 }, \bar { \beta }, c ) ^ { \top } ] \in \mathbb { R } ^ { l _ { E } \times l _ { E } }, \\ -. & \quad. \quad - \dots \quad - \dots \quad \dots \quad -$$

with ¯ E k r ∼ N (0 , Π ) k independent of ¯ γ , and ¯ H k +1 c ∼ N (0 Ω , k +1 ) independent of ¯ β . The state evolution is initialized with

$$\Pi ^ { 0 } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \widehat { \Pi } ^ { 0, c }, \quad \text{where} \quad \widehat { \Pi } ^ { 0, c } = \frac { 1 } { \delta } \mathbb { E } \left [ \bar { H } _ { c } ^ { 0 } ( \bar { H } _ { c } ^ { 0 } ) ^ { \top } \right ].$$

Theorem A.1 (State evolution for abstract matrix-AMP) . Consider the abstract matrix-AMP in (74) with the assumptions (C1)-(C4) being satisfied. For k ≥ 1 , and for r ∈ [ R ] , c ∈ [ C ] , the iterates of the abstract matrix AMP satisfy

$$( H ^ { k } _ { \mathcal { J } _ { c }, \cdot }, \beta _ { c } ) \stackrel { W _ { 2 } } { \to } ( \bar { H } ^ { k } _ { c }, \bar { \beta } ), \ \ ( E ^ { k } _ { \mathcal { I } _ { r }, \cdot }, \gamma _ { r } ) \stackrel { W _ { 2 } } { \to } ( \bar { E } ^ { k } _ { r }, \bar { \gamma } ),$$

where ¯ H k c is independent of ¯ β , and ¯ E k r is independent of ¯ γ .

Theorem A.1 is proved in Section A.2.

Proof of Theorem 5.2 using Theorem A.1. We reduce the SC-GAMP algorithm to the abstract matrix-AMP iteration. As given in (55), we can obtain ˜ X from ˜ X sc , which is a generalized white noise matrix (see Definition 5.1) with variance profile S ij = 1 for all ( i, j ) ∈ [ n ] × [ p ] . Next, we set γ := Ψ ˜ , the same β for both algorithms, and the functions ˜ f k : R C × R × [ C ] → R 2 R and ˜ g k : R 2 R × R × [ R ] → R C as follows:

$$\tilde { f } _ { k } ( H ^ { k } _ { j, \cdot }, \beta _ { j }, c ) = \left \lceil \beta _ { j } \left ( \sqrt { R \widetilde { W } _ { 1 c } }, \dots, \sqrt { R \widetilde { W } _ { R c } } \right ), \ f _ { k } ( H ^ { k } _ { j c } + \mu ^ { k } _ { \beta, c } \beta _ { j }, c ) \left ( \sqrt { R \widetilde { W } _ { 1 c } }, \dots, \sqrt { R \widetilde { W } _ { R c } } \right ) \right \rceil,$$

for j ∈ J c and H k j, : ∈ R C (i.e., l H = C ). We also set

$$\tilde { g } _ { k } ( E _ { i, \cdot } ^ { k }, \gamma _ { i }, r ) = g _ { k } ( E _ { i r } ^ { k }, \, q ( E _ { i, r + R }, \widetilde { \Psi } _ { i } ), \, r ) \left ( \sqrt { R \widetilde { W } _ { r 1 } }, \dots, \sqrt { R \widetilde { W } _ { r C } } \right ),$$

for i ∈ I r and E k i, : ∈ R 2 R (i.e., l E = 2 R ). The abstract matrix-AMP iteration is initialized with

$$\widehat { H } _ { j, \cdot } ^ { 0 } = \left [ \beta _ { j } \left ( \sqrt { R \widetilde { W } _ { 1 c } }, \dots, \sqrt { R \widetilde { W } _ { R c } } \right ), \, \hat { \beta } _ { j } ^ { 0 } \left ( \sqrt { R \widetilde { W } _ { 1 c } }, \dots, \sqrt { R \widetilde { W } _ { R c } } \right ) \right ], \quad \text{for } j \in \mathcal { J } _ { c }.$$

The state evolution parameters Π k ∈ R 2 R × 2 R and Ω k ∈ R C × C are recursively computed as follows. We have ¯ E k r ∼ N (0 , Π ) k independent of ¯ Ψ , and the entries of ̂ Ω k +1 , r are

$$\widehat { \Omega } _ { \text{cc} ^ { \prime } } ^ { k + 1, r } = \text{R} \cdot \mathbb { E } \left [ g _ { k } ( \bar { E } _ { r, r } ^ { k }, q ( \bar { E } _ { r, r + R }, \bar { \Psi } ), r ) ^ { 2 } \right ] \sqrt { \widetilde { W } _ { \text{cc} } \widetilde { W } _ { \text{cc} ^ { \prime } } }, \quad \text{for } c, c ^ { \prime } \in [ C ].$$

Next, we have ¯ H k c ∼ N (0 Ω ) , k independent of ¯ β and

$$\text{iveA, we move } \iota _ { c } \stackrel { \sim } { \sim } \mathcal { I } \left ( \mathcal { I }, \mathcal { I } \right ) \text{unique sequence on } \beta \text{ and } \\ \widehat { \Pi } _ { r s } ^ { k + 1, c } = \begin{cases} \frac { 1 } { \delta } \mathbb { R } \mathbb { E } [ \bar { \beta } ^ { 2 } ] \sqrt { \widetilde { W } _ { r c } \widetilde { W } _ { s c } } & r, s \in [ R ] \\ \frac { 1 } { \delta } \mathbb { R } \mathbb { E } [ \bar { \beta } f _ { k } ( \{ \bar { H } _ { c } ^ { k } \} _ { c } + \mu _ { \beta, c } ^ { k + 1 } \bar { \beta }, c ) ] \sqrt { \widetilde { W } _ { r c } \widetilde { W } _ { ( r - R ) c } } & r \in [ R ], R + 1 \leq r \leq 2 R, s \in [ R ] \\ \frac { 1 } { \delta } \mathbb { R } [ \bar { \beta } f _ { k } ( \{ \bar { H } _ { c } ^ { k } \} _ { c } + \mu _ { \beta, c } ^ { k + 1 } \bar { \beta }, c ) ] \sqrt { \widetilde { W } _ { ( r - R ) c } \widetilde { W } _ { s c } } & R + 1 \leq r \leq 2 R, s \in [ R ] \\ \frac { 1 } { \delta } \mathbb { R } [ f _ { k } ( \{ \bar { H } _ { c } ^ { k } \} _ { c } + \mu _ { \beta, c } ^ { k + 1 } \bar { \beta }, c ) ^ { 2 } ] \sqrt { \widetilde { W } _ { ( r - R ) c } \widetilde { W } _ { ( s - R ) c } } & R + 1 \leq r, s \leq 2 R. \end{cases}$$

The state evolution is initialized with

$$\widehat { \Pi } ^ { 0, \mathfrak { c } } = \frac { 1 } { \delta } \lim _ { p \to \infty } \frac { 1 } { p / \mathbb { C } } ( \widehat { H } ^ { 0 } _ { \mathcal { J } _ { \mathfrak { c } }, \cdot } ) ^ { \top } \widehat { H } ^ { 0 } _ { \mathcal { J } _ { \mathfrak { c } }, \cdot },$$

for c ∈ [ C ] , with ̂ H 0 J c , : ∈ R p/ C × 2 R . By assumption (A2) (see (56)), the entries of ̂ Π 0 , c are given by

$$\widehat { \Pi } _ { r s } ^ { 0, c } & = \begin{cases} \frac { 1 } { \delta } R \mathbb { E } [ \bar { \beta } ^ { 2 } ] \sqrt { \widetilde { W } _ { r c } \widetilde { W } _ { s c } } & r, s \in [ R ] \\ \frac { 1 } { \delta } \mathbb { E } [ \bar { \beta } \bar { \beta } _ { c } ^ { 0 } ] \sqrt { \widetilde { W } _ { r c } \widetilde { W } _ { ( s - R ) c } } & r \in [ R ], R + 1 \leq s \leq 2 R \\ \frac { 1 } { \delta } \mathbb { E } [ \bar { \beta } \bar { \beta } _ { c } ^ { 0 } ] \sqrt { \widetilde { W } _ { ( r - R ) c } \widetilde { W } _ { s c } } & R + 1 \leq s \leq 2 R, s \in [ R ] \\ \frac { 1 } { \delta } \mathbb { E } [ ( \bar { \beta } _ { c } ^ { 0 } ) ^ { 2 } ] \sqrt { \widetilde { W } _ { ( r - R ) c } \widetilde { W } _ { ( s - R ) c } } & R + 1 \leq s, r \leq 2 R. \end{cases} \\ \int \int \int \frac { 1 } { \delta } \equiv c \quad \widehat { \frac { \widehat { \Pi } } { \delta } } c$$

We then have Π = 0 1 ∑ C ̂ Π 0 , c .

C c =1 We can then show that

$$\begin{bmatrix} \Pi _ { r r } ^ { k } & \Pi _ { r ( r + R ) } ^ { k } \\ \Pi _ { ( r + R ) r } ^ { k } & \Pi _ { ( r + R ) ( r + R ) } \end{bmatrix} = \Sigma ^ { k, r } \ \text{ and } \ \Omega _ { \text{cc} } ^ { k + 1 } = ( \sigma _ { \beta, \text{c} } ^ { k + 1 } ) ^ { 2 },$$

implying that ( { ¯ E k r } r , { ¯ E k r } r + R ) d = ( Z , Z r k r ) and { ¯ H k c } c d = G k β, c . We can also show that for k ≥ 0 ,

$$& E _ { i, r } ^ { k } = \Theta _ { i } ^ { k }, \quad E _ { i, r + R } ^ { k } = \Theta _ { i }, \quad \text{for $i\in\mathcal{I}_{r}$, $r\in[R]$} \\ & H _ { j, c } ^ { k + 1 } + \mu _ { \beta, c } ^ { k + 1 } \beta _ { j } = \beta _ { j } ^ { k + 1 }, \quad \text{for $j\in\mathcal{J}_{c}$, $c\in[C]$.}$$

Both (77) and (78) are shown using steps identical to those in [35, Section 5.1.2], so we omit repeating the proof for brevity. Theorem 5.2 follows by using (77) and (78) in Theorem A.1.

## A.2 Proof of Theorem A.1 via Reduction to U-AMP

Theorem A.1 is proved by reducing the abstract matrix-AMP recursion to the U-AMP recursion which is defined as follows. Given a generalized white noise matrix ˜ X , for t ≥ 1 , the iterates of U-AMP, denoted by h t ∈ R p and e t ∈ R n , are produced using functions f v t : R t + L d +1 → R , f u t +1 : R t + L c +1 → R . Given an initializer u 1 ∈ R n , side information vectors c 1 , . . . , c L c ∈ R n and d , . . . , d 1 L d ∈ R p , all independent of ˜ X , the iterates of the U-AMP recursion are computed as:

$$, \dots, \infty \, & \ < a \,, \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \min \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \max \, \text{$\mathcal{ } } \\ h ^ { t } & = \sqrt { \delta } \widetilde { X } ^ { \top } u ^ { t } - \sum _ { s = 1 } ^ { t } b ^ { t } _ { s } v ^ { s }, \quad \ \ \ \ \ v ^ { t } = f ^ { v } _ { t + 1 } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { L _ { d } }, \mathcal { C } ) \\ e ^ { t } & = \sqrt { \delta } \widetilde { X } v ^ { t } - \sum _ { s = 1 } ^ { t } a ^ { t } _ { s } u ^ { s }, \quad \ \ \ u ^ { t + 1 } = f ^ { u } _ { t + 1 } ( e ^ { 1 }, \dots, e ^ { t }, c ^ { 1 }, \dots, e ^ { L _ { c } }, \mathcal { R } ),$$

where C R , were defined in (51), and the functions f v t and f u t +1 act row-wise. The coefficients a t s and b t s are defined later in (83) in terms of state evolution parameters.

Recalling the notification simplification for sub-blocks of vectors presented in the paragraph below (17), we have the following assumptions:

- (D1) As n, p → ∞ , we have n/p = δ &gt; 0 , for fixed L c and L d . Furthermore, for all c ∈ [ C ] and r ∈ [ R ] ,

$$( u _ { r } ^ { 1 }, c _ { r } ^ { 1 }, \dots, c _ { r } ^ { L _ { c } } ) \stackrel { W } { \to } ( \bar { u } _ { r } ^ { 1 }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } } ) \text{ and } ( d _ { c } ^ { 1 }, \dots, d _ { c } ^ { L _ { d } } ) \stackrel { W } { \to } ( \bar { d } _ { c } ^ { 1 }, \dots, \bar { d } _ { c } ^ { L _ { d } } ),$$

for joint limit laws (¯ u , c 1 r ¯ 1 r , . . . , c ¯ L c r ) and ( ¯ d , . . . , d 1 c ¯ L d c ) having finite moments of all orders, where E [(¯ u 1 ) 2 ] ≥ 0 . Multivariate polynomials are dense in the real L 2 -spaces of functions f : R L c +1 → R and g : R L d → R with the inner products

$$\langle f, \tilde { f } \rangle \coloneqq \mathbb { E } [ f ( \bar { u } _ { r } ^ { 1 }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } } ) \tilde { f } ( \bar { u } _ { r } ^ { 1 }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } } ) ] \text{ and } \langle g, \tilde { g } \rangle \coloneqq \mathbb { E } [ g ( \bar { d } _ { c } ^ { 1 }, \dots, \bar { d } _ { c } ^ { L _ { d } } ) \tilde { g } ( \bar { d } _ { c } ^ { 1 }, \dots, \bar { d } _ { c } ^ { L _ { d } } ) ].$$

- (D2) Each function f v t : R t + L d +1 → R and f u t +1 : R t + L c +1 → R is continuous, is Lipschitz in its first t arguments, and satisfies the polynomial growth condition in (5) for some order r ≥ 1 .
- (D3) ∥ ˜ X ∥ op &lt; C , for some constant C almost surely for all sufficiently large n and p .
- (D4) For any fixed polynomial functions f † : R L c +1 → R and f ‡ : R L d → R , almost surely as n, p →∞ ,

$$\max _ { j \in [ p ] } & \left | \langle f ^ { \dagger } ( u ^ { 1 }, c ^ { 1 }, \dots, c ^ { L _ { c } } ) \odot S _ { ;, j } \rangle - \langle f ^ { \dagger } ( u ^ { 1 }, c ^ { 1 }, \dots, c ^ { L _ { c } } ) \rangle \cdot \langle S _ { ;, j } \rangle \right | \rightarrow 0 \\ & \max _ { i \in [ n ] } \left | \langle f ^ { \dagger } ( d ^ { 1 }, \dots, d ^ { L _ { d } } ) \odot S _ { i, ; } \rangle - \langle f ^ { \dagger } ( d ^ { 1 }, \dots, d ^ { L _ { d } } ) \rangle \cdot \langle S _ { i, ; } \rangle \right | \rightarrow 0.$$

State evolution. The state evolution result below states that the joint empirical distribution of ( h , . . . , h 1 c t c ) converges to a Gaussian law N (0 Ξ ) , t , for c ∈ [ C ] . Similarly, the joint empirical distribution of ( e , . . . , e 1 r t r ) converges to N (0 , Γ ) t , for r ∈ [ R ] . The covariance matrices Ξ t , Γ t ∈ R t × t are iteratively defined as follows, starting from Ξ 1 = δ R ∑ R r =1 E [(¯ u 1 r ) 2 ] . Given Ξ t , for t ≥ 1 , let ( ¯ h , . . . , h 1 ¯ ) t ∼ N (0 , Ξ ) t be independent of ( ¯ d , . . . , d 1 c ¯ L d c ) and define

$$\bar { v } _ { \text{c} } ^ { s } = f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { s }, \bar { d } _ { \text{c} } ^ { 1 }, \dots, \bar { d } _ { \text{c} } ^ { L _ { d } }, \text{c} ), \quad s \in [ t ], \text{c} \in [ \text{C} ].$$

Then we have

$$\Gamma ^ { t } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \widehat { \Gamma } ^ { t, c } \quad \text{where} \quad \widehat { \Gamma } ^ { t, c } = \left ( \mathbb { E } [ \bar { v } _ { c } ^ { r } \bar { v } _ { c } ^ { s } ] \right ) _ { r, s = 1 } ^ { t }.$$

Next, let (¯ e , . . . , e 1 ¯ ) t ∼ N (0 , Γ ) t be independent of (¯ u , c 1 r ¯ 1 r , . . . , c ¯ L c r ) and define

$$\bar { u } _ { r } ^ { s + 1 } = f _ { s + 1 } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { s }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } }, r ), \ \ s \in [ t ], \, r \in [ \mathbf R ].$$

Then, we have

$$\Xi ^ { t + 1 } = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \widehat { \Xi } ^ { t + 1, r } \ \text{ where } \ \widehat { \Xi } ^ { t + 1, r } = \delta \left ( \mathbb { E } [ \bar { u } ^ { r } _ { r } \bar { u } ^ { s } _ { r } ] \right ) ^ { t + 1 } _ { r, s = 1 }.$$

We define the memory coefficients a t s and b t s in (79) as:

$$a _ { s } ^ { t } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \mathbb { E } \left [ \partial _ { s } f _ { t } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } _ { c } ^ { 1 } \dots, \bar { d } _ { c } ^ { L _ { d } }, c ) \right ], \quad b _ { s } ^ { t } = \frac { \delta } { R } \sum _ { r = 1 } ^ { R } \mathbb { E } \left [ \partial _ { s } f _ { t } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } }, r ) \right ],$$

where ∂ s denotes partial derivative in the s th argument. The following result gives the state evolution result for the U-AMP recursion.

Corollary A.2 (State evolution for U-AMP) . Let ˜ X ∈ R n × p be a generalized white noise matrix (as defined in Definition 5.1) with variance profile S ∈ R n × p , and let u , c 1 r 1 r , . . . , c L c r , d 1 c , . . . , d L d c be independent of ˜ X and satisfy Assumptions (D1)-(D4). Further assume that each matrix Ξ t and Γ t is non-singular. Then for any fixed t ≥ 1 , as n, p → ∞ , the iterates of the abstract AMP in (79) almost surely satisfy the following, for r ∈ [ R ] and c ∈ [ C ] :

$$& ( u _ { r } ^ { 1 }, c _ { r } ^ { 1 }, \dots, c _ { r } ^ { L _ { c } }, e _ { r } ^ { 1 }, \dots, e _ { r } ^ { t } ) \stackrel { W _ { 2 } } { \to } ( \bar { u } _ { r } ^ { 1 }, \bar { c } _ { r } ^ { 1 }, \dots, \bar { c } _ { r } ^ { L _ { c } }, \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t } ) \\ & ( d _ { \mathbf c } ^ { 1 }, \dots, d _ { \mathbf c } ^ { L _ { d } }, h _ { \mathbf c } ^ { 1 }, \dots, h _ { \mathbf c } ^ { t } ) \stackrel { W _ { 2 } } { \to } ( \bar { d } _ { \mathbf c } ^ { 1 }, \dots, \bar { d } _ { \mathbf c } ^ { L _ { d } }, \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t } ),$$

where ( ¯ h , . . . , h 1 ¯ ) t ∼ N (0 , Ξ ) t and (¯ e , . . . , e 1 ¯ ) t ∼ N (0 , Γ ) t are independent of ( ¯ d , . . . , d 1 c ¯ L d c ) and (¯ u , c 1 r ¯ 1 r . . . . , c ¯ L c r ) respectively.

Corollary A.2 is obtained from the AMP universality result in [36, Theorem 2.17]. A proof is provided in Section A.3.

Proof of Theorem A.1 using Corollary A.2. We reduce the abstract matrix-AMP iteration to the U-AMP. We set the initializer to be u 1 = 0 , L c = 1 , L d = l E +1 and the side information vectors to be

$$c ^ { 1 } = \gamma, \ \ d ^ { 1 } = \beta, \ \ d ^ { 2 } = \widehat { H } ^ { 0 } _ { \cdot ; 1 }, \ \ \dots, \ \ d ^ { l _ { E } + 1 } = \widehat { H } ^ { 0 } _ { \cdot ; l _ { E } }.$$

We show the reduction through induction.

Base case. We consider the case t = 0 , and our goal is to reduce ̂ H 0 , E 0 , H 1 , ̂ R 0 , ̂ H 1 , and E 1 to iterates of U-AMP defined via careful choices of the functions f v t and f u t +1 . We provide a summary of the reductions before giving their derivations.

- · For t = 1 , . . . , l E : We have

$$h ^ { 1 }, \dots, h ^ { l _ { E } } = 0, \ \ ( v ^ { 1 }, \dots, v ^ { l _ { E } } ) = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 0 }, \ \ ( e ^ { 1 }, \dots, e ^ { l _ { E } } ) = E ^ { 0 }, \ \ u ^ { 1 }, \dots, u ^ { l _ { E } } = 0. \quad \ ( 8 5 )$$

- · For t = l E +1 , . . . , l E + l H : We have

$$& ( h ^ { l _ { E } + 1 }, \dots, h ^ { l _ { E } + l _ { H } } ) = H ^ { 1 }, \quad v ^ { l _ { E } + 1 }, \dots, v ^ { l _ { E } + l _ { H } } = 0 \\ & e ^ { l _ { E } + 1 }, \dots, e ^ { l _ { E } + l _ { H } } = 0, \quad ( u ^ { l _ { E } + 1 }, \dots, u ^ { l _ { E } + l _ { H } } ) = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { 0 }.$$

- · For t = l E + l H +1 , . . . , 2 l E + l H : We have

$$h ^ { l _ { E } + l _ { H } + 1 }, \dots, h ^ { 2 l _ { E } + l _ { H } } = 0, \quad ( v ^ { l _ { E } + l _ { H } + 1 }, \dots, v ^ { 2 l _ { E } + l _ { H } } ) = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 1 } \\ ( e ^ { l _ { E } + l _ { H } + 1 }, \dots, e ^ { 2 l _ { E } + l _ { H } } ) = E ^ { 1 }, \quad u ^ { l _ { E } + l _ { H } + 1 }, \dots, u ^ { 2 l _ { E } + l _ { H } } = 0.$$

We now provide the derivations of (85)-(87). For t = 1 , . . . , l E -1 , we set

$$f _ { t } ^ { v } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { l _ { E } + 1 }, \mathcal { C } ) = \frac { 1 } { \sqrt { \delta } } d ^ { t + 1 }, \ \ f _ { t + 1 } ^ { u } ( e ^ { 1 }, \dots, e ^ { t }, c ^ { 1 }, \mathcal { R } ) = 0.$$

For t = 1 , we have our initializer u 1 = 0 , and using (84) gives

$$h ^ { 1 } = 0, \ \ v ^ { 1 } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 0 } _ { ; ; 1 }, \ \ e ^ { 1 } = E ^ { 0 } _ { ; ; 1 }, \ \ u ^ { 2 } = 0.$$

Following similar steps, for t = 2 , . . . , l E -1 , we have

$$h ^ { t } = 0, \ \ v ^ { t } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 0 } _ { \colon, t }, \ \ e ^ { t } = E ^ { 0 } _ { \colon, t }, \ \ u ^ { t + 1 } = 0.$$

For t = l E , set

$$f _ { l _ { E } } ^ { v } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { l _ { E } + 1 }, \mathcal { C } ) = \frac { 1 } { \sqrt { \delta } } d ^ { l _ { E } + 1 }, \quad f _ { l _ { E } + 1 } ^ { u } ( \underbrace { e ^ { 1 }, \dots, e ^ { l _ { E } } } _ { = E ^ { 0 } }, \underbrace { e ^ { 1 }, \mathcal { R } } _ { = \gamma } = \frac { 1 } { \sqrt { \delta } } \{ \tilde { g } _ { 0 } ( E ^ { 0 }, \gamma, \mathcal { R } ) \} _ { ; 1 },$$

which gives

$$h ^ { l _ { E } } = 0, \ \ v ^ { l _ { E } } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 0 } _ { \colon, l _ { E } }, \ \ e ^ { l _ { E } } = E ^ { 0 } _ { \colon, l _ { E } }, \ \ u ^ { l _ { E } + 1 } = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { 0 } _ { \colon, 1 }.$$

This completes the derivation for (85). For t = l E +1 , . . . , l E + l H -1 , we set

$$f _ { t } ^ { v } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { l _ { E } + 1 }, \mathcal { C } ) = 0, \ \ f _ { t + 1 } ^ { u } ( e ^ { 1 }, \dots, e ^ { t }, c ^ { 1 }, \mathcal { R } ) = \frac { 1 } { \sqrt { \delta } } \{ \tilde { g } _ { 0 } ( E ^ { 0 }, \gamma, \mathcal { R } ) \} _ { ; t + 1 - l _ { E } }.$$

For t = l E +1 , we have the following identity:

$$\text{for } \iota = \iota _ { E } + 1, \, \text{we have one following} \text{tenuity} \colon \\ \{ \widehat { H } ^ { 0 } ( \mathbb { D } ^ { 0 } ) ^ { \top } \} _ { \vdots ; 1 } & = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \left \{ \widehat { H } ^ { 0 } \left ( \mathbb { E } [ \tilde { g } _ { 0 } ^ { \prime } ( \bar { E } _ { r } ^ { 0 }, \bar { \gamma }, r ) ] \right ) ^ { \top } \right \} _ { \vdots, 1 } = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \widehat { H } ^ { 0 } \cdot \begin{bmatrix} \mathbb { E } [ \partial _ { 1 } \tilde { g } _ { 0, 1 } ( \bar { E } _ { r } ^ { 0 }, \bar { \gamma }, r ) ] \\ \vdots \\ \mathbb { E } [ \partial _ { l } \tilde { g } _ { 0, 1 } ( \bar { E } _ { r } ^ { 0 }, \bar { \gamma }, r ) ] \end{bmatrix} \\ & = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \sum _ { s = 1 } ^ { l _ { E } } \widehat { H } ^ { 0 } _ { \vdots, s } \mathbb { E } [ \partial _ { s } \tilde { g } _ { 0, 1 } ( \bar { E } _ { r } ^ { 0 }, \bar { \gamma }, r ) ].$$

Then, we have

$$h ^ { l _ { E } + 1 } & \overset { ( a ) } { = } \sqrt { \delta } \widetilde { X } ^ { \top } u ^ { l _ { E } + 1 } - \sum _ { s = 1 } ^ { l _ { E } } b ^ { l _ { E } + 1 } _ { s } v ^ { s } \overset { ( b ) } { = } \widetilde { X } ^ { \top } \widehat { R } ^ { 0 } _ { \cdot, 1 } - \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \sum _ { s = 1 } ^ { l _ { E } } \mathbb { E } [ \partial _ { s } \tilde { g } _ { 0, 1 } ( \bar { R } ^ { 0 } _ { r }, \bar { \gamma }, r ) ] \widehat { H } ^ { 0 } _ { \cdot, s } \\ & \overset { ( c ) } { = } \widetilde { X } ^ { \top } \widehat { R } ^ { 0 } _ { \cdot, 1 } - \{ \widehat { H } ^ { 0 } ( \mathbb { D } ^ { 0 } ) ^ { \top } \} _ { \cdot, 1 } = H ^ { 1 } _ { \cdot, 1 },$$

where we use (79) in (a), substitute the definitions of u l E +1 , b l E +1 s , and v s in (b), and apply (90) in (c). Next, we have

$$v ^ { l _ { E } + 1 } = 0, \ \ e ^ { l _ { E } + 1 } = 0, \ \ u ^ { l _ { E } + 2 } = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { 0 } _ { \cdot, 2 }.$$

Similarly for t = l E +2 , . . . , l E + l H -1 , we have

$$h ^ { t } = H ^ { 1 } _ { \colon, t - l _ { E } }, \ \ v ^ { t } = 0, \ \ e ^ { t } = 0, \ \ u ^ { t + 1 } = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { 0 } _ { \colon, t + 1 - l _ { E } }.$$

For t = l E + l H , we set f v t = 0 and f u t +1 = 0 , so that

$$h ^ { l _ { E } + l _ { H } } = H ^ { 1 } _ { ;, l _ { H } }, \ \ v ^ { l _ { E } + l _ { H } } = 0, \ \ e ^ { l _ { E } + l _ { H } } = 0, \ \ u ^ { l _ { E } + l _ { H } + 1 } = 0.$$

This completes the derivation for (86).

For t = l E + l H +1 , . . . , 2 l E + l H -1 , we set

$$f _ { t } ^ { v } ( h ^ { 1 }, \dots, h ^ { l _ { E } }, \underbrace { h ^ { l _ { E } + 1 }, \dots, h ^ { l _ { E } + l _ { H } } } _ { = H ^ { 1 } }, &, h ^ { l _ { E } + l _ { H } + 1 }, \dots, h ^ { t }, \underbrace { d ^ { 1 } } _ { = \beta }, d ^ { 2 }, \dots, d ^ { l _ { E } + 1 }, \mathcal { C } ) = \frac { 1 } { \sqrt { \delta } } \{ \tilde { f } _ { 1 } ( H ^ { 1 }, \beta, \mathcal { C } ) \} _ { ; 1 } \\ \tilde { \psi } _ { 1 } ( 1 + 1 ) \, \gg \, \quad \, \, \,$$

$$f _ { t + 1 } ^ { u } ( e ^ { 1 }, \dots, e ^ { t }, c ^ { 1 }, \mathcal { R } ) = 0.$$

For t = l E + l H +1 , we have

$$h ^ { l _ { E } + l _ { H } + 1 } = 0, \ \ v ^ { l _ { E } + l _ { H } + 1 } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 1 } _ { \colon, 1 }.$$

We pause to show an identity:

$$\tt v c { \, \rho a v e c \, \tt v a n w a n w a n w a n w a n }. \\ \{ \widehat { R } ^ { 0 } ( B ^ { 1 } ) ^ { T } \} _ { \, \cdots, 1 } & = \frac { 1 } { \delta C } \sum _ { c = 1 } ^ { C } \left \{ \widehat { R } ^ { 0 } \mathbb { E } [ \tilde { f } ^ { t } _ { 1 } ( \tilde { H } ^ { 1 } _ { c }, \tilde { \beta }, c ) ] ^ { T } \right \} _ { \, \cdots, 1 } = \frac { 1 } { \delta C } \sum _ { c = 1 } ^ { C } \widehat { R } ^ { 0 } \begin{bmatrix} \mathbb { E } [ \partial _ { 1 } \tilde { f } _ { 1, 1 } ( \tilde { H } ^ { 1 } _ { c }, \tilde { \beta }, c ) ] \\ \vdots \\ \mathbb { E } [ \partial _ { 1 } \tilde { f } _ { 1, 1 } ( \tilde { H } ^ { 1 } _ { c }, \tilde { \beta }, c ) ] \end{bmatrix} \\ \tt T h o n \, \ f r o n \, \ w a n h a n$$

Then, from (79), we have

$$e ^ { l _ { E } + l _ { H } + 1 } & = \sqrt { \delta } \widetilde { X } v ^ { l _ { E } + l _ { H } + 1 } - \sum _ { s = 1 } ^ { l _ { E } + l _ { H } + 1 } a _ { s } ^ { l _ { E } + l _ { H } + 1 } u ^ { s } \stackrel { ( a ) } { = } \widetilde { X } \widehat { H } ^ { 1 } _ { ; ; 1 } - \frac { 1 } { \delta C } \sum _ { c = 1 } ^ { C } \sum _ { s = 1 } ^ { l _ { H } } \widehat { R } ^ { 0 } _ { ; ; s } \mathbb { E } [ \partial _ { s } \widetilde { f } _ { 1, 1 } ( \widehat { H } ^ { 1 } _ { c ; ; }, \widehat { \beta }, c ) ] \\ & \stackrel { ( b ) } { = } \widetilde { X } \widehat { H } ^ { 1 } _ { ; ; 1 } - \left \{ \widehat { R } ^ { 0 } ( B ^ { 1 } ) ^ { \top } \right \} _ { ; ; 1 } = E ^ { 1 } _ { ; ; 1 },$$

where (a) uses the definitions of v l E + l H +1 , a l E + l H +1 s and u s , and (b) uses (92). Next, we have u l E + l H +1 = 0 . Similarly, for t = l E + l H +2 , . . . , 2 l E + l H -1 , we have

$$h ^ { t } = 0, \ \ v ^ { t } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 1 } _ { \colon t - l _ { E } - l _ { H } }, \ \ e ^ { t } = E ^ { 1 } _ { \colon ; t - l _ { E } - l _ { H } }, \ \ u ^ { t + 1 } = 0.$$

For t = 2 l E + l H , we set

$$f _ { 2 l _ { E } + l _ { H } } ^ { v } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { l _ { E } + 1 }, \mathcal { C } ) = \frac { 1 } { \sqrt { \delta } } \{ \tilde { f } _ { 1 } ( H ^ { 1 }, \beta, \mathcal { C } ) \} _ { \colon l _ { E } }$$

$$f _ { 2 l _ { E } + l _ { H } + 1 } ^ { u } ( e ^ { 1 }, \dots, e ^ { l _ { E } + l _ { H } }, \underbrace { e ^ { l _ { E } + l _ { H } + 1 }, \dots, e ^ { 2 l _ { E } + l _ { H } } } _ { = E ^ { 1 } }, \underbrace { c ^ { 1 } _ { \Lambda } } _ { = \gamma }, \mathcal { R } ) = \frac { 1 } { \sqrt { \delta } } \left \{ \tilde { g } _ { 1 } ( E ^ { 1 }, \gamma, \mathcal { R } ) \right \} _ { ; 1 }.$$

$$& d ^ { * E ^ { \top \ast } }, \mathcal { C } ) = \frac { } { \sqrt { \delta } } \{ f _ { 1 } ( H ^ { * }, \beta, \mathcal { C } ) \} _ { \colon, l _ { E } } \\ & \frac { _ { L E ^ { + } l _ { H } + 1 } } { _ { = E ^ { 1 } } }, \underbrace { c ^ { 1 } } _ { = \gamma }, \mathcal { R } ) = \frac { 1 } { \sqrt { \delta } } \{ \tilde { g } _ { 1 } ( E ^ { 1 }, \gamma, \mathcal { R } ) \} _ { \colon, 1 }.$$

This gives

$$h ^ { 2 l _ { E } + l _ { H } } = 0, \ \ v ^ { 2 l _ { E } + l _ { H } } = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { 1 } _ { \colon, l _ { E } }, \ \ e ^ { 2 l _ { E } + l _ { H } } = E ^ { 1 } _ { \colon, l _ { E } }, \ \ u ^ { 2 l _ { E } + l _ { H } + 1 } = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { 1 } _ { \colon, 1 },$$

completing the derivation of (87). This concludes the reduction of the abstract matrix-AMP iterates to the U-AMP iterates for the case of t = 0 .

We now show the convergence statements in Theorem A.1 for E , H 0 1 , and E 1 by reducing the abstract matrix-AMP SE parameters to the corresponding U-AMP SE parameters.

Convergence of ( E , γ 0 ) . From Assumption (D1) and Corollary A.2, we have ( d , . . . , d 1 c L d c , h , . . . , h 1 c l E c ) W → ( ¯ d , . . . , d 1 c ¯ L d c , h ¯ 1 , . . . , h ¯ l E ) , for c ∈ [ C ] . Recalling that

$$v _ { \mathfrak { c } } ^ { s } = f _ { s } ^ { v } ( h ^ { 1 }, \dots, h ^ { s }, d _ { \mathfrak { c } } ^ { 1 }, \dots, d _ { \mathfrak { c } } ^ { L _ { d } }, \mathfrak { c } ), \quad \bar { v } _ { \mathfrak { c } } ^ { s } = f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { s }, \bar { d } _ { \mathfrak { c } } ^ { 1 }, \dots, \bar { d } _ { \mathfrak { c } } ^ { L _ { d } }, \mathfrak { c } ), \quad \text{for } s \geq 1,$$

the convergence above implies that ( v , . . . , v 1 c l E c ) W → (¯ v , . . . , v 1 c ¯ l E c ) . Since we have shown in (85) that ( v , . . . , v 1 l E ) = 1 √ δ ̂ H 0 , we must have (¯ v , . . . , v 1 c ¯ l E c ) = 1 √ δ ¯ H 0 c , where the latter is given by Assumption (C2).

For r ∈ [ R ] , Corollary A.2 implies that ( e , . . . , e 1 r l E r ) W 2 → (¯ e , . . . , e 1 ¯ l E ) ∼ N (0 Γ , l E ) , where

$$\Gamma ^ { l _ { E } } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \widehat { \Gamma } ^ { l _ { E }, c }, \quad \widehat { \Gamma } ^ { l _ { E }, c } = \left ( \mathbb { E } [ \bar { v } _ { c } ^ { r } \bar { v } _ { c } ^ { s } ] \right ) _ { r, s = 1 } ^ { l _ { E } }.$$

Using (¯ v , . . . , v 1 c ¯ l E c ) = 1 √ δ ¯ H 0 c in (76), we have that ̂ Π 0 , c = Γ ̂ l E , c and Π = Γ 0 l E . Since ( e , . . . , e 1 r l E r ) = E 0 I r , : from (85), and Corollary A.2 guarantees that ( e , . . . , e 1 r l E r ) W 2 → (¯ e , . . . , e 1 ¯ l E ) ∼ N (0 , Γ l E ) , we have E 0 I r , : W 2 → ¯ E 0 r ∼ N (0 , Π ) 0 . Since c 1 = γ , and γ is independent of E 0 , we can use Assumption (C2) and apply Corollary A.2 to ( E 0 I r , : , γ r ) to obtain

$$( E _ { \mathcal { I } _ { r } ; } ^ { 0 }, \gamma _ { r } ) \stackrel { W _ { 2 } } { \to } ( \bar { E } _ { r } ^ { 0 }, \bar { \gamma } ).$$

Convergence of ( H ,β 1 ) . For r ∈ [ R ] , recall that u s +1 r = f u s +1 ( e , . . . , e 1 r s r , c 1 r , . . . , c L c r , r ) and ¯ u s +1 r = f u s +1 (¯ e , . . . , e 1 ¯ s , c ¯ 1 r , . . . , c ¯ L c r , r ) . Corollary A.2 implies ( u , . . . , u 1 r l E + l H r ) W 2 → (¯ u , . . . , u 1 r ¯ l E + l H r ) . Moreover, we have shown in (86) that ( u l E +1 , . . . , u l E + l H ) = 1 √ δ ̂ R 0 = 1 √ δ ˜ ( g 0 E , γ, 0 R ) . Hence, using (95) and noting that ˜ g 0 satisfies the polynomial growth condition in (5), we have (¯ u l E +1 r , . . . , u ¯ l E + l H r ) = 1 √ δ ˜ ( ¯ g 0 E , γ, 0 r ¯ r ) , for r ∈ [ R ] .

Corollary A.2 states that ( h , . . . , h 1 c 2 l E + l H c ) W 2 → ( ¯ h , . . . , h 1 ¯ 2 l E + l H ) ∼ N (0 Ξ , 2 l E + l H ) , where

$$\Xi ^ { 2 l _ { E } + l _ { H } } = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \widehat { \Xi } ^ { 2 l _ { E } + l _ { H }, r }, \quad \widehat { \Xi } ^ { 2 l _ { E } + l _ { H } } = \delta \left ( \mathbb { E } [ \bar { u } ^ { r } _ { r } \bar { u } ^ { s } _ { r } ] \right ) ^ { 2 l _ { E } + l _ { H } } _ { r, s = 1 }.$$

Then recalling the definition of Ω 1 from (75), and the functions f u l E +1 , . . . , f u l E + l H from (88)-(89), we have

$$\widehat { \Xi } _ { [ l _ { E } + 1 \colon l _ { E } + l _ { H } ], [ l _ { E } + 1 \colon l _ { E } + l _ { H } ] } ^ { 2 l _ { E } + l _ { H }, r } = \widehat { \Omega } ^ { 1, r } \implies \Xi _ { [ l _ { E } + 1 \colon l _ { E } + l _ { H } ], [ l _ { E } + 1 \colon l _ { E } + l _ { H } ] } ^ { 2 l _ { E } + l _ { H } } = \Omega ^ { 1 }.$$

Since we have shown that ( h l E +1 , . . . , h l E + l H ) = H 1 (see (86)), and Corollary A.2 states that

$$( d _ { \mathbf c } ^ { 1 }, h _ { \mathbf c } ^ { l _ { E } + 1 }, \dots, h _ { \mathbf c } ^ { l _ { E } + l _ { H } } ) \stackrel { W _ { 2 } } { \to } ( \bar { d } _ { \mathbf c } ^ { 1 }, \bar { h } ^ { l _ { E } + 1 }, \dots, \bar { h } ^ { l _ { E } + l _ { H } } ),$$

where ¯ d 1 c and ( ¯ h l E +1 , . . . , h ¯ l E + l H ) ∼ N ( 0 Ξ , l E + l H [ l E +1: l E + l H ] , l [ E +1: l E + l H ] ) are indepedent. Recalling that d 1 c = β c , by the equivalence of the covariance matrices in (97), we have

$$( H ^ { 1 } _ { \mathcal { J } _ { \epsilon } }, \beta _ { \epsilon } ) \stackrel { W _ { 2 } } { \to } ( \bar { H } ^ { 1 } _ { \epsilon }, \bar { \beta } ), \quad \text{where} \, \bar { \beta } \text{ is independent of} \, \bar { H } ^ { 1 } _ { \epsilon } \sim \mathcal { N } ( 0, \Omega ^ { 1 } ).$$

Convergence of ( E , γ 1 ) . Recall that v s c = f v s ( h , . . . , h 1 c t c , d 1 c , . . . , d L d c , c ) and ¯ v s c = f v s ( ¯ h , . . . , h 1 ¯ t , d ¯ 1 c , . . . , d ¯ L d c , c ) , for c ∈ [ C ] . Corollary A.2 implies that ( v , . . . , v 1 c 2 l E + l H c ) W 2 → (¯ v , . . . , v 1 c ¯ 2 l E + l H c ) . Moreover, we have shown in (87) that ( v l E + l H +1 , . . . , v 2 l E + l H ) = 1 √ δ ̂ H 1 = 1 √ δ ˜ ( f 1 H ,β, 1 C ) . Thus, using (98) and noting that ˜ f 1 satisfying the polynomial growth condition in (5), we have (¯ v l E + l H +1 c , . . . , v ¯ 2 l E + l H c ) = 1 √ δ ˜ ( ¯ f 1 H ,β, 1 c ¯ c ) , for c ∈ [ C ] .

Corollary states that A.2 that ( e , . . . , e 1 r 2 l E + l H r ) W 2 → (¯ e , . . . , e 1 ¯ 2 l E + l H ) ∼ N (0 , Γ 2 l E + l H ) , where

$$\Gamma ^ { 2 l _ { E } + l _ { H } } = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \widehat { \Gamma } ^ { 2 l _ { E } + l _ { H }, c }, \quad \widehat { \Gamma } ^ { 2 l _ { E } + l _ { H }, c } = \left ( \mathbb { E } \{ \bar { v } _ { c } ^ { r } \bar { v } _ { c } ^ { s } \} \right ) _ { r, s = 1 } ^ { 2 l _ { E } + l _ { H } }.$$

Then recalling the definition of Π 1 from (75), and the functions f v l E + l H +1 , . . . , f v 2 l E + l H from (91) and (93), we have

$$\widehat { \Gamma } _ { [ l _ { E } + l _ { H } + 1 ; 2 l _ { E } + l _ { H } ], [ l _ { E } + l _ { H } + 1 ; 2 l _ { E } + l _ { H } ] } ^ { 2 l _ { E } + l _ { H } } = \widehat { \Pi } ^ { 1, c } \, \Longrightarrow \, \Gamma _ { [ l _ { E } + l _ { H } + 1 ; 2 l _ { E } + l _ { H } ], [ l _ { E } + l _ { H } + 1 ; 2 l _ { E } + l _ { H } ] } ^ { 2 l _ { E } + l _ { H } } = \Pi ^ { 1 }. \quad ( 1 0 0 )$$

Since ( e l E + l H +1 , . . . , e 2 l E + l H ) = E 1 from (87), and Corollary A.2 states that

$$( c _ { r } ^ { 1 }, e _ { r } ^ { l _ { E } + l _ { H } + 1 }, \dots, e _ { r } ^ { 2 l _ { E } + l _ { H } } ) \stackrel { W _ { 2 } } { \to } ( \bar { c } _ { r } ^ { 1 }, \bar { e } ^ { l _ { E } + l _ { H } + 1 }, \dots, \bar { e } ^ { 2 l _ { E } + l _ { H } } ),$$

where ¯ c 1 r and (¯ e l E + l H +1 , . . . , e ¯ 2 l E + l H ) ∼ N ( 0 Γ , 2 l E + l H [ l E + l H +1:2 l E + l H ] , l [ E + l H +1:2 l E + l H ] ) are independent. Therefore, using c 1 r = γ r and (100), we have ( E 1 I r , : , γ r ) W 2 → ( ¯ E , γ 1 r ¯) , where ¯ E 1 r ∼ N (0 Π ) , 1 is independent of ¯ γ .

Inductive hypothesis. For k ≥ 1 , assume that we can reduce H k , ̂ R k -1 , ̂ H k , and E k to U-AMP in iterations t = l E k + l H ( k -1) + 1 , . . . , l E ( k +1) + l H k . In formulas, this means the following:

$$\text{$w=\cdot_{E}^{\prime}\nu\cdot\cdot\cdotN\vee\nolimits}^{\prime} \quad \gamma \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime}\cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \cdot_{E}^{\prime} \\ \bullet \text{For $t=l_{E}k+l_{H}(k-1)+1,\dots,h^{l_{E}k+kl_{H}$)} = H^{k}, \quad v^{l_{E}k+l_{H}(k-1)+1},\dots,v^{l_{E}k+kl_{H}=0} \\ e^{l_{E}k+l_{H}(k-1)+1},\dots,e^{l_{E}k+kl_{H}=0}, \quad ( u^{l_{E}k+l_{H}(k-1)+1},\dots, u^{l_{E}k+kl_{H}=kl_{H})} = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { k - 1 }.$$

$$\bullet & \text{ For } t = l _ { E } k + l _ { H } k + 1, \dots, l _ { E } ( k + 1 ) + l _ { H } k \colon \text{ We have} \\ & h ^ { l _ { E } k + l _ { H } k + 1 }, \dots, h ^ { l _ { E } ( k + 1 ) + l _ { H } k } = 0, \quad ( v ^ { l _ { E } k + l _ { H } k + 1 }, \dots, v ^ { l _ { E } ( k + 1 ) + l _ { H } k } ) = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { k } \\ & \quad \text{$l_{e} l_{E} k + l _ { H } k + 1 } \quad \text{$l_{E} ( k + 1 ) + l _ { H } k } \sum - \mathbb { R } ^ { k } \quad \text{$l_{E} k + l _ { H } k + 1 } \quad \text{$l_{E} ( k + 1 ) + l _ { H } k } - \mathbb { A }$$

$$h ^ { l _ { E } k + l _ { H } k + 1 }, \dots, h ^ { l _ { E } ( k + 1 ) + l _ { H } k } = 0, \quad ( v ^ { l _ { E } k + l _ { H } k + 1 }, \dots, v ^ { l _ { E } ( k + 1 ) + l _ { H } k } ) = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { k } \\ ( e ^ { l _ { E } k + l _ { H } k + 1 }, \dots, e ^ { l _ { E } ( k + 1 ) + l _ { H } k } ) = E ^ { k }, \quad u ^ { l _ { E } k + l _ { H } k + 1 }, \dots, u ^ { l _ { E } ( k + 1 ) + l _ { H } k } = 0.$$

Defining the index sets

$$\mathcal { I } _ { k } = [ l _ { E } k + l _ { H } k + 1 \, \colon l _ { E } ( k + 1 ) + l _ { H } k ], \quad \mathcal { J } _ { k } = [ l _ { E } k + l _ { H } ( k - 1 ) + 1 \, \colon l _ { E } k + l _ { H } k ],$$

we also assume that the follows convergence statements hold, for 1 ≤ s ≤ k :

$$& ( H ^ { s } _ { \mathcal { J } _ { c } ; }, \beta _ { c } ) \stackrel { W _ { 2 } } { \to } ( \bar { H } ^ { s } _ { c }, \bar { \beta } ), \quad \bar { \beta } \text{ independent of } \bar { H } ^ { s } _ { c } \sim \mathcal { N } ( 0, \Omega ^ { s } ), \\ & ( E ^ { s } _ { \mathcal { I } _ { r } ; }, \gamma _ { r } ) \stackrel { W _ { 2 } } { \to } ( \bar { E } ^ { s } _ { r }, \bar { \gamma } ), \quad \bar { \gamma } \text{ independent of } \bar { E } ^ { k } _ { r } \sim \mathcal { N } ( 0, \Pi ^ { s } ).$$

Inductive step. We need to show that we can reduce H k +1 , ̂ R k , ̂ H k +1 , and E k +1 to U-AMP in iterations t = l E ( k +1)+ l H k +1 , . . . , l E ( k +2)+ l H ( k +1) , and that the corresponding convergence statements hold. The choices for the functions ( f v t , f u t +1 ) are analogous to those in (89), (91), and the steps for reduction are very similar to the base case for t ∈ { l E + 1 , . . . , 2 L E + l H } , and are omitted for brevity. We provide the summary of the reductions below:

$$\text{omited for brevity. we province the summary of the recuions below: } \\ \bullet \text{ For } t = l _ { E } ( k + 1 ) + l _ { H } k + 1, \dots, l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) \colon \text{We have} \\ \left ( h _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } k + 1 }, \dots, h _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) } \right ) = H ^ { k + 1 }, \quad v _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } k + 1 }, \dots, v _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) } = 0 \\ \right ) _ { e ^ { l _ { E } ( k + 1 ) + l _ { H } k + 1 }, \dots, e ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) } = 0, \quad \left ( u _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } k + 1 }, \dots, u _ { E } ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) } \right ) = \frac { 1 } { \sqrt { \delta } } \widehat { R } ^ { k }. \\ \right ) _ { - } \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdclots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \cdots \quad \end{$$

$$& \bullet \text{ For } t = l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1, \dots, l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) \colon \text{ We have} \\ & \quad h ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1 }, \dots, h ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = 0, \quad \left ( v ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1 }, \dots, v ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } \right ) = \frac { 1 } { \sqrt { \delta } } \widehat { H } ^ { k + 1 } \\ & \quad \left ( e ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1 }, \dots, e ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } \right ) = E ^ { k + 1 }, \quad u ^ { l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1 }, \dots, u ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = 0.$$

$$( 1 0 5 )$$

We now show the convergence statements in Theorem A.1 for H k +1 , and E k +1 by reducing the abstract matrix-AMP SE parameters to the corresponding U-AMP SE parameters. Define the index sets

$$\mathcal { I } _ { k + 1 } = [ l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) + 1 \, \colon l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) ], \, \mathcal { J } _ { k + 1 } = [ l _ { E } ( k + 1 ) + l _ { H } k + 1 \, \colon l _ { E } ( k + 1 ) + l _ { H } ( k + 1 ) ].$$

Convergence of ( H k +1 , β ) . From (104) we have ( u l r ) l ∈J k +1 = 1 √ δ ̂ R k I r , : = 1 √ δ ˜ ( g k E k I r , : , γ r , r ) , and by the inductive hypothesis, ( E k I r , : , γ r ) W 2 → ( ¯ E ,γ k r ¯) . Since ˜ g k satisfies the polynomial growth condition in (5), Corollary A.2 implies that ( u l r ) l ∈J k +1 W 2 → (¯ ) u l r l ∈J k +1 = 1 √ δ ˜ ( ¯ g k E ,γ, k r ¯ r ) . Corollary A.2 states that ( h l c ) l ∈J k +1 W 2 → ( ¯ h l ) l ∈J k +1 ∼ N (0 Ξ , l E ( k +2)+ l H ( k +1) J k +1 , J k +1 ) , where

$$\Xi _ { \mathcal { J } _ { k + 1 }, \mathcal { J } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = \frac { 1 } { R } \sum _ { r = 1 } ^ { R } \Xi _ { \mathcal { J } _ { k + 1 }, \mathcal { J } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), r }, \quad \widehat { \Xi } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), r } = \delta \left ( \mathbb { E } [ \bar { u } _ { r } \bar { u } _ { r } ^ { s } ] \right ) _ { r, s \in \mathcal { J } _ { k + 1 } }.$$

Then comparing the above with the definition of Ω k +1 in (75), we get

$$\widehat { \Xi } _ { \mathcal { J } _ { k + 1 }, \mathcal { J } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), r } = \widehat { \Omega } ^ { k + 1, r } \implies \Xi _ { \mathcal { J } _ { k + 1 }, \mathcal { J } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = \Omega ^ { k + 1 }.$$

From (104), we have ( h l ) l ∈J k +1 = H k +1 , and Corollary A.2 states that ( d , 1 c ( h l c ) l ∈J k +1 ) W 2 → ( ¯ d , 1 c ( ¯ h l ) l ∈J k +1 ) , where ¯ d 1 c and ( ¯ h l ) l ∈J k +1 are independent. Thus, by the equivalence of covariance matrices in (106) and recalling d 1 c = β c , we have ( H k +1 J c , : , β c ) W 2 → ( ¯ H k +1 c , β ¯ ) , where ¯ β is independent of ¯ H k +1 c ∼ N (0 , Ω k +1 ) .

Convergence of ( E k +1 , γ ) . From (105), we have ( v l c ) l ∈I k +1 = 1 √ δ ̂ H k +1 J c , : = 1 √ δ ˜ f k +1 ( H k +1 J c , : , β c , c ) , and we have shown that ( H k +1 J c , : , β c ) W 2 → ( ¯ H k +1 c , β ¯ ) . Since ˜ f k +1 satisfies the polynomial growth condition, Corollary A.2 implies that ( v l c ) l ∈I k +1 W 2 → (¯ ) v l c l ∈I k +1 = 1 √ δ ˜ f k +1 ( ¯ H k +1 , β, ¯ c ) .

Corollary A.2 states that ( e l r ) l ∈I k +1 W 2 → (¯ ) e l l ∈I k +1 ∼ N (0 Γ , l E ( k +2)+ l H ( k +1) I k +1 , I k +1 ) , where

$$\Gamma _ { \mathcal { I } _ { k + 1 }, \mathcal { I } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = \frac { 1 } { \mathbb { C } } \sum _ { c = 1 } ^ { C } \widehat { \Gamma } _ { \mathcal { I } _ { k + 1 }, \mathcal { I } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), \epsilon }, \quad \widehat { \Gamma } _ { \mathcal { I } _ { k + 1 }, \mathcal { I } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), \epsilon } = \left ( \mathbb { E } [ \bar { v } _ { c } ^ { r } \bar { v } _ { c } ] \right ) _ { r, s \in \mathcal { I } _ { k + 1 } }.$$

Then comparing the above with the definition of Π k +1 in (75), we get

$$\widehat { \Gamma } _ { \mathcal { I } _ { k + 1 }, \mathcal { I } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ), \epsilon } = \widehat { \Pi } ^ { k + 1, \epsilon } \implies \Gamma _ { \mathcal { I } _ { k + 1 }, \mathcal { I } _ { k + 1 } } ^ { l _ { E } ( k + 2 ) + l _ { H } ( k + 1 ) } = \Pi ^ { k + 1 }$$

From (105), we have ( e l ) l ∈I k +1 = E k +1 and Corollary A.2 states that ( c 1 , ( e l r ) l ∈I k +1 ) W 2 → ( ¯ c 1 , (¯ ) e l l ∈I k +1 ) , where ¯ c 1 and (¯ ) e l l ∈I k +1 are independent. Thus, by the equivalence of covariance matrices in (107) and recalling c 1 r = γ r , we have ( E k +1 I r , : , γ r ) W 2 → ( ¯ E k +1 r , γ ¯) , where ¯ γ is independent of ¯ E k +1 r ∼ N (0 Π , k +1 ) . This completes the proof of the inductive step, and hence, of Theorem A.1.

## A.3 Proof of Corollary A.2

The abstract AMP recursion in (79), but without the block-wise dependence of the functions f v t and f u t +1 , was analyzed in [36]. We show how the block-wise dependence can be included without loss of generality, and thereby prove Corollary A.2 by appealing to the state evolution result of [36].

The abstract AMP iteration for a generalized white noise matrix ˜ X ∈ R n × p analyzed in [36] is as follows. Given an initializer u 1 ∈ R n , side information c 1 , . . . , c L ∗ c ∈ R n and d , . . . , d 1 L ∗ d ∈ R p , all independent of ˜ X , the iterates of the abstract AMP recursion are computed as:

$$\text{prename or } a, \, \text{and are always or in a conversion } a \, \text{compute as}. \\ h ^ { t } & = \sqrt { \delta } \widetilde { X } ^ { \top } u ^ { t } - \sum _ { s = 1 } ^ { t - 1 } b _ { s } ^ { t } v ^ { s }, \quad v ^ { t } = f _ { t } ^ { v } ( h ^ { 1 }, \dots, h ^ { t }, d ^ { 1 }, \dots, d ^ { L _ { d } ^ { d } } ), \\ e ^ { t } & = \sqrt { \delta } \widetilde { X } v ^ { t } - \sum _ { s = 1 } ^ { t } a _ { s } ^ { t } u ^ { s }, \quad u ^ { t + 1 } = f _ { t + 1 } ^ { u } ( e ^ { 1 }, \dots, e ^ { t }, c ^ { 1 }, \dots, c ^ { L _ { c } ^ { d } } ), \\ \text{and are } a, \, \text{and are } v \,. \, \text{$\mathfrak{d}$} + L _ { d } ^ { \prime } \cdots \, \mathfrak{D} \, \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} + \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \text{$\mathfrak{d}$} \,. \, \end{$$

where the functions f v t : R t + L ∗ d → R , f u t +1 : R t + L ∗ c → R act row-wise. The memory coefficients { b t s } s&lt;t and { a t s } s ≤ t are defined below in (109). We have the following assumptions:

(E1) When n, p →∞ , we have n/p = δ &gt; 0 , for fixed L ∗ c and L ∗ d . Furthermore, we have

$$( u ^ { 1 }, c ^ { 1 }, \dots, c ^ { L _ { c } ^ { * } } ) \stackrel { W } { \to } ( \bar { u } ^ { 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } } ) \text{ and } ( d ^ { 1 }, \dots, d ^ { L _ { d } ^ { * } } ) \stackrel { W } { \to } ( \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } } ),$$

for joint limit laws (¯ u , c 1 ¯ 1 , . . . , c ¯ L ∗ c ) and ( ¯ d , . . . , d 1 ¯ L ∗ d ) having finite moments of all orders, where E [(¯ u 1 ) 2 ] ≥ 0 . Multivariate polynomials are dense in the real L 2 -spaces of functions f : R L ∗ c +1 → R and g : R L ∗ d → R with the inner products

$$\langle f, \tilde { f } \rangle \coloneqq \mathbb { E } [ f ( \bar { u } ^ { 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } } ) \tilde { f } ( \bar { u } ^ { 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } } ) ] \text{ and } \langle g, \tilde { g } \rangle \coloneqq \mathbb { E } [ g ( \bar { u } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } } ) \tilde { g } ( \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } } ) ].$$

(E2) , (E3) , (E4) These are identical to (D2) , (D3) , (D4) , with L ∗ d replacing L d .

The state evolution covariance matrices Ξ t , Γ t ∈ R t × t are iteratively defined as follows, starting from Ξ 1 = δ E [(¯ u 1 ) 2 ] ∈ R 1 × 1 . Given Ξ t , for t ≥ 1 , let ( ¯ h , . . . , h 1 ¯ ) t ∼ N (0 Ξ ) , t independent of ( ¯ d , . . . , d 1 ¯ L ∗ d ) and define

$$\bar { v } ^ { s } = f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { s }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } } ), \ \ s \in [ t ].$$

Then, Γ t = ( E [¯ v v r ¯ ]) s t r,s =1 ∈ R t × t . Next, let (¯ e , . . . , e 1 ¯ ) t ∼ N (0 Γ ) , t independent of (¯ u , c 1 ¯ 1 , . . . , ¯ c L ∗ c ) and define

$$\bar { u } ^ { s + 1 } = f _ { s + 1 } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { s }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } } ), \ \ s \in [ t ].$$

Then, Ξ t +1 = ( δ · E [¯ u u r ¯ ]) s t +1 r,s =1 ∈ R ( +1) t × ( +1) t . The memory coefficients in (108) are then defined as

$$a _ { s } ^ { t } = \mathbb { E } [ \partial _ { s } f _ { t } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } } ) ] \text{ and } b _ { s } ^ { t } = \delta \cdot \mathbb { E } [ \partial _ { s } f _ { t } ^ { u } ( \bar { e } ^ { 1 } \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } } ) ],$$

where ∂ s denotes partial derivative in the s th argument. The following theorem gives the state evolution result for the abstract AMP recursion.

Theorem A.3. [36, Theorem 2.17] Let ˜ X ∈ R n × p be a generalized white noise matrix (as defined in Definition 5.1) with variance profile S ∈ R n × p , and let u , c 1 1 , . . . , c L ∗ c , d 1 , . . . , d L ∗ d be independent of ˜ X and satisfy Assumptions (E1)-(E4). Further assume that each matrix Ξ t and Γ t is non-singular. Then for any fixed t ≥ 1 , almost surely as n, p → ∞ with n/p = δ ∈ (0 , ∞ ) , the iterates of the abstract AMP in (108) satisfy

$$& ( u ^ { 1 }, c ^ { 1 }, \dots, c ^ { L _ { c } ^ { * } }, e ^ { 1 }, \dots, e ^ { t } ) \stackrel { W _ { 2 } } { \to } ( \bar { u } ^ { 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } ^ { * } }, \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t } ) \\ & ( d ^ { 1 }, \dots, d ^ { L _ { d } ^ { * } }, h ^ { 1 }, \dots, h ^ { t } ) \stackrel { W _ { 2 } } { \to } ( \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } ^ { * } }, \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t } ),$$

where ( ¯ h , . . . , h 1 ¯ ) t ∼ N (0 , Ξ ) t and (¯ e , . . . , e 1 ¯ ) t ∼ N (0 Γ ) , t are independent of (¯ u , c 1 ¯ 1 . . . . , c ¯ L ∗ c ) and ( ¯ d , . . . , d 1 ¯ L ∗ d ) .

To obtain the U-AMP recursion (79) from the abstract AMP recursion in (108), we choose L ∗ c = L c +1 and L ∗ d = L d +1 , and the side information vectors d L d +1 ∈ R p and c L c +1 ∈ R n are set as

$$d ^ { L _ { d } + 1 } = \mathcal { C }, \quad c ^ { L _ { c } + 1 } = \mathcal { R }.$$

The functions f v t and f u t +1 , as well as the initializer u 1 are the same as those in (108).

The memory coefficients in (109) can then be expressed as:

With this choice, the empirical distribution of d L d +1 converges to ¯ d L d +1 ∼ Uniform ([ C ]) , and the empirical distribution of c L c +1 converges to ¯ c L c +1 ∼ Uniform ([ R ]) . Moreover, Assumption (E1) is equivalent to Assumption (D1) of Corollary A.2. To see this, for r ∈ [ R ] , let (¯ u , c 1 r ¯ 1 r , . . . , c ¯ L c r ) be random variables whose joint law equals the conditional law of (¯ u , c 1 ¯ 1 , . . . , c ¯ L c ) given ¯ c L c +1 = r . Similarly, for c ∈ [ C ] , let ( ¯ d , . . . , d 1 c ¯ L d c ) be jointly distributed according to the conditional law of ( ¯ d , . . . , d 1 ¯ L d ) given ¯ d L d +1 = c .

$$\text{$1$} \text{ memory connections in (10)$} \text{ can even we expressed as} \colon \\ a ^ { t } _ { s } = \mathbb { E } [ \partial _ { s } f ^ { v } _ { t } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { h } ^ { 1 }, \dots, \bar { d } ^ { L d + 1 } ) ] & = \mathbb { E } \left [ \mathbb { E } [ \partial _ { s } f ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L d + 1 } ) \, | \, \bar { d } ^ { L d + 1 } ] \right ] \\ & = \frac { 1 } { \mathbb { C } } \sum _ { c = 1 } ^ { c } \mathbb { E } [ \partial _ { s } f ^ { v } _ { t } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L d }, c ) \, | \, \bar { d } ^ { L d + 1 } = c ], \\ & = \frac { 1 } { \mathbb { C } } \sum _ { c = 1 } ^ { c } \mathbb { E } [ \partial _ { s } f ^ { v } _ { t } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 } _ { c }, \dots, \bar { d } ^ { L d } _ { c }, c ) ]. \\ \text{throw for that last enough if you need to set that } \bar { h } ^ { 1 } _ { t } \, \text{$\bar{ }h } ^ { 1 } \, \text{$\bar{ }i$} \, \text{information} \, of $f(\bar{ }l_{1}$} } \quad \bar { h } ^ { L d + 1 } \, \text{$\bar{ }l_{1}$} \, \text{$Similarity}$$

where for the last equality we used fact that ( ¯ h , . . . , h 1 ¯ ) t is independent of ( ¯ d , . . . , d 1 ¯ L d +1 ) . Similarly, we have

$$b _ { s } ^ { t } = \delta \mathbb { E } [ \partial _ { s } f _ { t } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } + 1 } ) ] & = \delta \mathbb { E } \left [ \mathbb { E } [ \partial _ { s } f _ { t } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L _ { c } + 1 } ) \, | \, \bar { c } ^ { L _ { c } + 1 } ] \right ] \\ & = \frac { \delta } { R } \sum _ { r = 1 } ^ { R } \mathbb { E } [ \partial _ { s } f _ { t } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 } _ { r }, \dots, \bar { c } ^ { L _ { c } + 1 } _ { r } ) ].$$

$$\begin{smallmatrix} 1 \end{smallmatrix}$$

Next, for r, s ∈ [ t ] , the ( r, s ) th element of the state evolution matrix Γ t ∈ R t × t is

$$\left ( \Gamma ^ { t } \right ) _ { r, s } & = \mathbb { E } \left [ f _ { r } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } + 1 } ) f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } + 1 } ) \right ] \\ & = \mathbb { E } \left [ \mathbb { E } \left [ f _ { r } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } + 1 } ) f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 }, \dots, \bar { d } ^ { L _ { d } + 1 } ) \left | \bar { d } ^ { L _ { d } + 1 } \right ] \right ] \\ & = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \mathbb { E } \left [ f _ { r } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 } _ { c }, \dots, \bar { d } ^ { L _ { d } } _ { c }, c ) f _ { s } ^ { v } ( \bar { h } ^ { 1 }, \dots, \bar { h } ^ { t }, \bar { d } ^ { 1 } _ { c }, \dots, \bar { d } ^ { L _ { d } } _ { c }, c ) \right ] = \frac { 1 } { C } \sum _ { c = 1 } ^ { C } \mathbb { E } \left [ v _ { c } ^ { v } v _ { c } ^ { c } \right ], \\ \cdot \quad \ \ - \quad \ \frac { \sim } { c } \sum _ { c = 1 } ^ { C } \frac { \sim } { c } \sum _ { c = 1 } ^ { C } \frac { \sim } { c } \sum _ { c = 1 } ^ { C } \frac { \sim } { c } \sum _ { c = 1 } ^ { C } \frac { \sim } { c } \cdot \quad \ \.$$

where ¯ v s c = f v s ( ¯ h , . . . , h 1 ¯ t , d ¯ 1 c , . . . , d ¯ L d c , c ) . Similarly, we have

$$( \Xi ^ { t } ) _ { r, s } = \frac { \delta } { R } \sum _ { r = 1 } ^ { R } \mathbb { E } [ f _ { r } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L c + 1 } = r ) f _ { s } ^ { u } ( \bar { e } ^ { 1 }, \dots, \bar { e } ^ { t - 1 }, \bar { c } ^ { 1 }, \dots, \bar { c } ^ { L c + 1 } = r ) ] = \frac { \delta } { R } \sum _ { r = 1 } ^ { R } \mathbb { E } [ \bar { u } ^ { r } _ { r } \bar { u } ^ { s } _ { r } ],$$

where for s ≥ 1 , we have ¯ u s +1 r = f u s +1 (¯ e , . . . , e 1 ¯ s , c ¯ 1 r , . . . , c ¯ L c r , r ) . We have shown that with the choice of side information in (110), the AMP recursion (108) matches with that in (79), and the corresponding state evolution recursions also match. Applying Theorem A.3 and recalling the definitions of R C , from (51) gives us Corollary A.2.

## B Proof of Lemma 3.6

For δ &gt; 0 , evaluating U b ( 0 ; δ ) in (31) with b 0 := δσ 2 gives

$$U ( b _ { 0 } ; \delta ) = - \frac { \delta } { 2 } + \delta \log 2 + 2 I \left ( \bar { \beta }, \sqrt { \frac { 1 } { 2 \sigma ^ { 2 } } } \bar { \beta } + G \right ).$$

We pause to state the following auxiliary result.

Lemma B.1. [34, Proposition 7.15] For a discrete distribution P ¯ β with finite alphabet, we have

$$\lim _ { s \to \infty } \frac { I ( \bar { \beta } ; \sqrt { s } \bar { \beta } + G ) } { \frac { 1 } { 2 } \log s } = 0.$$

Lemma B.1 implies that for any ∆ &gt; 0 , we have I ( ¯ ; β √ sβ ¯ + G ) ≤ ∆ 2 log s for all sufficiently large s . Taking s = 1 2 σ 2 further implies that for sufficiently small σ , we have

$$I \left ( \bar { \beta } ; \sqrt { \frac { 1 } { 2 \sigma ^ { 2 } } } \bar { \beta } + G \right ) \leq \frac { \Delta } { 2 } \log \left ( \frac { 1 } { 2 \sigma ^ { 2 } } \right ) \\ \Longrightarrow & U ( b _ { 0 } ; \delta ) \leq - \frac { \delta } { 2 } + \delta \log 2 + \Delta \log \left ( \frac { 1 } { 2 \sigma ^ { 2 } } \right ).$$

Hence, for any ∆ &gt; 0 , there exists σ 0 (∆) &gt; 0 such that for all σ &lt; σ 0 (∆) we have the following for all b ∈ (0 , Var( ¯ )] β and δ &gt; 0 :

$$\dots \nu \subset \{ \nu, \dots \nu \nmid \nu \nmid \nu \times \nu. \\ U ( b ; \delta ) - U ( b _ { 0 } ; \delta ) & \geq - \delta \left ( 1 - \frac { \delta \sigma ^ { 2 } } { b + \delta \sigma ^ { 2 } } \right ) + \delta \log \left ( 1 + \frac { b } { \delta \sigma ^ { 2 } } \right ) + 2 I \left ( \tilde { \beta } ; \sqrt { \frac { \delta } { b + \delta \sigma ^ { 2 } } } \tilde { \beta } + G \right ) \\ & + \frac { \delta } { 2 } - \delta \log 2 - \Delta \log \left ( \frac { 1 } { 2 \sigma ^ { 2 } } \right ) \\ & \geq \delta \log \left ( \frac { b } { \delta \sigma ^ { 2 } } \right ) - \Delta \log \left ( \frac { 1 } { 2 \sigma ^ { 2 } } \right ) - \frac { \delta } { 2 } - \delta \log 2, \\ \dots \dots \nmid \nu \dots \nmid \nu \nmid \dots \nmid \nu \nmid \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \dots \nmid \nu \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \ delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \dta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \Delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta \delta$$

where (a) uses (31) and (113), and (b) uses the non-negativity of mutual information. Now, for σ &lt; σ 0 (∆) and

$$b > 2 \delta e ^ { \frac { 1 } { 2 } } \left ( \frac { 1 } { 2 } \right ) ^ { \frac { \Delta } { \delta } } \sigma ^ { 2 - \frac { 2 \Delta } { \delta } },$$

the lower bound in (114) is strictly positive. Therefore, U b δ ( ; ) -U b ( 0 ; δ ) &gt; 0 for b satisfying (115), implying that these values of b cannot be minimizers of U b δ ( ; ) . Therefore, for any δ, ∆ &gt; 0 and σ &lt; σ 0 (∆) , we have:

$$\max \left \{ \underset { b \in ( 0, \text{Var} ( \bar { \beta } ) ] } { \arg \min } U ( b ; \delta ) \right \} \leq 2 \delta e ^ { \frac { 1 } { 2 } } \left ( \frac { 1 } { 2 } \right ) ^ { \frac { \Delta } { \delta } } \sigma ^ { 2 - \frac { 2 \Delta } { \delta } } < \frac { 7 } { 2 } \delta ( \sigma ^ { 2 - \frac { 2 \Delta } { \delta } } ).$$

## C Implementation Details

SC-AMP denoiser and state evolution parameters for QGT. The Bayes-optimal denoiser f k +1 in (19) can be computed using the prior ¯ β ∼ Bernoulli ( π ) . For j ∈ J c and k ≥ 1 , we have

$$f _ { k } ( s, c ) & = \mathbb { E } [ \bar { \beta } \, | \, ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G = s ] = \frac { \mathbb { P } \left [ \bar { \beta } = 1 \right ] \cdot \mathbb { P } \left [ ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G = s | \bar { \beta } = 1 \right ] } { \sum _ { \bar { \beta } \in \{ 0, 1 \} } \mathbb { P } [ \bar { \beta } ] \cdot \mathbb { P } [ ( \chi _ { c } ^ { k } ) ^ { 2 } \bar { \beta } + \chi _ { c } ^ { k } G = s | \bar { \beta } ] } \\ & = \frac { \pi \phi ( ( s - ( \chi _ { c } ^ { k } ) ^ { 2 } ) / \chi _ { c } ^ { k } ) } { \pi \phi ( ( s - ( \chi _ { c } ^ { k } ) ^ { 2 } ) / \chi _ { c } ^ { k } ) },$$

where ϕ x ( ) is the standard normal density. Instead of precomputing the state evolution parameters ( χ k c ) , they can be estimated from the SC-AMP iterates as:

$$\left ( \widehat { \chi } _ { c } ^ { k } \right ) ^ { 2 } = \sum _ { r = 1 } ^ { R } \widetilde { W } _ { r c } \left ( \left \| \widetilde { \Theta } _ { r } ^ { k } \right \| _ { 2 } ^ { 2 } \right ) ^ { - 1 } \quad \text{ for } c \in [ C ],$$

where ˜ Θ k r = (Θ ) ˜ k i i ∈I r is the restriction of ˜ Θ k to indices I r . The derivative ∂ f 1 k ( β , k j c ) , required for b k in (15), can be obtained by applying the Quotient rule to the last expression in (116).

SC-AMP denoiser and state evolution parameters for pooled data. The Bayes-optimal denoiser f k : R L × R → R L in the SC-AMP algorithm in (40) is computed as:

$$f _ { k } ( s, c ) & = \mathbb { E } [ \bar { B } \, | \, \bar { B } + G _ { \epsilon } ^ { k } = s ] = \sum _ { l = 1 } ^ { L } e _ { l } \frac { \mathbb { P } [ \bar { B } = e _ { l } ] \mathbb { P } [ \bar { B } + G _ { \epsilon } ^ { k } = s | \bar { B } = e _ { l } ] } { \mathbb { P } [ \bar { B } + G _ { \epsilon } ^ { k } = s ] } \\ & = \frac { ( a ) } { \sum _ { l = 1 } ^ { L } \pi _ { l } e _ { l } \exp \left ( - \frac { 1 } { 2 } ( e _ { l } - s ) ^ { \intercal } \left ( \mathbb { T } _ { \epsilon } ^ { k } \right ) ^ { - 1 } ( e _ { l } - s ) \right ) } } { \sum _ { l = 1 } ^ { L } \pi _ { l } \exp \left ( - \frac { 1 } { 2 } ( e _ { l } - s ) ^ { \intercal } \left ( \mathbb { T } _ { \epsilon } ^ { k } \right ) ^ { - 1 } ( e _ { l } - s ) \right ) }, \\ & \quad \quad \quad$$

where (a) follows since G k c ∼ N ( 0 T , k c ) . The state evolution parameters ϕ k r for r ∈ [ R ] are estimated from matrix SC-AMP iterates as

$$\hat { \phi } _ { r } ^ { k } = \frac { 1 } { n / R } \sum _ { i \in \mathcal { I } _ { r } } ( \widetilde { \Theta } _ { i } ^ { k } ) ^ { \top } \widetilde { \Theta } _ { i } ^ { k }, \quad \text{ for } r \in [ R ].$$

The Jacobian f ′ k ( B k j, : , c ) in (42) can be computed for all j ∈ [ p ] can be computed by applying the Quotient rule to (118), following the method in [14, App. D.1].

Potential function. To generate the curves in Figure 2, for each δ , the potential function U b δ ( ; ) in in (31) is evaluated at 500 data points between 0 and Var ( ¯ ) β (for π = 0 1 . , Var ( ¯ ) = β π -π 2 = 0 09 . ). To analyze the noiseless QGT model, we set σ = 1 × 10 -30 to avoid computational instability. The mutual information term in (31) is computed via numerical integration (instead of Monte Carlo methods) to ensure that the curves are smooth.

## References

| [1]   | M. Aldridge, O. Johnson, and J. Scarlett, 'Group testing: an information theory perspective,' Founda- tions and Trends® in Communications and Information Theory , vol. 15, no. 3-4, pp. 196-392, 2019.                                                               |
|-------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [2]   | C. L. Chan, S. Jaggi, V. Saligrama, and S. Agnihotri, 'Non-adaptive group testing: Explicit bounds and novel algorithms,' IEEE Transactions on Information Theory , vol. 60, no. 5, pp. 3019-3035, 2014.                                                              |
| [3]   | O. Gebhard, O. Johnson, P. Loick, and M. Rolvien, 'Improved bounds for noisy group testing with constant tests per item,' IEEE Transactions on Information Theory , vol. 68, no. 4, pp. 2604-2621, 2021.                                                              |
| [4]   | O. Gebhard, M. Hahn-Klimroth, O. Parczyk, M. Penschuck, M. Rolvien, J. Scarlett, and N. Tan, 'Near- optimal sparsity-constrained group testing: Improved bounds and algorithms,' IEEE Transactions on Information Theory , vol. 68, no. 5, pp. 3253-3280, 2022.       |
| [5]   | N. Tan, W. Tan, and J. Scarlett, 'Performance bounds for group testing with doubly-regular designs,' IEEE Transactions on Information Theory , vol. 69, no. 2, pp. 1224-1243, 2023.                                                                                   |
| [6]   | N. H. Bshouty, 'Optimal algorithms for the coin weighing problem with a spring scale,' Conference on Learning Theory , 2009.                                                                                                                                          |
| [7]   | C.-C. Cao, C. Li, and X. Sun, 'Quantitative group testing-based overlapping pool sequencing to identify rare variant carriers,' BMC Bioinformatics , vol. 15, no. 1, pp. 1-14, 2014.                                                                                  |
| [8]   | G. De Marco, T. Jurdziński, and D. R. Kowalski, 'Optimal channel utilization with limited feedback,' Journal of Computer and System Sciences , vol. 119, pp. 21-33, 2021.                                                                                             |
| [9]   | C. Wang, Q. Zhao, and C.-N. Chuah, 'Group testing under sum observations for heavy hitter detection,' Information Theory and Applications Workshop , pp. 149-153, 2015.                                                                                               |
| [10]  | I.-H. Wang, S.-L. Huang, K.-Y. Lee, and K.-C. Chen, 'Data extraction via histogram and arithmetic mean queries: Fundamental limits and algorithms,' IEEE International Symposium on Information Theory , pp. 1386-1390, 2016.                                         |
| [11]  | J. Scarlett and V. Cevher, 'Phase transitions in the pooled data problem,' Advances in Neural Infor- mation Processing Systems , vol. 30, 2017.                                                                                                                       |
| [12]  | V. Grebinski and G. Kucherov, 'Optimal reconstruction of graphs under the additive model,' Algorith- mica , vol. 28, pp. 104-124, 2000.                                                                                                                               |
| [13]  | A. El Alaoui, A. Ramdas, F. Krzakala, L. Zdeborová, and M. I. Jordan, 'Decoding from pooled data: Sharp information-theoretic bounds,' SIAM Journal on Mathematics of Data Science , vol. 1, no. 1, pp. 161-188, 2019.                                                |
| [14]  | N. Tan, P. Pascual Cobo, J. Scarlett, and R. Venkataramanan, 'Approximate message passing with rigorous guarantees for pooled data and quantitative group testing,' to appear in SIAM Journal on Mathematics of Data Science , 2024. arXiv:2309.15507.                |
| [15]  | Y. Kabashima, 'A CDMA multiuser detection algorithm on the basis of belief propagation,' Journal of Physics A: Mathematical and General , vol. 36, no. 43, pp. 11 111-11 121, Oct 2003.                                                                               |
| [16]  | M. Bayati and A. Montanari, 'The dynamics of message passing on dense graphs, with applications to compressed sensing,' IEEE Transactions on Information Theory , vol. 57, pp. 764-785, 2011.                                                                         |
| [17]  | D. L. Donoho, A. Maleki, and A. Montanari, 'Message passing algorithms for compressed sensing,' Proceedings of the National Academy of Sciences , vol. 106, pp. 18 914-18 919, 2009.                                                                                  |
| [18]  | F. Krzakala, M. Mézard, F. Sausset, Y. Sun, and L. Zdeborová, 'Probabilistic reconstruction in com- pressed sensing: algorithms, phase diagrams, and threshold achieving matrices,' Journal of Statistical Mechanics: Theory and Experiment , vol. 2012, no. 8, 2012. |

| [19]   | S. Rangan, 'Generalized approximate message passing for estimation with random linear mixing,' IEEE International Symposium on Information Theory , 2011.                                                                                            |
|--------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [20]   | J. Ma, J. Xu, and A. Maleki, 'Optimization-based AMP for phase retrieval: The impact of initialization and ℓ 2 regularization,' IEEE Transactions on Information Theory , vol. 65, no. 6, pp. 3600-3629, 2019.                                       |
| [21]   | A. Maillard, B. Loureiro, F. Krzakala, and L. Zdeborová, 'Phase retrieval in high dimensions: Statistical and computational phase transitions,' in Neural Information Processing Systems , 2020.                                                     |
| [22]   | M. Mondelli and R. Venkataramanan, 'Approximate message passing with spectral initialization for generalized linear models,' International Conference on Artificial Intelligence and Statistics , pp. 397- 405, 2021.                                |
| [23]   | N. Tan and R. Venkataramanan, 'Mixed regression via approximate message passing,' Journal of Ma- chine Learning Research , vol. 24, pp. 1-44, 2023.                                                                                                  |
| [24]   | Y. Deshpande and A. Montanari, 'Information-theoretically optimal sparse PCA,' in IEEE International Symposium on Information Theory , 2014, pp. 2197-2201.                                                                                          |
| [25]   | A. K. Fletcher and S. Rangan, 'Iterative reconstruction of rank-one matrices in noise,' Information and Inference: A Journal of the IMA , vol. 7, no. 3, pp. 531-562, 2018.                                                                          |
| [26]   | Y. Kabashima, F. Krzakala, M. Mézard, A. Sakata, and L. Zdeborová, 'Phase transitions and sample complexity in Bayes-optimal matrix factorization,' IEEE Transactions on Information Theory , vol. 62, no. 7, pp. 4228-4265, 2016.                   |
| [27]   | T. Lesieur, F. Krzakala, and L. Zdeborová, 'Constrained low-rank matrix estimation: Phase transi- tions, approximate message passing and applications,' Journal of Statistical Mechanics: Theory and Experiment , vol. 2017, no. 7, p. 073403, 2017. |
| [28]   | A. Montanari and R. Venkataramanan, 'Estimation of low-rank matrices via approximate message passing,' Annals of Statistics , vol. 45, no. 1, pp. 321-345, 2021.                                                                                     |
| [29]   | G. Li, W. Fan, and Y. Wei, 'Approximate message passing from random initialization with applica- tions to Z 2 synchronization,' Proceedings of the National Academy of Sciences , vol. 120, no. 31, p. e2302930120, 2023.                            |
| [30]   | R. Rossetti and G. Reeves, 'Approximate message passing for the matrix tensor product model,' 2023, arXiv: 2306.15580.                                                                                                                               |
| [31]   | O. Y. Feng, R. Venkataramanan, C. Rush, and R. J. Samworth, 'A unifying tutorial on approximate message passing,' Foundations and Trends in Machine Learning , 2022.                                                                                 |
| [32]   | S. Kudekar and H. D. Pfister, 'The effect of spatial coupling on compressive sensing,' in 48th Annual Allerton Conf. Commun., Control, and Computing , 2010.                                                                                         |
| [33]   | K. Takeuchi, T. Tanaka, and T. Kawabata, 'Performance improvement of iterative multiuser detec- tion for large sparsely spread CDMA systems by spatial coupling,' IEEE Transactions on Information Theory , vol. 61, no. 4, pp. 1768-1794, 2015.     |
| [34]   | D. L. Donoho, A. Javanmard, and A. Montanari, 'Information-theoretically optimal compressed sensing via spatial coupling and approximate message passing,' IEEE Transactions on Information Theory vol. 59, no. 11, pp. 7434-7464, Nov. 2013.        |
| [35]   | P. Pascual Cobo, K. Hsieh, and R. Venkataramanan, 'Bayes-optimal estimation in generalized linear models via spatial coupling,' in IEEE International Symposium on Information Theory (ISIT) , 2023, arXiv:2309.08404.                               |
| [36]   | T. Wang, X. Zhong, and Z. Fan, 'Universality of approximate message passing algorithms and tensor networks,' 2022, to appear in Annals of Applied Probability . arXiv:2206.13037.                                                                    |

| [37]   | A. Yedla, Y.-Y. Jian, P. S. Nguyen, and H. D. Pfister, 'A simple proof of Maxwell saturation for coupled scalar recursions,' IEEE Transactions on Information Theory , vol. 60, no. 11, pp. 6943-6965, 2014.                                     |
|--------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [38]   | A. J. Felstrom and K. S. Zigangirov, 'Time-varying periodic convolutional codes with low-density parity- check matrix,' IEEE Transactions on Information Theory , vol. 45, no. 6, pp. 2181-2191, Sept. 1999.                                     |
| [39]   | S. Kudekar, T. Richardson, and R. L. Urbanke, 'Spatially coupled ensembles universally achieve capacity under belief propagation,' IEEE Transactions on Information Theory , vol. 59, no. 12, pp. 7761-7813, Dec 2013.                           |
| [40]   | A. Coja-Oghlan, O. Gebhard, M. Hahn-Klimroth, and P. Loick, 'Optimal group testing,' Conference on Learning Theory , pp. 1374-1388, 2020.                                                                                                        |
| [41]   | A. Coja-Oghlan, M. Hahn-Klimroth, L. Hintze, D. Kaaser, L. Krieg, M. Rolvien, and O. Scheftelowitsch, 'Noisy group testing via spatial coupling,' 2024, arXiv:2402.02895.                                                                        |
| [42]   | M. M. Mashauri, A. G. i Amat, and M. Lentmaier, 'Low-density parity-check codes and spatial coupling for quantitative group testing,' in IEEE International Symposium on Information Theory (ISIT) , 2023.                                       |
| [43]   | --, 'Threshold saturation for quantitative group testing with low-density parity-check codes,' in IEEE International Symposium on Information Theory (ISIT) , 2024.                                                                              |
| [44]   | E. Karimi, F. Kazemi, A. Heidarzadeh, K. R. Narayanan, and A. Sprintson, 'Sparse graph codes for non-adaptive quantitative group testing,' in IEEE Information Theory Workshop , 2019.                                                           |
| [45]   | --, 'Non-adaptive quantitative group testing using irregular sparse graph codes,' in Annual Allerton Conference on Communication, Control, and Computing , 2019.                                                                                 |
| [46]   | M. Hahn-Klimroth and N. Müller, 'Near optimal efficient decoding from pooled data,' Conference on Learning Theory , vol. 178, pp. 3395-3409, 2022.                                                                                               |
| [47]   | O. Gebhard, M. Hahn-Klimroth, D. Kaaser, and P. Loick, 'Information-theoretic and algorithmic as- pects of parallel and distributed reconstruction from pooled data,' Journal of Parallel and Distributed Computing , vol. 180, p. 104718, 2023. |
| [48]   | M. Soleymani and T. Javidi, 'A non-adaptive algorithm for the quantitative group testing problem,' in Proceedings of Thirty Seventh Conference on Learning Theory , 2024.                                                                        |
| [49]   | M. Hahn-Klimroth and D. Kaaser, 'Distributed reconstruction of noisy pooled data,' IEEE 42nd In- ternational Conference on Distributed Computing Systems , pp. 89-99, 2022.                                                                      |
| [50]   | Y.-H. Li and I.-H. Wang, 'Combinatorial quantitative group testing with adversarially perturbed mea- surements,' IEEE Information Theory Workshop , pp. 1-5, 2021.                                                                               |
| [51]   | M. Hahn-Klimroth, D. Kaaser, and M. Rau, 'Efficient approximate recovery from pooled data using doubly regular pooling scheme,' 2023, arXiv:2303.00043.                                                                                          |
| [52]   | M. Soleymani and T. Javidi, 'Quantitative group testing with tunable adaptation,' in IEEE Interna- tional Symposium on Information Theory , 2024.                                                                                                |
| [53]   | C. Rush, K. Hsieh, and R. Venkataramanan, 'Capacity-achieving spatially coupled sparse superposition codes with AMP decoding,' IEEE Transactions on Information Theory , vol. 67, no. 7, pp. 4446-4484, 2021.                                    |
| [54]   | G. Reeves and H. D. Pfister, 'The replica-symmetric prediction for random linear estimation with Gaussian matrices is exact,' IEEE Transactions on Information Theory , vol. 65, no. 4, pp. 2252-2283, 2019.                                     |
| [55]   | J. Barbier, N. Macris, M. Dia, and F. Krzakala, 'Mutual information and optimality of approximate message-passing in random linear estimation,' IEEE Transactions on Information Theory , vol. 66, no. 7,                                        |

- [56] P. Pascual Cobo, 'Code for spatially coupled AMP for pooled data and quantitative group testing,' https://github.com/PabloPasc/SC\_AMP\_QGT, 2024.
- [57] S. Foucart and H. Rauhut, An invitation to compressive sensing . Springer, 2013.
- [58] X. Liu, P. Pascual Cobo, and R. Venkataramanan, 'Many-user multiple access with random user activity,' in IEEE International Symposium on Information Theory (ISIT) , 2024.
- [59] R. Vershynin, High-dimensional probability: An introduction with applications in data science . Cambridge University Press, 2018, vol. 47.
- [60] Z. Bao, Q. Han, and X. Xu, 'A leave-one-out approach to approximate message passing,' 2023, arXiv:2312.05911.
- [61] R. Berthier, A. Montanari, and P.-M. Nguyen, 'State evolution for approximate message passing with non-separable functions,' Information and Inference: A Journal of the IMA , vol. 9, no. 1, pp. 33-79, 2020.
- [62] P. Nikolopoulos, S. R. Srinivasavaradhan, T. Guo, C. Fragouli, and S. N. Diggavi, 'Community-aware group testing,' IEEE Transactions on Information Theory , vol. 69, no. 7, pp. 4361-4383, 2023.
- [63] S. Ahn, W.-N. Chen, and A. Özgür, 'Adaptive group testing on networks with community structure: The stochastic block model,' IEEE Transactions on Information Theory , vol. 69, no. 7, pp. 4758-4776, 2023.
Nelvin Tan Pablo Pascual Cobo Ramji Venkataramanan
Quantitative Group Testing and Pooled Data in the Linear Regime with Sublinear Tests
In the pooled data problem, the goal is to identify the categories associated
with a large collection of items via a sequence of pooled tests. Each pooled
test reveals the number of items in the pool belonging to each category. A
prominent special case is quantitative group testing (QGT), which is the case
of pooled data with two categories. We consider these problems in the
non-adaptive and linear regime, where the fraction of items in each category is
of constant order. We propose a scheme with a spatially coupled Bernoulli test
matrix and an efficient approximate message passing (AMP) algorithm for
recovery. We rigorously characterize its asymptotic performance in both the
noiseless and noisy settings, and prove that in the noiseless case, the AMP
algorithm achieves almost-exact recovery with a number of tests sublinear in
the number of items. For both QGT and pooled data, this is the first efficient
scheme that provably achieves recovery in the linear regime with a sublinear
number of tests, with performance degrading gracefully in the presence of
noise. Numerical simulations illustrate the benefits of the spatially coupled
scheme at finite dimensions, showing that it outperforms i.i.d. test designs as
well as other recovery algorithms based on convex programming.