Deep Generalized Unfolding Networks for Image Restoration

In plain words

Image restoration means taking a damaged picture (rainy, blurry, noisy, or rebuilt from a few measurements) and recovering the clean image. A classic approach writes this as an optimisation problem and solves it step by step, alternating between "match what was measured" and "look like a natural image". Plain neural networks restore images well but work as black boxes.

Deep unfolding keeps the steps of the classic procedure and turns each step into a learnable layer, so the whole chain trains end to end and still follows the procedure. This paper's network, DGUNet, adds two ideas. First, it learns the "match what was measured" step even when the exact way the picture was damaged is unknown. Second, it lets the intermediate features of one step pass into the next, so detail that is normally lost when a stage turns features back into an image is kept.

The authors test it on deraining, deblurring, denoising and compressive sensing and report higher scores than the methods they compare with on the deraining, deblurring and denoising test sets.

TL;DR

DGUNet is a 7-stage network built by unfolding Proximal Gradient Descent §3.1 §3.2. Each stage has a gradient module, which uses the known degradation matrix when it exists and two learned residual blocks otherwise eq. 6 eq. 7, and an informative proximal mapping module, an encoder-decoder with channel attention §3.2. Inter-stage feature fusion carries encoder and decoder features across stages through spatially varying scale and shift eq. 8. Author-reported: DGUNet and DGUNet+ (independent stages) beat the compared methods on the deraining, deblurring and denoising test sets, and DGUNet+ is best on most of them tab. 1 tab. 2 tab. 3. Compressive sensing results are mixed tab. 4. The official code differs from the paper in several details, so read the Paper vs code section before implementing.

Problem

Image restoration is modelled as y = A x + n, where A is the degradation matrix and n is noise, and is ill-posed eq. 1 §1. Model-based methods use handcrafted priors that limit their representation, and deep plug-and-play methods are slow at inference §1. Deep unfolding networks (DUN) train end to end and infer quickly, but usually need a known degradation process or handcrafted degradation assumptions. Their stage-to-stage interface is an image, so features are distorted at every stage §1 §2.

The paper states that A is exactly known in compressive sensing §4.5. It states that A is unknown in some degradation problems §3.2 without listing which ones. The official deraining code uses only the learned branch (see Paper vs code).

Method

Proximal Gradient Descent (PGD) splits each iteration into a gradient step and a proximal step eq. 4. The gradient step is v = x - rho A^T (A x - y) and the proximal step is x = prox(v) eq. 5a eq. 5b §3.1. ISTA is the special case with an l1 regulariser, whose proximal map is soft-thresholding §3.1.

DGUNet unfolds K PGD iterations into K stages. The default is K = 7 §3.2. Stages 2 to 6 share parameters, and stages 1 and 7 have their own. The plus version, DGUNet+, makes all stages independent §3.2.

1. Flexible gradient descent module (FGDM), one per stage §3.2 fig. 2. a. If A is known, use A and A^T directly, with a trainable step size rho^k per stage eq. 6 §3.2. b. If A is unknown, replace A and A^T with two independent residual blocks F_A^k and F_{A^T}^k, and use F_{A^T}^k(F_A^k(x) - y) as the gradient eq. 7 §3.2.
2. Informative proximal mapping module (IPMM), one per stage §3.2 fig. 2. a. A channel attention block extracts shallow features §3.2. b. Residual blocks without batch normalisation extract features at three scales, down with 2x2 max-pooling and up with bilinear upsampling followed by a convolution §3.2. c. A global pathway adds the input to the output §3.2. d. A supervised attention module (SAM, from MPRNet) produces the stage output and clean features §3.2. e. The features are injected into the next stage through a subspace projection (from NBNet) §3.2.
3. Inter-stage feature fusion (ISFF), at each encoder scale §3.2 fig. 3. Previous-stage encoder and decoder features are embedded and summed, then produce a spatial scale and shift for the current encoder features eq. 8.
4. Loss: an l2 error summed over the outputs of all stages, with no other terms §3.3 eq. 10.

Key equations:

\mathbf{v}^k = \hat{\mathbf{x}}^{k-1} - \rho^k \mathbf{A}^\top(\mathbf{A}\hat{\mathbf{x}}^{k-1} - \mathbf{y}) eq. 6

