Compressed Sensing with Deep Image Prior and Learned Regularization

In plain words

Some sensors record only a few numbers about a picture instead of every pixel, and the job is to rebuild the picture from those numbers. The usual fixes need either a hand-designed image rule or a neural network trained on many similar pictures. This method uses neither. It takes a neural network that is still random, feeds it a fixed random input, and adjusts the network's internal settings until the picture it produces, when measured the same way as the real one, matches the recorded numbers. The network's structure already favours smooth, natural-looking images, so the result tends to look like a picture rather than noise.

The method also has a second, optional ingredient. If a handful of earlier pictures measured the same way are available, their fit gives a rough guess of what typical network settings look like. Nudging the new fit toward that guess helps most when the measurements are noisy or very few. Without it, nothing needs to be trained at all.

TL;DR

For a linear measurement process y = Ax + noise, fit the weights of an untrained DCGAN-style generator so that A applied to its output matches y, using a fixed random latent, total variation, and stopping after a fixed number of steps. Optionally add a quadratic penalty toward a Gaussian prior on the weights, estimated from about ten earlier measurement sets. The paper reports gains over unlearned baselines on MNIST, chest X-rays and retinopathy images, with runtime comparisons. All numbers are author-reported.

Problem

The task is to recover a signal x from m measurements y = Ax + η, with m smaller than the number of pixels n §3. Unlearned classical recovery methods (sparse bases, total variation, BM3D-AMP) are the baseline here §2 §5.1. Learned recovery usually needs a generator pre-trained on a very large dataset, with tens of thousands of similar images for CelebA-scale models, and that is often not available §2 §3.2.2. The method targets settings with little or no training data, and it assumes the measurement operator A is known and differentiable §3.1.

Method

Setup. Let n be the pixel count, A an m × n measurement matrix, and G(z; w) a generator with fixed input z and weights w. The paper fits the weights, not the latent code, which differs from the earlier approach of Bora et al. that optimises z through a pre-trained generator §3.1.

Step 1: data fit (CS-DIP). Start from an untrained generator with fixed random z and solve for weights that explain y §3.1 eq. 1. The paper argues that convolutional generators are biased toward smooth, natural images, so the architecture alone acts as a prior §3.1. Adding total variation weighted by λ_T, and a learned term weighted by λ_L, gives the final objective §3.1 eq. 2. Setting λ_L = 0 leaves a fully unlearned method §3.1.

w^* = \arg\min_w \lVert y - A\,G(z;w) \rVert^2 \qquad [\text{eq. 1}]

w^* = \arg\min_w \lVert y - A\,G(z;w) \rVert^2 + \lambda_T\,TV(G(z;w)) + \lambda_L\,LR(w) \qquad [\text{eq. 2}]

Step 2: learned regularisation as MAP. The likelihood of y given w is Gaussian with the standard quadratic data term, and the prior on w is a Gaussian with mean μ and covariance Σ §3.2. The MAP estimate then has the quadratic penalty LR(w) = (w − μ)ᵀ Σ⁻¹ (w − μ) with weight λ_L §3.2 eq. 3 eq. 4. With μ = 0 and Σ = I it reduces to ℓ2 weight decay §3.2.

LR(w) = (w-\mu)^{T}\,\Sigma^{-1}\,(w-\mu) \qquad [\text{eq. 4}]

Step 3: estimating (μ, Σ). Run Step 1 on Q measurement sets, each made with its own matrix A, from Q different images. Only the measurements are used, not ground truth §3.2.1. Each layer's weights are then modelled with one mean and one variance, and cross-layer covariances are set to zero §3.2.1. The estimate follows the procedure in alg. 1, which draws S weights per layer from a randomly chosen fitted network, computes per-round mean and covariance, and averages over T rounds alg. 1.

# Estimate a layer-wise prior (mu, Sigma) from fitted weight sets  [alg. 1]
# Input:  weights_list (Q fitted networks), S samples per layer, T rounds
# Output: mu (L,), Sigma (L, L)

def estimate_prior(weights_list, S, T):
    # Accumulate the per-round means and covariances  [alg. 1]
    mu_sum = zeros(L)
    sigma_sum = zeros(L, L)
    for t in range(T):
        # Pick one fitted network at random  [alg. 1]
        q = uniform_index(len(weights_list))
        # Sample S weights from each layer of that network  [alg. 1]
        samples = zeros(L, S)
        for l in range(L):
            samples[l, :] = sample_layer_weights(weights_list[q], l, S)
        mu_t = samples.mean(axis=1)                         # (L,)
        sigma_t = (samples @ samples.T) / S - outer(mu_t, mu_t)  # (L, L)
        mu_sum += mu_t
        sigma_sum += sigma_t
    # Average over the T rounds  [alg. 1]
    return mu_sum / T, sigma_sum / T

