Exercises

$\hat{\mathbf{c}}^{(k+1)} = \mathcal{S}_{\lambda/L}\!\bigl((\mathbf{I} - \tfrac{1}{L}\mathbf{A}^{H}\mathbf{A})\hat{\mathbf{c}}^{(k)} + \tfrac{1}{L}\mathbf{A}^{H}\mathbf{y}\bigr).$$

At initialisation, the untrained LISTA network performs exact ISTA. Training then refines the parameters. $\square$

$\mathbf{c}$-update: $\hat{\mathbf{c}}^{(k+1)} = (\mathbf{A}^{H}\mathbf{A} + \rho\mathbf{I})^{-1}(\mathbf{A}^{H}\mathbf{y} + \rho(\mathbf{z}^{(k)} - \mathbf{u}^{(k)}))$

$\mathbf{z}$-update: $\mathbf{z}^{(k+1)} = \mathcal{S}_{\lambda/\rho}(\hat{\mathbf{c}}^{(k+1)} + \mathbf{u}^{(k)})$

Dual update: $\mathbf{u}^{(k+1)} = \mathbf{u}^{(k)} + \hat{\mathbf{c}}^{(k+1)} - \mathbf{z}^{(k+1)}$

The $\mathbf{c}$-update solves $(\mathbf{G} + \rho\mathbf{I})\mathbf{c} = \mathbf{b}$ where $\mathbf{G} = \mathbf{A}^{H}\mathbf{A}$. Since $\mathbf{G} + \rho\mathbf{I}$ is positive definite for $\rho > 0$, the solution is unique. For Kronecker-structured $\mathbf{A}$, this can be solved via FFT. $\square$

Per layer: $800^2 + 800 \times 200 + 1 = 800{,}001$. Total: $15 \times 800{,}001 = 12{,}000{,}015$ parameters.

Shared: $640{,}000 + 160{,}000 = 800{,}000$. Per-layer thresholds: $15$. Total: $800{,}015$ parameters --- a 15$\times$ reduction. $\square$

Input: Previous estimate $\hat{\mathbf{c}}^{(t-1)}$, measurements $\mathbf{y}$

LMMSE Block (uses $\mathbf{A}$, not learned): $\mathbf{r}^{(t)} = \hat{\mathbf{c}}^{(t-1)} + \mathbf{W}_{\text{LE}}^{(t)}(\mathbf{y} - \mathbf{A}\hat{\mathbf{c}}^{(t-1)})$

State Evolution (uses singular values of $\mathbf{A}$, not learned): Compute $\sigma_t^2$

ProxNet (does NOT use $\mathbf{A}$, LEARNED): $\hat{\mathbf{c}}^{(t)} = \mathcal{D}_{\theta_t}(\mathbf{r}^{(t)};\, \sigma_t^2)$

Output: Updated estimate $\hat{\mathbf{c}}^{(t)}$

Only the denoiser is learned. This is why unrolled OAMP has far fewer parameters than LISTA. $\square$

Layer 1: $\text{mmse}(2.0) = 0.1$. $\sigma_1^2 = (0.01 + 0.1)/0.5 = 0.22$.

Layer 2: $\text{mmse}(0.22) = 0.022$. $\sigma_2^2 = (0.01 + 0.022)/0.5 = 0.064$.

Layer 3: $\text{mmse}(0.064) = 0.0064$. $\sigma_3^2 = (0.01 + 0.0064)/0.5 = 0.0328$.

Layer 4: $\text{mmse}(0.0328) = 0.00328$. $\sigma_4^2 = (0.01 + 0.00328)/0.5 = 0.02656$.

Layer 5: $\text{mmse}(0.02656) = 0.00266$. $\sigma_5^2 = (0.01 + 0.00266)/0.5 = 0.02531$.

The trajectory converges to $\sigma_\infty^2 \approx 0.025$ (NMSE $\approx -16$ dB). Convergence is rapid in the first 3 layers and slows near the fixed point. A ProxNet denoiser achieving lower mmse at each step would shift the fixed point downward. $\square$

$\frac{\partial [\mathcal{S}_\tau(\mathbf{z})]_i}{\partial z_j} = \delta_{ij} \cdot \mathbf{1}_{|z_i| > \tau}.$$

$\frac{\partial [\mathcal{S}_\tau(\mathbf{z})]_i}{\partial \tau} = -\operatorname{sign}(z_i) \cdot \mathbf{1}_{|z_i| > \tau}.$$