\mathbf{v}^k = \hat{\mathbf{x}}^{k-1} - \rho^k \mathcal{F}^k_{\mathbf{A}^\top}\big(\mathcal{F}^k_{\mathbf{A}}(\hat{\mathbf{x}}^{k-1}) - \mathbf{y}\big) eq. 7

\begin{aligned} \mathbf{H}^{k-1}_n &= \mathrm{Conv}(\mathbf{F}^{k-1}_{enc,n}) + \mathrm{Conv}(\mathbf{F}^{k-1}_{dec,n}) \\ \bm{\alpha}^k_n, \bm{\beta}^k_n &= \mathrm{Conv}_\alpha(\mathbf{H}^{k-1}_n), \mathrm{Conv}_\beta(\mathbf{H}^{k-1}_n) \\ \mathbf{F}^k_{enc,n} &= \hat{\mathbf{F}}^k_{enc,n} \odot \bm{\alpha}^k_n + \bm{\beta}^k_n \end{aligned} eq. 8

\hat{\mathbf{x}}^k, \mathbf{F}^k = \mathtt{prox}_{\bm{\theta}^k}(\mathbf{v}^k, \mathbf{F}^{k-1}) eq. 9

\mathcal{L}(\bm{\Omega}) = \sum_{k=1}^{K} \left\| \mathbf{x} - \hat{\mathbf{x}}^k \right\|_2^2 eq. 10

The algorithm, as printed, is:

# DGUNet forward pass  [alg. 1]
# Input:  degraded image y (B, 3, H, W)
# Output: estimates of stages 1..K, each (B, 3, H, W)

def dgunet(y):
    # Start from the input with no inter-stage features  [alg. 1]
    estimate = y
    inter_stage = None
    estimates = []
    for k in range(1, K + 1):
        # Gradient step with step size rho_k  [§3.2] [eq. 6] [eq. 7]
        # Use the known A when it exists, else the learned pair
        v = gradient_step_k(estimate, y)     # (B, 3, H, W)
        # Proximal mapping; also returns features for stage k + 1  [eq. 9]
        estimate, inter_stage = prox_k(v, inter_stage)
        estimates.append(estimate)
    return estimates

# One proximal mapping (IPMM) with inter-stage fusion  [§3.2] [eq. 8]
# Input:  v (B, 3, H, W); features from the previous stage
# Output: estimate (B, 3, H, W); features for the next stage

def prox_k(v, inter_stage):
    # Shallow features with channel attention  [§3.2]
    shallow = channel_attention(v)
    # Encoder at three scales, modulated by the previous stage  [eq. 8]
    encoder_feats, bottleneck = encoder(shallow, inter_stage)
    # Decoder with skip connections  [§3.2]
    decoder_feats = decoder(bottleneck, encoder_feats)
    # Supervised attention gives the estimate and clean features  [§3.2]
    clean_feats, estimate = sam(decoder_feats[-1], v)
    return estimate, (encoder_feats, decoder_feats, clean_feats)

# Inter-stage fusion at one encoder scale  [eq. 8]
# Input:  previous encoder and decoder features; current encoder features
# Output: modulated current encoder features

def fuse(enc_prev, dec_prev, enc_now):
    # Embed both previous-stage maps and add them  [eq. 8]
    fused = embed_enc(enc_prev) + embed_dec(dec_prev)
    # Spatial scale and shift, not per-channel vectors  [eq. 8]
    alpha = conv_alpha(fused)
    beta = conv_beta(fused)
    return enc_now * alpha + beta

The paper does not specify a learning-rate decay, and it does not state architectural widths; see Hyperparameters.

Hyperparameters

Values that the paper does not state are marked "not stated". Code values are from commit bae845c and apply to the Deraining folder unless noted.

Results

Author-reported. Setup: each task trained on the dataset named in §4.2-§4.5 with the training schedule in Hyperparameters, and evaluated with PSNR and SSIM §4.1. The compared methods are listed in the text of each subsection. The best and second-best marking of the source tables is not reproduced here; values are copied unchanged.

Image deraining tab. 1. Cells are PSNR / SSIM. Compared with §4.2.

Image deblurring tab. 2. Cells are PSNR / SSIM. "-" as in the source. Compared with §4.3.

Image denoising tab. 3. Cells are PSNR / SSIM. Compared with §4.4.