Step 3b: the fit used in experiments. Take the fixed random z, the untrained DCGAN and the objective of eq. 2. Optimise all network weights with RMSProp for a fixed number of steps and keep the final iterate §5.1 app. A. The pseudocode below follows the paper's description; the official code differs on several points, listed under Paper vs code.

# CS-DIP fit for one measurement vector  [§3.1] [eq. 2] [§5.1] [app. A]
# Input:  y (m,); A (m, n); prior (mu, Sigma) or None; lambda_T, lambda_L
# Output: x_hat (H, W, C), values in [-1, 1]

def cs_dip(y, A, prior, lambda_T, lambda_L, num_iters):
    # Fixed random latent input, never optimised  [§3.1] [app. A]
    z = sample_gaussian(shape=(z_dim,))
    # Untrained DCGAN with random initial weights  [§3.1] [§5.1]
    G = build_dcgan(z_dim)
    # RMSProp, learning rate 1e-3, momentum 0.9  [§5.1]
    opt = RMSProp(G.weights, lr=1e-3, momentum=0.9)
    for step in range(num_iters):
        x = G(z)                               # generated image
        # Measurement term ||y - A x||^2  [eq. 1]
        data = sum((y - A @ flatten(x)) ** 2)
        # Total variation penalty  [§3.1] [app. A]
        tv = total_variation(x)
        # Learned penalty toward (mu, Sigma), only if a prior is given  [eq. 4]
        lr = quadratic_penalty(G.weights, prior) if prior else 0
        loss = data + lambda_T * tv + lambda_L * lr
        opt.step(loss)
    # Return the output of the final iterate  [app. A]
    return G(z)

The latent dimension is 128 app. A, and the paper says 128 is a standard DCGAN choice app. A. Pixel values are in [−1, 1] §5.1.

Hyperparameters

Results

Author-reported. Setup: MNIST (first 100 test images), chest X-ray (60 random images from the Shenzhen dataset, 256 × 256), retinopathy (STARE, 128 × 128 × 3); baselines BM3D-AMP, TVAL3, Lasso with DCT basis; models in PyTorch on a GTX 1080 Ti for the proposed method and CPU-based baselines §5.1 app. B.

Learned regularisation (x-ray, λ_L = 100 vs 0). Table entries are percent decrease in error from adding learned regularisation, on chest X-rays with Gaussian measurements tab. 1 §5.2.1. The paper states that the gain is largest for noisy data and few measurements tab. 1 §5.2.1.

Runtime (x-ray, seconds). Lower is better. tab. 2 app. B

The paper does not give the MSE of the unlearned comparisons as numbers in the text. Its qualitative claims are these, each in §5.2.2 unless stated:

• MNIST: CS-DIP outperforms the unlearned baselines in almost all cases tested.
• Chest X-ray: CS-DIP beats all unlearned baselines except BM3D-AMP at higher sampling rates. Below about 3% sampling, BM3D-AMP often does not converge.
• Against the pre-trained DCGAN of Bora et al., the pre-trained method is better at small m, and CS-DIP is better at larger m, where the pre-trained method saturates §5.2.2 app. B.
• Fourier (radial-line) sampling on x-rays: CS-DIP is better than all baselines at the tested counts of 381 to 2500 coefficients app. B.
• Retinopathy: CS-DIP beats all baselines except BM3D-AMP at higher m, but green and purple patches remain app. B.
• Additive noise: comparable to the baselines app. B.

Paper vs code

Official repository: github.com/davevanveen/compsensing_dip, first author's account, commit 028b11cd27d0afc0026842f97b6d7abeb4099746 (git rev-parse HEAD on a depth-1 clone). The repository has no LICENSE file, so no reuse licence is stated. Only the pseudocode above is used here, with no verbatim code.

Found via the GitHub link in the paper's experiments section §5.1; the owner is the first author. No other code was needed.

Discrepancies found after reading the paper and the code:

