Efficient and Explicit Modelling of Image Hierarchies for Image Restoration

In plain words

Image restoration means turning a damaged picture back into a clean one. The paper covers motion and defocus deblurring, denoising, super-resolution, JPEG artifact removal and demosaicking, plus real-world super-resolution.

Good restoration has to use information from far across the picture, but comparing every pixel with every other pixel gets too expensive on large images. The network, GRL, handles three ranges in every layer, separately and in parallel. Fine detail goes through a convolution with a channel-weighting step. Nearby areas go through attention inside small square windows. Large-scale context goes through attention along thin horizontal and vertical stripes. Each stripe sends its information through a small set of summary points, called anchors, so the attention stays cheap. The stripe direction alternates, so information can spread across the whole picture.

The authors report new best scores on several of the seven tasks, with three model sizes (tiny, small and base).

TL;DR

GRL is a transformer whose layer has three parallel parts: window self-attention for the regional range, anchored stripe self-attention for long-range dependencies, and a channel-attention enhanced convolution for local structure §4. The stripe attention replaces the full attention map with two smaller maps through a set of anchors eq. 2 eq. 3 eq. 4. Author-reported: GRL-B sets a new best PSNR on GoPro and HIDE deblurring tab. 2, on RealBlur-R tab. 3 and on several denoising, super-resolution, JPEG and demosaicking settings tab. 4 tab. 5 tab. 6 tab. 7 tab. 8. The official code differs from the paper in several places: the attention score, the optimizer and loss, the layer structure of the tiny and small models, and the anchor pooling factor. Read Paper vs code before implementing.

Problem

Self-attention models long-range dependencies explicitly, but its cost grows with the square of the number of tokens, and for images the token count is height times width §3.1. Window attention is cheap but limited to a small regional range; the authors cite evidence that a larger range helps restoration §2 §3.1. The task is to model dependencies beyond the regional range while keeping the cost under control, for high-resolution inputs §1 §3.1.

Method

The network has three parts §4. (1) A 3x3 convolution turns the input into feature maps. (2) A representation-learning module stacks transformer stages, each a stack of transformer layers that ends in a convolution, with a skip connection around the stage and around the whole module §4. (3) A reconstruction module produces the restored image §4. Feature dimension is kept constant across the representation module §4.

Anchored self-attention. Standard attention is written as

\mathbf{Y} = \mathrm{Softmax}\left( \mathbf{Q} \cdot \mathbf{K}^T / \sqrt{d} \right) \cdot \mathbf{V} eq. 1

with Q, K and V linear projections of the input X, N tokens of dimension d §3.1. The anchored version adds M anchors A (M much smaller than N) and computes

\mathbf{Y} = \mathbf{M}_e \cdot \left( \mathbf{M}_d \cdot \mathbf{V} \right) eq. 2

\mathbf{M}_d = \mathrm{Softmax}( \mathbf{A} \cdot \mathbf{K}^T / \sqrt{d}) eq. 3

\mathbf{M}_e = \mathrm{Softmax}( \mathbf{Q} \cdot \mathbf{A}^T / \sqrt{d}) eq. 4

Anchors act as a bottleneck. The product is evaluated right to left: first the anchors collect the values, then each query reads from the anchors §3.2. The stated cost is O(NMd) time and O(NM) space §3.2.

Anchored stripe self-attention. The attention is applied inside vertical and horizontal stripes, not over the whole image §3.3. The four modes are horizontal, vertical, and shifted horizontal and vertical stripes. Horizontal and vertical stripes alternate across the transformer §3.3. The anchor projection is an average pooling layer followed by a linear projection, which down-scales the feature map by a factor s in both directions §4. The factor s is not given in the paper text.

Transformer layer. Each layer runs in parallel §4 fig. 6:

• a shared linear projection to Q, K and V;
• the channels are split equally: one half goes to window attention (Swin V2 style) and the other half to anchored stripe attention; the two outputs are concatenated §4;
• a convolution branch with channel attention, which captures local structure §4;
• the outputs of the attention module and the convolution branch are added to the layer input, and the sum goes through an MLP §4.

