Deep unfolding

Overview

Deep unfolding, also called algorithm unrolling, takes an iterative solver for an inverse problem and runs a fixed number of its iterations as network stages. Each stage keeps the solver's structure: a step that uses the measurement model to fit the data, followed by a prior step. The learned modules replace the hand-set parts. The chain is then trained end to end 2204.13348 §3.1 2204.13348 §3.2 1704.02422 §4.

The target is a linear measurement model y = A x + n, where A is a degradation operator and n is noise. The problem is ill-posed 2204.13348 eq. 1 2204.13348 §1. The MRI card uses the same form, with A as a masked Fourier operator F_u and e as noise 1704.02422 eq. 1. Unfolding keeps the measurement model inside the network. In the restoration card each stage applies A and its transpose, or a learned stand-in for them [arxiv:2204.13348 eq. 6] 2204.13348 eq. 7. In the MRI card the measurement enters through the sampled k-space values and the mask 1704.02422 §3. The invariant all variants share is this split into a data-fit step that uses the measurement model and a learned prior step.

Two cards on wikillm cover this concept. One is for image restoration: deraining, deblurring, denoising and compressive sensing. The other is for dynamic cardiac MRI reconstruction. Everything below is what those sources state or report. Other unfolding work is mentioned in the restoration paper's related-work section, but it has no card here, so it is not summarised.

Variants

Known operator

When A is known, the gradient step uses A and its transpose directly, with a trainable step size per stage 2204.13348 eq. 6 2204.13348 §3.2. This is the unrolled form of Proximal Gradient Descent 2204.13348 eq. 4 2204.13348 eq. 5a. With an l1 regulariser the proximal map is soft-thresholding, which gives ISTA 2204.13348 §3.1. The paper treats A as exactly known in compressive sensing, where the sampling matrix is learned jointly 2204.13348 §4.5. Card: Deep Generalized Unfolding Networks for Image Restoration.

Closed-form data consistency

In the MRI cascade the data-fit step is not a gradient step. A data-consistency (DC) layer returns the CNN's k-space values where nothing was measured, and uses the scanner's values where it sampled 1704.02422 §3 1704.02422 eq. 6. In noise-free experiments the paper takes the limit in which the sampled entries are replaced exactly 1704.02422 §3. The step follows from the closed-form update of the DLMRI dictionary-learning method that the cascade unrolls 1704.02422 §2 1704.02422 §3. A CNN block is followed by a DC layer, and the pair repeats n_c times 1704.02422 §4. The DC layer is a fixed linear map, so gradients pass through it to the CNN blocks 1704.02422 §3 1704.02422 eq. 10. Card: CNN cascade for dynamic MRI.

Learned operator pair

When A is unknown, the paper replaces A and its transpose with two independent residual blocks, one for each direction, and uses their composition in the gradient step 2204.13348 eq. 7 2204.13348 §3.2. The paper says this is for unknown degradations but does not list the tasks that use it 2204.13348 §3.2.

Stage sharing

DGUNet unfolds K = 7 stages. Stages 2 to 6 share weights, and stages 1 and 7 have their own. DGUNet+ gives every stage independent weights 2204.13348 §3.2. In the paper's ablation on Rain100H, DGUNet+ scores 29.67 dB with 3 stages, 30.43 dB with 5, 31.06 dB with 7 and 31.27 dB with 9 2204.13348 tab. 5. The MRI card does not state whether its CNN blocks share weights across stages. Its greedy study starts each deeper cascade from the shallower one 1704.02422 §7.2.2.

Inter-stage features

Each stage also passes features to the next. The paper fuses the previous stage's encoder and decoder features into the current encoder features with a spatially varying scale and shift 2204.13348 eq. 8 [arxiv:2204.13348 §3.2]. Ablations on Rain100H, all on DGUNet+ and author-reported, give 31.06 dB for the full model, 30.09 dB without this fusion, 30.92 dB with plain addition instead of spatial scale and shift, and 30.51 dB with the learned gradient module replaced by a plain U-Net cascade 2204.13348 tab. 6.

Data sharing across frames

For dynamic sequences, the MRI paper's data-sharing layer averages the k-space samples of neighbouring frames at the locations the current frame missed. It feeds the averages to the CNN as extra input channels 1704.02422 §5 1704.02422 §6. The layer applies to sequences only, not 2D frames 1704.02422 §6. In the 9-fold 3D ablation, the version with data sharing has lower train and test error than the version without it. The paper reports this only in a figure 1704.02422 fig. 10 1704.02422 §7.3.1.