$\frac{\partial \mathcal{L}}{\partial \tau_k} = \frac{\partial \mathcal{L}}{\partial \hat{\mathbf{c}}^{(K)}} \left(\prod_{j=k+1}^{K} \frac{\partial \hat{\mathbf{c}}^{(j)}}{\partial \hat{\mathbf{c}}^{(j-1)}}\right) \frac{\partial \hat{\mathbf{c}}^{(k)}}{\partial \tau_k}.$$The intermediate Jacobians involve products of$\mathbf{W}j$and the diagonal mask$\operatorname{diag}(\mathbf{1}{|z_i^{(j)}| > \tau_j})$. Gradient flow is through the **active set** only.$\square$

$\mathbf{W}_{\text{LE}}\mathbf{A} = \mathbf{V}\boldsymbol{\Sigma}^H(\boldsymbol{\Sigma}\boldsymbol{\Sigma}^H + v\mathbf{I})^{-1}\boldsymbol{\Sigma}\mathbf{V}^H$.

$\frac{1}{N}\operatorname{tr}(\mathbf{W}_{\text{LE}}\mathbf{A}) = \frac{1}{N}\sum_{i=1}^{\min(M,N)} \frac{\sigma_i^2}{\sigma_i^2 + v}.$$This is a smooth function of$v$ranging from$M/N$(as$v \to 0$) to$0$(as$v \to \infty$).$\square$

$\mathbf{A} \in \mathbb{C}^{M \times N}$ with $M = M_1 M_2$, $N = N_1 N_2$. Cost of $\mathbf{A}\mathbf{c}$: $O(MN) = O(M_1 M_2 N_1 N_2)$.

Reshape $\mathbf{c}$ into $\mathbf{C} \in \mathbb{C}^{N_2 \times N_1}$. Then: $\operatorname{vec}((\mathbf{A}_1 \otimes \mathbf{A}_2)\mathbf{c}) = \operatorname{vec}(\mathbf{A}_2 \mathbf{C} \mathbf{A}_1^T)$.

Cost: $O(N_1(M_2 N_2 + M_1 M_2))$.

For square problems ($M_1 = M_2 = m$, $N_1 = N_2 = n$): without $O(m^2 n^2)$, with $O(mn(m+n))$. Speedup $\approx O(n/2)$ for $m \approx n$. For $n = 256$: $\sim 128\times$ speedup. $\square$

$\frac{\partial \mathcal{L}}{\partial \theta_1}$ passes through $K-1$ Jacobians. If each has spectral norm $\rho < 1$, the gradient decays as $\rho^{K-1}$. For $K = 20$, $\rho = 0.9$: $0.9^{19} \approx 0.14$.

The direct term $w_k \partial \mathcal{L}(\hat{\mathbf{c}}^{(k)})/\partial \theta_k$ provides a gradient signal that bypasses the vanishing chain.

Use linearly increasing weights: $w_k = k/K$. This gives the most weight to the final output while providing non-trivial supervision at every layer. Alternative: $w_k = \alpha^{K-k}$ with $\alpha = 0.5$. $\square$

$L = \sigma_{\max}^2 = 100$. Step size $1/L = 0.01$. Contraction rate: $\rho_{\text{ISTA}} = 1 - \sigma_{\min}^2/L = 1 - 0.01/100 = 0.9999$. After 100 iterations: $0.9999^{100} \approx 0.99$ --- almost no progress.

OAMP applies per-singular-value weights. For the bulk of singular values near $\sigma_i \approx 5$, the weight $w_i = 5/(25 + v)$ is much larger than $1/L = 0.01$.

The effective noise after the LMMSE step involves $\frac{1}{N}\sum_i v/(\sigma_i^2 + v)$, which is dominated by the bulk spectrum, not the extreme $\sigma_{\min} = 0.1$.

ISTA is bottlenecked by the worst singular value; OAMP adapts to the full spectrum. For this matrix, OAMP converges in $\sim 10$ iterations while ISTA needs $> 10{,}000$. $\square$

$\operatorname{prox}_R(\mathbf{r}) = \arg\min_{\mathbf{c}} \frac{1}{2}\|\mathbf{c} - \mathbf{r}\|^2 + \sum_i \tau_i |c_i|$.

The objective separates across components: $\min_{c_i} \frac{1}{2}(c_i - r_i)^2 + \tau_i |c_i|$.

The solution is $c_i^* = \mathcal{S}_{\tau_i}(r_i) = \operatorname{sign}(r_i)\max(|r_i| - \tau_i, 0)$.