# Anchored stripe self-attention for one stripe  [§3.2] [eq. 2-4]
# Input:  tokens (N, d) of one stripe, anchors (M, d)
# Output: (N, d)

def anchored_stripe_attention(tokens, anchors):
    # Q, K, V are plain linear projections of the tokens  [§4]
    queries = linear_query(tokens)               # (N, d)
    keys = linear_key(tokens)                    # (N, d)
    values = linear_value(tokens)                # (N, d)
    # Anchors read from the keys  [eq. 3]
    anchor_to_key = softmax(anchors @ keys.T / sqrt(d))   # (M, N)
    # Anchors collect the values  [eq. 2]
    anchor_values = anchor_to_key @ values       # (M, d)
    # Queries read from the anchors  [eq. 4]
    query_to_anchor = softmax(queries @ anchors.T / sqrt(d))  # (N, M)
    # Expand the anchor summary back to the tokens  [eq. 2]
    return query_to_anchor @ anchor_values       # (N, d)

# One GRL transformer layer  [§4] [fig. 6b]
# Input: x (B, L, C) with L = H * W

def grl_layer(x):
    # Shared Q, K, V projection  [§4]
    qkv = linear_qkv(x)                          # (B, L, 3C)
    # Split the channels equally between the two attention modules  [§4]
    window_part, stripe_part = split_channels(qkv)
    # Regional range: window attention  [§4]
    regional = window_attention(window_part)     # (B, L, C/2)
    # Beyond the regional range: anchored stripe attention  [§4]
    anchors = avgpool_then_linear(x)             # anchors from the input
    global_part = anchored_stripe(stripe_part, anchors)   # (B, L, C/2)
    # Concatenate the two halves and project  [§4]
    attended = output_projection(concat(regional, global_part))
    # Local structure: channel-attention enhanced convolution  [§4]
    local = channel_attention_conv(x)            # (B, L, C)
    # Add both branches to the input, then the MLP  [§4]
    x = x + attended + local
    x = mlp(x)
    return x

Notes on the pseudocode: the paper does not describe the output projection, the normalisation layers or the MLP residual, and the code adds them (see Paper vs code and Pitfalls).

Hyperparameters

As stated in the paper. Main text only; the supplementary material is not in the arXiv source.

As set in the official code. These are the released training configs, not the paper. Each config trains one task. Each uses 8 GPUs with batch size 1 per GPU, except demosaicking (batch 4) grl_p256.yaml.

Model shapes in the code grl_tiny.yaml grl_small.yaml grl_base.yaml:

The paper does not state these values.

Results

Author-reported. Setup: the paper's own tables, each task against the baselines in its table. The paper does not give the training protocol in the main text. Table numbers below count captions in source order: Tab. 1 defocus, Tab. 2 GoPro and HIDE, Tab. 3 RealBlur, Tab. 4 denoising, Tab. 5 classical super-resolution, Tab. 6 grayscale JPEG, Tab. 7 colour JPEG, Tab. 8 demosaicking, Tab. 9 similarity ablation, Tab. 10 anchor ablation. Figure numbers are Fig. 1 motivation, Fig. 2 teaser, Fig. 3 cross-scale similarity, Fig. 4 stripe features, Fig. 5 attention map visualisation, Fig. 6 network, Fig. 7 real-world SR.

Motion deblurring, PSNR / SSIM tab. 2

The text gives 1.01 dB over Restormer on GoPro and 0.43 dB on HIDE §5.1.

Motion deblurring, RealBlur tab. 3

Defocus deblurring, combined scenes, PSNR / SSIM / MAE / LPIPS tab. 1

The text gives 0.20 dB and 0.38 dB over Restormer, and 1.39 dB and 1.05 dB over Uformer and IFAN for dual-pixel §5.1.

Colour denoising, PSNR at sigma 15 / 25 / 50 tab. 4