Compressive sensing tab. 4. Cells are PSNR / SSIM. Compared with §4.5.

Figures (described, not reproduced): fig. 4 shows deraining results on the test images, where the authors report that DGUNet and DGUNet+ remove faint rain lines better than MSPFN and MPRNet §4.2. fig. 5 shows deblurring results on two blurred test samples with per-image PSNR labels, comparing the blurred input, ground truth, Gao et al., DBGAN, MT-RNN, MPRNet, DGUNet and DGUNet+ §4.3. fig. 6 shows real denoising on DND and SIDD samples, with DGUNet and DGUNet+ at the end of the row §4.4. fig. 7 shows compressive sensing at a 10 % ratio on Set11, with the per-image PSNR labels 26.13 dB for DGUNet and 26.41 dB for DGUNet+ §4.5.

Headline comparisons stated in the text: DGUNet+ is 1.85 dB above MPRNet on Rain100L, 0.65 dB on Rain100H and 0.73 dB on the deraining average §4.2; 0.51 dB and 0.44 dB above MPRNet on GoPro and HIDE §4.3; 0.32 dB and 0.20 dB above MPRNet on DND and SIDD §4.4; about 2 dB above OPINENet+ on Set11 at 1 % §4.5. The authors note that DGUNet sometimes beats DGUNet+ in compressive sensing because the training set is small §4.5.

Ablation on Rain100H, PSNR, author-reported, all changes made on DGUNet+ §4.6.

Paper vs code

Repository: the code link in the frontmatter, at commit bae845c (2022-05-13). The paper links it in the abstract and it is written by the authors, so it is the official code. The repository has no LICENSE file, so no code is quoted; citations are by path and line only. I read the deraining model, training script and loss in full, and the deraining test script in part. For denoising, deblurring and compressive sensing I checked only the training configuration and loss.

• Loss. Paper: l2 summed over stages §3.3 eq. 10. Deraining: Charbonnier loss summed over the 7 stage outputs, plus 0.05 times an edge loss train.py losses.py losses.py. Deblurring: same train.py. Denoising: Charbonnier only, output clamped to [0, 1] train.py. Compressive sensing: mean squared error summed over stages by default train.py train.py. Only compressive sensing matches the paper's stated loss.
• Downsampling and upsampling. Paper: 2x2 max-pooling; bilinear upsampling plus convolution §3.2. Code: strided 4x4 convolution DGUNet.py DGUNet.py; transposed 2x2 convolution DGUNet.py. The bilinear helper classes are defined but not used DGUNet.py.
• Normalisation in the proximal mapping. Paper: residual blocks without batch normalisation §3.2. Code: encoder blocks are conv, LeakyReLU, conv with a 1x1 identity path and instance normalisation on half the channels (HIN); the decoder has no HIN DGUNet.py DGUNet.py DGUNet.py. The ResBlock without normalisation appears only in the gradient module DGUNet.py DGUNet.py.
• Inter-stage fusion. Paper: two 1x1 embeddings, summed, then alpha and beta applied as F = F_hat * alpha + beta §3.2 eq. 8. Code: 3x3 embeddings, LeakyReLU(0.1) on the sum, and output = out * sigmoid(phi(s)) + gamma(s) + out, so alpha is a sigmoid and the block's own output is added back DGUNet.py DGUNet.py. These are different functions.
• SAM. Paper: the SAM from MPRNet §3.2. Code: two convolutions with residual sums; this file has no sigmoid attention step DGUNet.py.
• Subspace projection. Paper: from NBNet, into the next stage §3.2. Code: a 16-dimensional subspace from a learned convolution, L1-normalised, then a 16x16 matrix inverse; applied at the start of each stage to the previous stage's SAM features DGUNet.py DGUNet.py DGUNet.py.
• Step size. Paper: rho^k trainable per stage §3.2. Code: one scalar per gradient module, initialised to 0.5, and shared by stages 2-6 DGUNet.py DGUNet.py.
• Stage sharing. Matches: shared stages 2-6 and independent stage 1 and 7 in the default model DGUNet.py; independent stages in DGUNet+ DGUNet_plus.py.
• Output order. Paper: the algorithm returns [x^1, ..., x^K] alg. 1. Code returns the list reversed, final stage first DGUNet.py; validation reads index 0 train.py.
• Compressive sensing operator. Paper: A exactly known, sampling matrix learned §4.5. Code: a learned matrix with 1024 columns, consistent with 32x32 blocks (inferred from the shape), and the exact gradient via the matrix's transpose DGUNet.py DGUNet.py. Consistent with eq. 6. The training patch size differs (see Hyperparameters).
• Test-time augmentation. Paper: not stated. Code: plain forward pass and clamp; the 8-way self-ensemble is defined but commented out test.py test.py.
• Training schedule. Paper: 4e5 iterations, 3 days on two V100s §4.1. Code runs in epochs with other totals (see Hyperparameters).

