Exercises

$\mathbf{A}_{1} \otimes \mathbf{A}_{2} \in \mathbb{C}^{24 \times 64}$.

Full product: $24 \times 64 = 1{,}536$ entries. Factors: $4 \times 8 + 6 \times 8 = 32 + 48 = 80$ entries. Ratio: $1{,}536 / 80 = 19.2\times$.

$\mathbf{A}_{1} \otimes \mathbf{A}_{2} = \begin{pmatrix} a\mathbf{A}_{2} & b\mathbf{A}_{2} \\ c\mathbf{A}_{2} & d\mathbf{A}_{2} \end{pmatrix}$.

Multiplying by $\text{vec}(\mathbf{I}) = (1,0,0,1)^T$:

$$\begin{pmatrix} ae + b f \\ ag + bh \\ ce + df \\ cg + dh \end{pmatrix}.$$

$\mathbf{A}_{2} \mathbf{I} \mathbf{A}_{1}^{T} = \mathbf{A}_{2} \mathbf{A}_{1}^{T} = \begin{pmatrix} ea + fb & ec + fd \\ ga + hb & gc + hd \end{pmatrix}$.

Vectorizing column-by-column: $(ea + fb, ga + hb, ec + fd, gc + hd)^T$.

This equals the LHS. $\checkmark$

From TSpectral Properties of Kronecker Products, the SVD of $\mathbf{A}_{1} \otimes \mathbf{A}_{2}$ is $(\mathbf{U}_1 \otimes \mathbf{U}_2) (\boldsymbol{\Sigma}_1 \otimes \boldsymbol{\Sigma}_2) (\mathbf{V}_1 \otimes \mathbf{V}_2)^H$.

The singular values are $\{\sigma_i^{(1)} \sigma_j^{(2)}\}$.

$\|\mathbf{A}_{1} \otimes \mathbf{A}_{2}\| = \max_{i,j} \sigma_i^{(1)} \sigma_j^{(2)} = (\max_i \sigma_i^{(1)})(\max_j \sigma_j^{(2)}) = \|\mathbf{A}_{1}\| \cdot \|\mathbf{A}_{2}\|.$$$\blacksquare$

$L_f = \|\mathbf{A}\|^2 = (\|\mathbf{A}_{1}\| \cdot \|\mathbf{A}_{2}\|)^2 = \|\mathbf{A}_{1}\|^2 \cdot \|\mathbf{A}_{2}\|^2$.

$\mu = 1/L_f = 1/(\|\mathbf{A}_{1}\|^2 \|\mathbf{A}_{2}\|^2)$.

Adding a second factor multiplies $L_f$ by $\|\mathbf{A}_{2}\|^2$, which reduces the step size by the same factor. Larger sensing operators lead to smaller step sizes and potentially slower convergence.

import numpy as np

def kron_forward(c, A1, A2):
    N2, N1 = A2.shape[1], A1.shape[1]
    C = c.reshape(N2, N1)
    return (A2 @ C @ A1.T).ravel()

M1, N1, M2, N2 = 4, 8, 6, 8
A1 = np.random.randn(M1, N1) + 1j*np.random.randn(M1, N1)
A2 = np.random.randn(M2, N2) + 1j*np.random.randn(M2, N2)
c = np.random.randn(N1*N2) + 1j*np.random.randn(N1*N2)

y_naive = np.kron(A1, A2) @ c
y_fast = kron_forward(c, A1, A2)
print(np.allclose(y_naive, y_fast))  # True

Oversampled grid: $(\sigma N)^2 = (2 \times 128)^2 = 65{,}536$ complex entries. At 8 bytes each (complex64): $\approx 0.5$ MB.

$O((\sigma N)^2 \log(\sigma N)) = O(65{,}536 \times \log 256) = O(65{,}536 \times 8) \approx 5.2 \times 10^5$.

NUFFT total: $\sim 5.2 \times 10^5 + M \times W = 5.2 \times 10^5 + 2048 \times 12 \approx 5.5 \times 10^5$.