• Latent dimension. The paper sets z to 128 for all experiments app. A. The code uses 64 for x-ray, 128 for MNIST and 32 for retinopathy configs.json configs.json configs.json.
• TV weight. The paper sets λ_T = 0.01 app. A. The code uses 0.01 for MNIST, 0.05 for x-ray and 0.02 for retinopathy utils.py.
• Learned-regularisation weight. The paper uses λ_L = 100 for the x-ray experiments app. A §5.2.1. The code uses 100 for x-ray, 1000 for retinopathy and 0 for MNIST utils.py. The MNIST value is 0, so the MNIST runs are unlearned in the code, although the paper's MNIST results are in the unlearned section §5.2.2.
• Measurement matrix per image. The paper draws one matrix per image app. A. The code calls get_A once per measurement count, before the image loop, so every image at a given m shares one matrix comp_sensing.py comp_sensing.py.
• Stopping rule. The paper stops after 1000 steps and reports that iterate app. A. The code keeps the 51 final steps (indices 949 to 999 of 1000) and returns the one with the lowest total loss, which includes the TV and learned terms, not only the measurement term cs_dip.py cs_dip.py.
• Restarts. The paper does not mention restarts. The code runs NUM_RESTARTS fits per measurement set and keeps the one with the lowest measurement loss. It uses 5 for MNIST and 1 for x-ray and retinopathy cs_dip.py cs_dip.py configs.json configs.json.
• Learned regularisation form. The paper's penalty is quadratic in all weights, with a covariance over layers and independent layers §3.2 §3.2.1. The code penalises the quadratic form of the vector of per-tensor means (layer.mean() over all parameters, which includes batch-norm parameters), and it uses the full inverse of the stored covariance matrix cs_dip.py utils.py.
• Estimating (μ, Σ). alg. 1 is not in the repository. The README says the generation of new (μ, Σ) "will be added soon", and the repository ships only the resulting arrays, for fixed measurement counts README.md .
• Lasso regularisation. The paper grid-searches the baseline parameters §5.1. The code uses a fixed value of 1e-5 baselines.py.
• TV form. The paper gives no formula. The code uses the sum of absolute differences along both image axes, without normalisation cs_dip.py.
• Noise and Fourier experiments. The noise variable in comp_sensing.py is set to zero comp_sensing.py, and the measurement matrix is Gaussian only baselines.py. The noise sweep (tab. 1) and the Fourier experiments have no code path in the shipped scripts.
• MATLAB baselines. BM3D-AMP and TVAL3 are in a MATLAB folder, and the Python entry point raises an error for them comp_sensing.py.

Pitfalls

• The theorem does not cover the experiments. Theorem 4.1 §4 has a single hidden ReLU layer, with the output layer V fixed and only the input layer W trained. It assumes A has orthonormal rows, and V is Gaussian with a variance set from y and z. The constant C in d ≥ Cm is not stated, and the probability bound has several terms. The experiments train a deep DCGAN with all weights, so the guarantee is only a motivation for early stopping §4.
• Early stopping is a choice, not a rule. The paper stops at 1000 steps, inside the plateau that the MSE reaches before it overfits, which the paper places near 10,000 steps for its setting app. A. This is a fixed step budget, not a validation criterion. Results depend on it. A different network, step size or number of measurements may move the plateau.
• Hyperparameters are tuned on measurements only. The paper selects parameters by the lowest error on the observed measurements, not on ground truth, using ten held-out images app. A. This does not guarantee that a setting that fits the measurements also generalises to other images.
• Random latent and small images. For small n (MNIST) the paper says the result varies with the random z; for large n it does not app. A.
• Learned prior needs ten measurement sets. The fit of (μ, Σ) uses the measurements of about ten images, each run through the full fit §3.2.2 app. A. The number of samples S and rounds T in alg. 1 are not stated.
• Matrix inversion. The code inverts the stored covariance directly, with no regularisation of the diagonal utils.py. A near-singular Σ would make the penalty unstable. The paper does not discuss this.
• Penalty scale. The code's data term is a sum over the m measurements cs_dip.py, and the λ values are set per dataset utils.py. Inferred from the code: changing m or the image size changes the balance between the data term and the penalties, so a λ_L from one setting may not transfer. The paper tunes λ_L per setting app. A.
• Pixel range. Images are in [−1, 1], and MSE can exceed 1 §5.1. Comparing numbers with papers that use [0, 1] requires rescaling.
• Runtime comparison is not like for like. The proposed method runs on a GPU and the baselines on CPU. The paper says so and does not claim a fair speed comparison app. B.
• Retinopathy artefacts. Some reconstructions show green and purple patches even at high m app. B.
• Fourier sampling. The paper's Fourier results are not reproducible from the shipped scripts (see Paper vs code).

Lineage

• Builds on: Deep Image Prior, Ulyanov et al. 2017, cited as [ulyanov2017deep], with the untrained-network idea. This paper changes the generator to a DCGAN and applies it to linear measurements §3.1.
• Compared against: Bora et al. (compressed sensing with a pre-trained DCGAN, optimising z) §2 §5.2.2; BM3D-AMP; TVAL3; Lasso with DCT and wavelet bases §5.1. No cards for these exist yet.
• Proof techniques borrowed from: Oymak et al. 2019 and Du et al. 2018 §4.
• Superseded by: not stated.