Pitfalls

• Algorithm 1 has its branches reversed. The text assigns eq. 6 to a known A and eq. 7 to an unknown A §3.2. The printed algorithm uses eq. 6 when A is unknown alg. 1. Follow the text.
• Index inconsistency in eq. 8. The text says the fused result of stage k is H^k_n §3.2, but the equation defines H^{k-1}_n from stage k-1 features. Implement the equation: the fusion uses the previous stage's features.
• The loss differs by task. Implementing the paper's l2 loss for deraining will not reproduce the code's training. The deraining code uses Charbonnier with eps 1e-3 (computed as sqrt(d^2 + 1e-6)) losses.py plus 0.05 times the edge term.
• Model selection uses a reported test set. The deraining validation folder is Rain100L training.yml, which is also one of the test sets in tab. 1. The checkpoint is chosen by PSNR on that folder train.py. The paper does not say which checkpoint it reports. Treat the Rain100L numbers as optimistic unless re-checked on a held-out set.
• "Twelve test sets" does not match. The introduction and conclusion say twelve §1 §5. Counting the named sets gives 11: five deraining, two deblurring, two denoising and two compressive sensing §4.2-§4.5.
• Learning rate and schedule are not fully specified. Warm-up length, decay and initial step-size values are not stated §4.1 §3.2. Code schedules differ between tasks, so copying one task's settings to another will not match any reported result.
• Feature widths are missing. The paper gives no channel counts, residual block width, attention reduction or subspace dimension. The code uses 40 base channels with +20 per encoder level, a 64-wide residual block, reduction 4 and a 16-dimensional subspace DGUNet.py DGUNet.py DGUNet.py. The num_cab argument is accepted but not used in the forward pass DGUNet.py DGUNet.py.
• Step-size sharing. The notation rho^k and the parameter set Omega suggest one step size per stage, but the code shares one scalar across stages 2-6 DGUNet.py DGUNet.py. With shared weights, sharing is faithful to the architecture, but the notation is ambiguous.
• Numerically sensitive step. The subspace projection normalises by an L1 norm with a 1e-6 offset and then inverts a 16x16 matrix with torch.inverse DGUNet.py. Near-singular cases are not handled explicitly in the code.
• Hidden cost. Each of the 7 stages contains a full encoder-decoder across three scales plus two residual blocks for the gradient step, so the network is heavy. The paper reports complexity only in a figure, with no numbers in the text fig. 1. Budget memory and time from the code, not from the paper.
• Supplementary material is referenced but not in the source. The LaTeX mentions supplementary results in commented-out text §4 and does not include them.
• Baseline numbers may not share one protocol. The paper does not say whether the baseline rows were rerun under the authors' setup or copied from the cited papers §4.2 §4.3. Compare with them only as reported.

Lineage

• Builds on: Proximal Gradient Descent and ISTA, as the algorithmic basis §3.1. The deep unfolding networks reviewed in §2 (deep plug-and-play and ADMM/HQS unfolding) are the closest prior work the paper discusses. The source gives no arXiv ids for them, so no link is given.
• Uses from other work: the supervised attention module from MPRNet §3.2, subspace projection from NBNet §3.2, and spatially-adaptive normalisation §3.2. The source gives no arXiv ids for these.
• Compared against: deraining: DerainNet, SEMI, DIDMDN, UMRL, RESCAN, PreNet, MSPFN, MPRNet §4.2. Deblurring: DBGAN, MT-RNN, DMPHN, Suin et al., MPRNet and others in tab. 2 §4.3. Denoising: SADNet, DANet+, CycleISP, DeamNet, MPRNet §4.4. Compressive sensing: ISTANet+, CSNet, AdapRecon, OPINENet+, AMPNet §4.5.
• Superseded by: not stated in the source.