Direct: $M \times N^2 = 2048 \times 16{,}384 \approx 3.4 \times 10^7$.

Speedup: $\sim 60\times$.

$\mathbf{A}_{3} \times_1$: contract mode-1 of the tensor ($N_3$ dimension) with $\mathbf{A}_{3}$ ($M_3 \times N_3$). Cost: $M_3 \times N_3 \times N_2 \times N_1$.

After mode-1, the tensor is $M_3 \times N_2 \times N_1$. Mode-2 contraction with $\mathbf{A}_{2}$ ($M_2 \times N_2$): cost $M_3 \times M_2 \times N_2 \times N_1$.

After mode-2, the tensor is $M_3 \times M_2 \times N_1$. Mode-3 contraction with $\mathbf{A}_{1}$ ($M_1 \times N_1$): cost $M_3 \times M_2 \times M_1 \times N_1$.

Total: $N_1 N_2 N_3 M_3 + N_1 N_2 M_2 M_3 + N_1 M_1 M_2 M_3$.

For the balanced case $M_k = M$, $N_k = N$: factored cost $\sim 3 M N^3$, naive cost $M^3 N^3$. Speedup $= M^2 / 3$. For $M \geq 18$, speedup exceeds $100\times$.

$4096 \times 1.5 \times 10^9 \times 8 = 49.2$ real TFLOP/s $= 6.1$ complex TFLOP/s.

$32 \times 128 \times 128 = 524{,}288$ complex multiply-adds. Time: $524{,}288 / (6.1 \times 10^{12}) \approx 86$ ns.

Memory bandwidth limits: loading the two matrices ($32 \times 128 + 128 \times 128 = 20{,}480$ complex entries $= 160$ KB) at $\sim 1$ TB/s takes $\sim 160$ ns — already longer than the compute time. Plus kernel launch overhead ($\sim 5$--$10\;\mu$s) dominates for such small matrices.

CuPy mirrors the NumPy API almost exactly. An existing reconstruction pipeline written with numpy can often be GPU-accelerated by simply replacing import numpy as np with import cupy as np. No code restructuring needed.

Training an unrolled network requires backpropagation through the reconstruction iterations. PyTorch's autograd engine automatically builds and differentiates the computation graph. CuPy has no native AD support, so implementing backpropagation would require manual gradient derivation.

$n$ JVP evaluations, each costing $\sim 2$--$3 \times C_f$. Total: $\sim 2n \cdot C_f = 20{,}000 \cdot C_f$.

One forward pass ($C_f$) plus one backward pass ($\sim 2$--$4 \times C_f$, depending on the computation). Total: $\sim 3$--$5 \times C_f$. Independent of $n$.

Central differences: $f(x + h e_i)$ and $f(x - h e_i)$ for each $i = 1, \ldots, n$. Total: $2n \cdot C_f = 20{,}000 \cdot C_f$. Plus: subject to truncation and round-off error from the step size $h$.

Reverse AD is $\sim 4{,}000\times$ cheaper than forward AD or finite differences for this problem. This is why backpropagation is the standard in deep learning.

$T \times Q \times 16$ bytes (complex128) $= 50 \times 16{,}384 \times 16 \approx 13.1$ MB. In complex64: $\approx 6.6$ MB.

$\lceil\sqrt{50}\rceil = 8$ checkpoints. Memory: $8 \times Q \times 16 = 8 \times 16{,}384 \times 16 \approx 2.1$ MB (complex128), or $1.05$ MB (complex64). Reduction: $\sim 6.3\times$.

During backpropagation, each segment between checkpoints ($\sim 7$ iterations) must be recomputed. In the worst case, each of the $T = 50$ iterations is recomputed once (during the backward pass through its segment). Total forward computation: $\sim 2T$ (original $T$ + recomputed $T$). Overhead: $\sim 100\%$ of the original forward pass, but memory is reduced $6.3\times$.

