Unrolled OAMP with ProxNet

Algorithm unrolling (or deep unfolding) converts a $T$-iteration algorithm into a $T$-layer neural network by:

1. Truncating the iteration to $T$ steps.
2. Parameterising each iteration's operators with learnable parameters $\theta_t$.
3. Training the parameters end-to-end by minimising

$$\min_{\theta_1, \ldots, \theta_T} \; \mathbb{E}_{(\mathbf{c}, \mathbf{y})}\bigl[ \mathcal{L}\bigl(\hat{\mathbf{c}}^{(T)}(\theta),\; \mathbf{c}\bigr) \bigr]$$

where $\hat{\mathbf{c}}^{(T)}$ is the output of the $T$-th layer and $\mathcal{L}$ is a task loss (e.g., MSE).

The Orthogonal AMP (OAMP) algorithm from Ch 17 iterates two steps:

Linear Estimation (LE): $$\mathbf{r}^{(t)} = \hat{\mathbf{c}}^{(t)} + \mathbf{W}_{\text{LE}}\bigl(\mathbf{y} - \mathbf{A}\hat{\mathbf{c}}^{(t)}\bigr)$$

Nonlinear Estimation (NLE): $$\hat{\mathbf{c}}^{(t+1)} = \eta_t\bigl(\mathbf{r}^{(t)}\bigr)$$

where $\mathbf{W}_{\text{LE}}$ satisfies the divergence constraint: $\frac{1}{N}\operatorname{tr}(\mathbf{W}_{\text{LE}} \mathbf{A}) = 1$. This ensures $\mathbf{r}^{(t)} = \mathbf{c} + \text{approximately Gaussian noise}$ in the large-system limit, regardless of the structure of $\mathbf{A}$.

Learned OAMP unrolls $T$ iterations of OAMP into a feedforward neural network.

Layer $t$ ($t = 1, \ldots, T$):

1. Linear module (LMMSE): Compute the linear estimate using the known $\mathbf{A}$ and the current noise estimate $\sigma_t^2$: $$\mathbf{r}^{(t)} = \hat{\mathbf{c}}^{(t)} + \mathbf{W}_{\text{LE}}^{(t)}\bigl(\mathbf{y} - \mathbf{A}\hat{\mathbf{c}}^{(t)}\bigr).$$
2. Orthogonalisation (fixed): The analytical orthogonalisation formula from Ch 17.3 ensures the divergence-free condition.
3. ProxNet denoiser: Replace the hand-designed $\eta_t$ with a neural network $\mathcal{D}_{\theta_t}(\mathbf{r}^{(t)};\, \sigma_t^2)$, typically a small CNN or DnCNN with parameters $\theta_t$.

Training: Minimise the end-to-end loss:

$$\mathcal{L}(\theta_1, \ldots, \theta_T) = \mathbb{E}\left[\|\hat{\mathbf{c}}^{(T)} - \mathbf{c}\|^2\right]$$

where the expectation is over training data $\{(\mathbf{y}_j, \mathbf{c}_{j})\}$.

ProxNet replaces the fixed denoiser $\eta_t$ in OAMP with a learned neural network $\mathcal{D}_{\theta_t}$:

$$\hat{\mathbf{c}}^{(t+1)} = \mathcal{D}_{\theta_t}\bigl(\mathbf{r}^{(t)};\; \sigma_t^2\bigr)$$

where $\sigma_t^2$ is the estimated noise variance at layer $t$ (provided by the state evolution). The network $\mathcal{D}_{\theta_t}$ is typically a small U-Net or DnCNN that takes $\mathbf{r}^{(t)}$ as input and produces a denoised estimate.

The divergence of $\mathcal{D}_{\theta_t}$ must be computed for the state evolution update:

$$\operatorname{div}(\mathcal{D}_{\theta_t}) = \frac{1}{N}\sum_{i=1}^N \frac{\partial [\mathcal{D}_{\theta_t}(\mathbf{r})]_i}{\partial r_i}.$$

In practice, Monte Carlo divergence estimation (Hutchinson's trace estimator) provides an efficient single-backward-pass approximation using a random probe vector $\mathbf{b}$: $\operatorname{div}(\mathcal{D}) \approx \mathbf{b}^T \nabla \mathcal{D} \cdot \mathbf{b}$.

When the sensing matrix has Kronecker structure $\mathbf{A} = \mathbf{A}_1 \otimes \mathbf{A}_2$ (common in 2D imaging with separable Tx/Rx arrays and frequency grids), the OAMP linear estimator decomposes as:

$$\mathbf{W}_{\text{LE}} = \mathbf{W}_1 \otimes \mathbf{W}_2$$

where $\mathbf{W}_i$ is the LMMSE estimator for $\mathbf{A}_i$: $\mathbf{W}_i = \mathbf{A}_i^H(\mathbf{A}_i\mathbf{A}_i^H + \sigma_t^2 \mathbf{I})^{-1}$.

This reduces the computational cost from $O(N^2)$ to $O(N^{3/2})$ (for square images with $N = n^2$), and from $O(N^2)$ to $O(N\log N)$ when $\mathbf{A}_i$ are partial DFT matrices (computed via FFT).

The state evolution for Kronecker-structured OAMP involves the joint singular value distribution of $\mathbf{A}_1$ and $\mathbf{A}_2$.