This is exactly the spatially-varying soft-thresholding operator. Therefore HST $= \operatorname{prox}_R$ for the weighted $\ell^1$ norm $R(\mathbf{c}) = \sum_i \tau_i |c_i|$. $\square$

$\|\mathcal{S}_\tau(\mathbf{u}) - \mathcal{S}_\tau(\mathbf{v})\| \leq \|\mathbf{u} - \mathbf{v}\|$.

$\|\mathcal{S}_\tau(\mathbf{W}\mathbf{a} + \mathbf{S}\mathbf{y}) - \mathcal{S}_\tau(\mathbf{W}\mathbf{b} + \mathbf{S}\mathbf{y})\| \leq \|\mathbf{W}(\mathbf{a} - \mathbf{b})\| \leq \|\mathbf{W}\|\,\|\mathbf{a} - \mathbf{b}\|.$$If$|\mathbf{W}| < 1$, each layer contracts by$\rho = |\mathbf{W}|$.

By Banach's theorem, the unique fixed point satisfies: $\mathbf{c}^* = \mathcal{S}_\tau(\mathbf{W}\mathbf{c}^* + \mathbf{S}\mathbf{y})$. The $K$-layer network approximates this with error $\leq \rho^K \|\hat{\mathbf{c}}^{(0)} - \mathbf{c}^*\|$. $\square$

ADMM is a special case of Douglas-Rachford splitting. If both proximal operators are firmly nonexpansive, the Douglas-Rachford operator is averaged, guaranteeing convergence.

If $\mathcal{G}_{\theta_k^z}$ is firmly nonexpansive and $\mathcal{F}_{\theta_k^x}$ is the true proximal of a convex function, both are resolvents of maximal monotone operators. The ADMM convergence theory applies layer by layer, with varying $\rho_k$ as preconditioning. $\square$

After each gradient update, compute $\sigma_1(\mathbf{W}_k)$ via power iteration (1--2 steps). Normalise: $\bar{\mathbf{W}}_k = (1-\epsilon)\mathbf{W}_k / \max(\sigma_1(\mathbf{W}_k), 1-\epsilon)$. This ensures $\|\bar{\mathbf{W}}_k\| \leq 1 - \epsilon < 1$.

Convergence rate: $\prod_{k=1}^K \|\mathbf{W}_k\| \leq (1-\epsilon)^K$. For $\epsilon = 0.05$, $K = 20$: $(0.95)^{20} \approx 0.36$.

Smaller $\epsilon$: faster per-layer but weaker contraction guarantee. Larger $\epsilon$: strong contraction but limits expressivity. $\epsilon = 0.01$--$0.1$ works in practice. $\square$

Each $\frac{\partial [\mathcal{D}_\theta]_i}{\partial r_i}$ requires one backward pass with seed $\mathbf{e}_i$. Total: $N$ backward passes. For $N = 65{,}536$ ($256 \times 256$ image), this is prohibitive.

Draw $\mathbf{b} \sim \mathcal{N}(0, \mathbf{I}_N)$. $\widehat{\operatorname{div}} = \mathbf{b}^T \mathbf{J}_\theta(\mathbf{r}) \mathbf{b}$. This requires only one backward pass via vector-Jacobian product.

$\mathbb{E}[\widehat{\operatorname{div}}] = \operatorname{tr}(\mathbf{J}) = \operatorname{div}(\mathcal{D}_\theta)$. $\operatorname{Var}(\widehat{\operatorname{div}}) = 2\|\mathbf{J}\|_F^2$ for symmetric $\mathbf{J}$. Averaging $L$ independent draws reduces variance by $1/L$. $\square$

Total parameters: $P = K \cdot P_{\text{layer}}$. Generalisation gap: $\propto \sqrt{K \cdot P_{\text{layer}} / n}$.

Total parameters: $P = P_{\text{layer}} + K \approx P_{\text{layer}}$. Gap: $\propto \sqrt{P_{\text{layer}} / n}$.

Ratio of gaps: $\sqrt{K} = \sqrt{20} \approx 4.5\times$. Weight tying reduces the generalisation gap by a factor of $\sim 4.5$ for $K = 20$ layers. $\square$

After support identification, the contraction rate on support $S$ is $\rho(S) = \|(\mathbf{I} - \mathbf{W}\mathbf{G})_S\|$. We seek: $\mathbf{W}^* = \arg\min_\mathbf{W} \max_{|S| = s} \rho(S)$.