Minimize $f(\mathbf{c}) + g(\mathbf{z})$ subject to $\mathbf{c} = \mathbf{z}$, where $f(\mathbf{c}) = \frac{1}{2}\|\mathbf{A}\mathbf{c} - \mathbf{y}\|^2$ and $g(\mathbf{z}) = \lambda\|\mathbf{z}\|_1$.

$r^{(t)} = \mathbf{c}^{(t)} - \mathbf{z}^{(t)}$ (constraint violation).

$s^{(t)} = \rho(\mathbf{z}^{(t)} - \mathbf{z}^{(t-1)})$ (since $\mathbf{B} = \mathbf{I}$, the formula simplifies).

$\mathbb{E}[\|\mathbf{w}\|^2] = M\sigma^2 = 2048 \times 0.0025 = 5.12$.

$\delta \approx \sigma\sqrt{M} = 0.05 \times \sqrt{2048} = 0.05 \times 45.25 \approx 2.26$.

$\tau\delta = 1.2 \times 2.26 = 2.72$.

Stop when $\|\mathbf{A}\hat{\mathbf{c}} - \mathbf{y}\| \leq 2.72$.

$T \approx \kappa \log_2(1/\varepsilon) = 5000 \times \log_2(10^4) \approx 5000 \times 13.3 \approx 66{,}500$.

Each iteration: $\sim 2 \times 10^6$ FLOPs (forward + adjoint). Total: $66{,}500 \times 2 \times 10^6 \approx 1.3 \times 10^{11}$.

Forming $\mathbf{A}_{1} \otimes \mathbf{A}_{2}$: $M_1 M_2 N_1 N_2 = 33.6 \times 10^6$ entries. Each iteration (dense matvec): $2 \times 33.6 \times 10^6 \approx 6.7 \times 10^7$. Total: $66{,}500 \times 6.7 \times 10^7 \approx 4.5 \times 10^{12}$.

Vec trick saves $\sim 35\times$.

The $q$-th diagonal entry with $q = (q_1 - 1)N_2 + q_2$:

$$[\mathbf{A}^{H}\mathbf{A}]_{q,q} = \|\text{col}_q(\mathbf{A})\|^2 = \|\text{col}_{q_1}(\mathbf{A}_{1})\|^2 \cdot \|\text{col}_{q_2}(\mathbf{A}_{2})\|^2.$$

This is exactly the $(q_1, q_2)$ entry of the outer product of the column-norm-squared vectors of $\mathbf{A}_{1}$ and $\mathbf{A}_{2}$.

1. Compute $\mathbf{d}_1 = \text{diag}(\mathbf{A}_{1}^{H}\mathbf{A}_{1})$ and $\mathbf{d}_2 = \text{diag}(\mathbf{A}_{2}^{H}\mathbf{A}_{2})$.
2. Compute adjoint: reshape $\mathbf{y}$ into $\mathbf{Y}$, then $\mathbf{A}_{2}^{H} \mathbf{Y} \overline{\mathbf{A}_{1}}$.
3. Divide element-wise by $\mathbf{d}_2 \mathbf{d}_1^T$.

Diagonals: $O(M_1 N_1 + M_2 N_2)$. Adjoint: $O(N_1(M_2 N_2 + M_1 M_2))$ (same as before). Division: $O(N_1 N_2)$. Total: dominated by the adjoint, $\sim 10^6$ FLOPs.

def ista_with_monitor(A, y, lam, mu, T):
    c = np.zeros(A.shape[1], dtype=complex)
    obj_prev = np.inf
    for t in range(T):
        grad = A.conj().T @ (A @ c - y)
        c_new = soft_threshold(c - mu * grad, mu * lam)
        obj = 0.5*np.linalg.norm(A@c_new - y)**2 \
              + lam*np.linalg.norm(c_new, 1)
        fpr = np.linalg.norm(c_new - c)
        data_res = np.linalg.norm(A @ c_new - y)
        if obj > obj_prev:
            print(f"WARNING: obj increased at t={t}")
        obj_prev = obj
        c = c_new
    return c