Grayscale denoising (Set12, BSD68, Urban100) and the other baselines are in tab. 4. The text says GRL-B cuts parameters by 24% against Restormer §5.2; the table gives 19.81 M against 26.13 M.

Classical super-resolution, PSNR at x4 tab. 5

At x2, GRL-B reaches 38.67 on Set5 and 35.06 on Urban100 tab. 5. SSIM is in the same table. The text gives 0.44 dB and 0.22 dB over the 0.92 M EDT on Urban100 and Manga109 at x4 §5.2.

JPEG artifact removal, PSNR tab. 6 tab. 7

The caption gives 71.92 M parameters for FBCNN and 3.12 M for GRL-S tab. 6.

Demosaicking, PSNR tab. 8

Real-world super-resolution. Visual comparison only, no metric §5.3 fig. 7.

Ablations. Dot product beats Euclidean distance in all 18 cells of tab. 9 (two test sets, colour and grayscale denoising at three sigmas, and x2, x3, x4 super-resolution). Among anchor projections tab. 10, all are within 0.05 dB on Set5 at about 35.0 dB, and average pooling with a linear layer (3.12 M) is kept.

Attention similarity. The authors report a Pearson correlation of 0.9505 between the approximate and exact attention maps §3.4 fig. 5.

Paper vs code

Official repository, owner Yawei Li (first author), at commit 3123fcadfee2107d813c5e8fff58227f14c989bc. Licence MIT LICENSE. The repository README calls it the official implementation.

1. Attention score. Paper: scaled dot product, Q K^T / sqrt(d) eq. 1 eq. 3 eq. 4. Code: cosine similarity, because queries and keys are L2-normalised, multiplied by a learned per-head logit scale that starts at 10 and is clamped, and a continuous relative position bias passed through 16 x sigmoid mixed_attn_block_efficient.py mixed_attn_block_efficient.py mixed_attn_block_efficient.py mixed_attn_block_efficient.py. There is no 1/sqrt(d) factor. The paper mentions none of these.
2. Euclidean option. The paper's ablation compares dot product with Euclidean distance §3.4 tab. 9. In the code, the Euclidean branch uses torch.norm(q - k) as the logit, with no minus sign, so larger distances get larger weights mixed_attn_block_efficient.py. The released default is off grl_base.yaml. Check the sign before using this option to reproduce tab. 9.
3. Anchors are per stripe. The paper defines one anchor set A eq. 2 §4. The code partitions the anchor map with the same stripe grid as the queries, so each stripe gets its own anchors mixed_attn_block_efficient.py.
4. Optimizer and loss. Paper: Adam, L1, initial learning rate 2e-4, for real and synthetic restoration §5. Code: AdamW for all tasks except real-world SR, which uses Adam. The loss is L1 for super-resolution and denoising, and Charbonnier for defocus, motion, JPEG and demosaicking grl_p256.yaml grl_p480.yaml grl.yaml. Learning rates are per task (table above). The Charbonnier loss uses eps = 1e-3 losses.py.
5. Real-world SR uses a GAN. The paper describes no adversarial training for real-world SR §5.3. The released config uses a GAN loss and a discriminator, starting from a PSNR-trained checkpoint grl.yaml grl.yaml grl_base_bsr.yaml.
6. Tiny and small models have no conv branch. The paper describes the convolution branch as part of the transformer layer §4. The released tiny and small configs turn it off grl_tiny.yaml grl_small.yaml. The base config turns it on grl_base.yaml.
7. Anchor pooling factor. The paper leaves s unstated §4. The model configs use 4 grl_base.yaml. The denoising test script uses 2 for GRL-B, together with window size 32 grl_test.md.
8. Stripe sizes. The paper does not give stripe sizes. The base model config can derive stripe sizes from the input resolution, using stripe groups [null, 4] grl_base.yaml mixed_attn_block_efficient.py. The task configs fix the stripe sizes instead.
9. Stripe direction and shift. The paper says horizontal and vertical stripes alternate, with shifted modes §3.3. The code agrees. Layer i is horizontal when i is even, and stripes are shifted when i mod 4 is 2 or 3 grl.py. Window attention is shifted on even layers grl.py.
10. Channel attention. The paper says only "channel attention enhanced convolution" §4. The code uses two 3x3 convolutions with a 4x channel squeeze and GELU, followed by a channel-attention gate with reduction 18 mixed_attn_block.py.
11. Anchor projection. Matches: average pooling, then a linear layer to half the width §4 mixed_attn_block.py.
12. Output of the stripe branch. The paper does not describe an output projection. The code applies a linear layer to the concatenated window and stripe outputs mixed_attn_block_efficient.py.
13. Text against table. The text says dot product "outperforms a distance metric for a couple of settings" §5.4. Table 9 shows dot product ahead in all 18 cells.
14. JPEG claim. The text says GRL-S beats the previous state of the art "across different datasets and quality factors" §5.2. Table 6 has the BSD500 QF 30 SSIM at 0.912 for GRL-S and 0.913 for FBCNN.
15. Parameter counts differ between tasks. GRL-S is 3.12 M in denoising and JPEG tables but 3.34 M and 3.49 M at x2 and x4 tab. 4 tab. 5 tab. 6. The upsampler differs by task, so one model name covers several sizes.