Upper-bounding by $\|\mathbf{I} - \mathbf{W}\mathbf{G}\|$ and writing $\mathbf{G} = \mathbf{V}\boldsymbol{\Lambda}\mathbf{V}^H$, the optimal $\mathbf{W}$ has the Tikhonov form $w_i = \lambda_i/(\lambda_i + \mu)$.

$\mathbf{W}^* = \mathbf{I} - \mathbf{A}^{H}(\mathbf{A}\mathbf{A}^{H} + \mu\mathbf{I})^{-1}\mathbf{A}$.

The equation for $\mu^*$: $\frac{1}{N}\sum_{i=1}^N \frac{\lambda_i^2}{(\lambda_i + \mu)^2} = \frac{N - s}{N}$.

This balances contraction on the support (small $\mu$) with stability off the support (large $\mu$). $\square$

For the ALISTA weight with optimal $\mu$, the restricted matrix satisfies $\|(\mathbf{I} - \mathbf{W}\mathbf{G})_S\| \leq \delta_{2s}$ for any support $S$ with $|S| \leq 2s$. After soft-thresholding, the effective contraction rate involves $\delta_{2s}/(1 - \delta_{2s})$.

The approximation error from non-exact sparsity contributes $C\|\mathbf{c}^* - \mathbf{c}^*_s\|_1/\sqrt{s}$ (standard LASSO bound). The noise term is $D\|\mathbf{W}\mathbf{A}^{H}\mathbf{w}\| \leq D\sqrt{\sigma^2}$.

By induction on $K$, the error after $K$ layers satisfies: $\|\hat{\mathbf{c}}^{(K)} - \mathbf{c}^*\| \leq \rho^K \|\hat{\mathbf{c}}^{(0)} - \mathbf{c}^*\| + \frac{1}{1-\rho}(C_{\text{tail}} + C_{\text{noise}})$. For $\delta_{2s} < \sqrt{2}-1$: $\rho < 1$. $\square$

Since groups $\{\mathcal{G}_\ell\}$ partition $\{1, \ldots, N\}$: $\operatorname{prox}_R(\mathbf{r}) = \bigoplus_\ell \operatorname{prox}_{\lambda_\ell \|\cdot\|_2}(\mathbf{r}_{\mathcal{G}_\ell})$.

For group $\ell$: $\operatorname{prox}_{\lambda_\ell \|\cdot\|_2}(\mathbf{r}_\ell) = \left(1 - \frac{\lambda_\ell}{\|\mathbf{r}_\ell\|_2}\right)_+ \mathbf{r}_\ell$.

This is block soft-thresholding: if $\|\mathbf{r}_\ell\|_2 < \lambda_\ell$, the entire group is set to zero. Otherwise, the group is shrunk toward zero by factor $\lambda_\ell / \|\mathbf{r}_\ell\|_2$.

Entire angular groups with small energy (noise only) are suppressed. Groups with signal energy are preserved and shrunk. The per-group thresholds $\lambda_\ell$ can be learned (as in HST) to adapt to the expected energy distribution. $\square$

Generate 10{,}000 synthetic scenes:

• 5{,}000 with 5--20 point scatterers (varying amplitudes)
• 3{,}000 with extended targets (rectangles, circles)
• 2{,}000 with mixed point + extended targets

Add noise at SNR $\in [5, 30]$ dB (uniformly sampled). Split: 8{,}000 train / 1{,}000 val / 1{,}000 test.

Combined loss: $\mathcal{L} = \|\hat{\mathbf{c}} - \mathbf{c}\|^2 + 0.1 \cdot \mathcal{L}_{\text{SE}}$

where $\mathcal{L}_{\text{SE}} = \sum_{t=1}^T (\hat{\sigma}_t^2 - \sigma_t^{\text{SE}})^2$ regularises the learned noise variances against the state evolution prediction.

Phase 1: Layer-wise pre-training (5 epochs per layer, learning rate $10^{-3}$, Adam optimiser).

Phase 2: End-to-end fine-tuning (50 epochs, learning rate $10^{-4}$, cosine annealing, batch size 16).

Gradient checkpointing every 3 layers: store activations at layers $\{0, 3, 6, 9\}$ and recompute the rest. Memory: $\sim 4 \times 128^2 \times 8$ bytes $\times 4$ checkpoints $= 2$ MB (negligible vs. ProxNet activations).

In-distribution: same scene types and SNR range. Out-of-distribution:

• Different SNR range ($[0, 5]$ dB and $[30, 40]$ dB)
• Different scene types (Shepp-Logan phantom, letter shapes)
• Different number of scatterers ($> 20$)

Report NMSE, SSIM, and per-pixel uncertainty if available. $\square$