When to use which

Implementation sketch

Our own pseudocode of the common core, with notes on where the variants differ. Citations in the comments point to the paper's equations.

# Unfolded network with K stages  [arxiv:2204.13348 §3.2] [arxiv:1704.02422 §4]
# Input:  degraded image y
# Output: one estimate per stage; the last one is the result

def unfolded_restore(y, K=7):
    # Start from the input with no features from earlier stages
    x, feats = y, None
    outputs = []
    for k in range(K):
        # Data step: a gradient (restoration) or a closed-form
        # replacement of sampled values (MRI)  [eq. 6] [arxiv:1704.02422 §3]
        v = x - step[k] * data_gradient(x, y, k)
        # Learned prior step; also returns features for stage k + 1
        x, feats = prox(v, feats, k)
        outputs.append(x)
    return outputs

The MRI variant replaces the two data-step lines with a DC layer, and its prior step is a residual CNN block with no features passed between stages 1704.02422 §4 1704.02422 §6.

# Gradient term with a known operator A  [arxiv:2204.13348 eq. 6]
def grad_known(x, y, A):
    # Residual in measurement space, mapped back with A transpose
    return A.T @ (A @ x - y)

# Gradient term with a learned pair  [arxiv:2204.13348 eq. 7]
def grad_learned(x, y, f_forward, f_back):
    # f_forward stands in for A, f_back for A transpose
    return f_back(f_forward(x) - y)

# Training loss over all stage outputs  [arxiv:2204.13348 eq. 10]
def loss(outputs, x_clean):
    # Squared error on every stage, not only the last one
    return sum(squared_error(o, x_clean) for o in outputs)

Contradictions

• Within the paper. The introduction and conclusion say twelve test sets 2204.13348 §1 2204.13348 §5. The named sets in the experiments total eleven: five deraining, two deblurring, two denoising and two compressive-sensing 2204.13348 §4.2 2204.13348 §4.3 2204.13348 §4.4 2204.13348 §4.5.
• Paper versus official code. The paper states an l2 loss summed over stages 2204.13348 eq. 10. The official deraining code uses a Charbonnier loss plus an edge term [arxiv:2204.13348 code: Deraining/train.py#L122-L127]. The fusion in the code is a different function from eq. 8 2204.13348 code: Deraining/DGUNet.py#L174-L178. A network built from the paper's equations is not the code's network.
• Evaluation caveat. The deraining checkpoint is chosen by PSNR on Rain100L 2204.13348 code: Deraining/train.py#L148-L156, which is also one of the reported test sets 2204.13348 tab. 1. The paper does not say which checkpoint it reports, so treat the Rain100L numbers as possibly optimistic.
• Noisy DC branch, paper versus code. Eq. 6 leaves unsampled entries unchanged in the noisy case 1704.02422 §3. The official code scales them by 1/(1+λ) too [arxiv:1704.02422 code: cascadenet/network/layers/data_consistency.py#L68-L69]. The noise-free branch matches the paper.
• Trainable λ, paper versus code. The paper trains λ in the noise experiments 1704.02422 §3 1704.02422 §7.3.3. The code passes λ as a fixed Python float, and the 3D build passes none, so it always takes the noise-free branch 1704.02422 code: cascadenet/network/model.py#L52 1704.02422 code: cascadenet/network/model.py#L105-L106.
• Test-time patches, paper versus code. The paper applies the 3D CNN to the full volume at test time 1704.02422 §7.3. The demo's test loop uses fixed 32-column patches 1704.02422 code: main_3d.py#L74.
• Not comparable across the two cards. The restoration results are PSNR in dB, and the MRI results are MSE ×10⁻³ on a different task. Their numbers cannot be compared 2204.13348 tab. 1 1704.02422 tab. 1.

Papers

• CNN cascade for dynamic MRI: a CNN cascade with closed-form data-consistency layers and a data-sharing layer for dynamic cardiac MRI.
• Deep Generalized Unfolding Networks for Image Restoration: DGUNet and DGUNet+, unfolded proximal gradient descent with learned gradient modules and inter-stage fusion, evaluated on deraining, deblurring, denoising and compressive sensing.