Pitfalls

• Supplementary material is missing. The arXiv source has no appendix. The paper says the training datasets, protocols and extra visual results are in the supplementary material §5. Patch size, batch size, steps and data are not stated in the paper. The values above come from the official configs at the commit in the front matter, not from the paper.
• Score function. To match the released weights, implement the cosine score with learned scale and position bias mixed_attn_block_efficient.py. To follow the paper's equations, use sqrt(d) scaling. The two differ (Paper vs code, item 1).
• Multi-head attention is not in the paper. The equations are single-head. The code uses 2 to 3 heads, split across the window and stripe halves grl_base.yaml. Head dimension is not stated.
• Normalisation layers and residuals. The paper does not describe them. The code applies a LayerNorm to the attention output before the residual add mixed_attn_block_efficient.py, applies LayerNorm around the MLP mixed_attn_block_efficient.py, and adds a residual scale of 1.0 grl.py. The layer also has a stochastic depth rate of 0.1 that rises across depth by default grl.py grl.py.
• Stem and tail. The code adds a LayerNorm at the start and end of the body grl.py, and a global residual from the first convolution to the end of the body grl.py. The paper states the skip connections §4 but not these norms.
• Training configs default to the tiny model for SR and denoising. The super-resolution and denoising configs set grl/grl_tiny grl_p256.yaml and grl_p256.yaml. Override model= to train the other sizes.
• Window size changes per task and at test time. Window size is 8 to 36 across task configs (see the training table). The denoising test script changes the window size to 32 and the anchor factor to 2 for GRL-B grl_test.md.
• Tiling at inference. The configs use tiles and overlaps, for example 480 with 48 overlap for defocus, and 576 with 72 for JPEG grl_p480.yaml grl_p288.yaml. The paper does not state an inference protocol.
• Cost. The anchor step reduces the attention cost, but the position-bias MLP and per-stripe mask and index tables are extra work and memory that the paper does not discuss.
• Table numbers. Numbers in this card follow caption order in the LaTeX source. Check against the compiled PDF before you cite them.
• Author-reported numbers. All results are from the authors. The base model's large gains are shown on a few benchmarks, and the abstract claims state of the art for "several" of the seven tasks.

Lineage

• Builds on: Swin Transformer V2 window attention §4 (arXiv id not in this card). The anchors and stripe attention are the paper's own contribution §3.
• Compared against: Restormer, SwinIR, EDT, DRUNet, IPT, MPRNet, Stripformer, MAXIM, FBCNN, RNAN and others in the tables tab. 1-8. None of these have cards in this wiki yet.
• Superseded by: not known from this card.