def soft_threshold(x, tau):
    return np.sign(x) * np.maximum(np.abs(x) - tau, 0)

FISTA evaluates the gradient at the extrapolated point $\mathbf{q}^{(t)}$ rather than the current iterate $\mathbf{c}^{(t)}$. The extrapolation is designed to accelerate convergence in the long run, but it can overshoot in the short run.

Specifically, $\mathbf{q}^{(t)}$ may lie in a region where $f(\mathbf{q}^{(t)}) > f(\mathbf{c}^{(t)})$, and the gradient step from $\mathbf{q}^{(t)}$ may produce $\mathbf{c}^{(t+1)}$ with $F(\mathbf{c}^{(t+1)}) > F(\mathbf{c}^{(t)})$.

Take $\mathbf{A} = \begin{pmatrix} 2 & 0 \\ 0 & 0.1 \end{pmatrix}$, $\mathbf{y} = (1, 0.1)^T$, $\lambda = 0.01$. The condition number $\kappa = 400$.

Starting from $\mathbf{c}^{(0)} = (0, 0)^T$ with $\mu = 1/\|\mathbf{A}\|^2 = 0.25$, FISTA's momentum causes the objective to oscillate for the first $\sim 10$ iterations before settling into monotone decrease.

• Fixed-point residual $\|\mathbf{c}^{(t+1)} - \mathbf{c}^{(t)}\|$.
• Data residual $\|\mathbf{A}\mathbf{c}^{(t)} - \mathbf{y}\|$.
• PSNR vs iteration (in simulation with ground truth).
• Constraint residual $\|\mathbf{c}^{(t)} - \mathbf{z}^{(t)}\|$.

If $\mathcal{D}_\sigma$ is firmly nonexpansive (i.e., $\|\mathcal{D}_\sigma(\mathbf{x}) - \mathcal{D}_\sigma(\mathbf{y})\|^2 + \|(\mathbf{I} - \mathcal{D}_\sigma)(\mathbf{x}) - (\mathbf{I} - \mathcal{D}_\sigma)(\mathbf{y})\|^2 \leq \|\mathbf{x} - \mathbf{y}\|^2$), then PnP ADMM converges to a fixed point and $\text{FPR}^{(t)} \to 0$.

Power iteration: starting from random $\mathbf{v}_0$, iterate $\mathbf{v}_{k+1} = \mathbf{J}_{\mathcal{D}}^T \mathbf{J}_{\mathcal{D}} \mathbf{v}_k / \|\cdots\|$.

Each iteration requires one JVP ($\mathbf{J}_{\mathcal{D}} \mathbf{v}_k$, via forward-mode AD) and one VJP ($\mathbf{J}_{\mathcal{D}}^T \mathbf{w}$, via reverse-mode AD).

Typically $\sim 20$ iterations suffice for convergence of the leading singular value.

In the far field, the steering vector for a UPA with $N_x \times N_y$ elements factors as $\mathbf{a}(\phi, \theta) = \mathbf{a}_{x}(\phi) \otimes \mathbf{a}_{y}(\theta)$. This gives the sensing matrix Kronecker structure.

In the near field, the phase $e^{+j\kappa\|\mathbf{p}_n - \mathbf{p}_q\|}$ depends on the full distance, which does not factor into separate $x$ and $y$ components.

1. Kronecker preconditioner: Use the far-field Kronecker factorization as a preconditioner for the near-field operator. This dramatically improves the condition number while each preconditioner application is fast (vec trick).

2. Low-rank correction: Write $\mathbf{A} = \mathbf{A}_{1} \otimes \mathbf{A}_{2} + \mathbf{E}$ where $\mathbf{E}$ is the near-field correction. If $\mathbf{E}$ has low rank (which it often does for moderate Fresnel numbers), the product $\mathbf{E}\mathbf{c}$ is cheap.

(a) No — uses only consecutive iterates. (b) Yes — PSNR requires the ground-truth image. (c) No — uses only the noise level $\delta$. (d) No — uses the primal and dual objectives. (e) No — uses only the objective function value.