Transformer for Single Image Super-Resolution

In plain words

Super-resolution means taking a small, blurry image and producing a larger, sharper one. The network has to guess fine detail that is not in the input. A useful trick is that images repeat themselves: bricks in a wall, windows on a building or letters in a comic look alike, so a clear copy of a pattern in one place can help reconstruct a blurred copy elsewhere. Transformers, with their attention mechanism, are good at finding such look-alikes, but on full images they need far too much GPU memory.

ESRT splits the work in two. A small convolutional network first does most of the work cheaply: it shrinks the feature maps to half size to save compute, but before shrinking it subtracts a blurred copy from the original and keeps the leftover "edges and texture" signal, which it adds back after upsampling. Then a single, slimmed-down Transformer lets every position look at similar positions, but only within one quarter of the image at a time and with half the usual feature width. That keeps memory to about a quarter of a standard Transformer layer.

The result is a model with well under a million parameters that matches or beats other small super-resolution networks on standard test sets, with the biggest gains on Urban100, a dataset full of repeating building facades.

TL;DR

ESRT is a hybrid CNN + Transformer for lightweight single-image super-resolution (SISR) §3. The CNN part (LCB) stacks 3 High Preserving Blocks that process features at half resolution and carry a high-frequency residual (feature minus its blurred, upsampled copy) around the downsampling §3.1 §3.2. The Transformer part (LTB) is one Efficient Transformer whose tokens are unfolded 3×3 neighbourhoods (one per pixel); attention halves the channel width and splits the token sequence into s=4 segments that attend only within themselves §3.3. Author-reported ×4: 751K parameters, Set5 32.19 dB, Urban100 26.39 dB; training memory 4,191M vs 16,057M with an original Transformer tab. 1 tab. 4. The paper has several internal inconsistencies (parameter counts, a probably mislabeled SwinIR table, no ×2 numbers), and the code adds a channel-attention layer inside every ARFB and uses 3×3 instead of 1×1 convs in the residual units.

Problem

CNN-based SISR models reach good quality with many layers but are expensive; lightweight CNNs cut cost at the price of capacity, and pure CNNs only see local context §1 §2.1. Image patches often recur within an image and could serve as references for each other, which suggests self-attention §1 fig. 1. But vision Transformers (ViT, Swin, IPT, SwinIR) need large GPU memory and, in the authors' view, lots of training data; SR inputs are large images, so the token count N is large and the N×N attention matrix is the bottleneck §1 §2.2 §3.3.2. Assumption: a predicted SR pixel depends mostly on a local neighbourhood of the LR image, so attention can be restricted to a local segment §1 §3.3.2.

Method

Four stages: shallow 3×3 conv, Lightweight CNN Backbone (LCB, a chain of HPBs), Lightweight Transformer Backbone (LTB, ETs over the concatenated HPB outputs), reconstruction with two PixelShuffle branches §3 fig. 2.

Paper equations eq. 1–4:

F_0 = f_s(I_{LR}) eq. 1

F_n = \zeta^{n}(\zeta^{n-1}(\dots(\zeta^{1}(F_0)))) eq. 2

F_d = \phi^{n}(\phi^{n-1}(\dots(\phi^{1}([F_1, F_2, \dots, F_n])))) eq. 3

I_{SR} = f(f_p(f(F_d))) + f(f_p(F_0)) eq. 4

with \zeta^n the n-th HPB, \phi an ET, [\cdot] concatenation, f a conv and f_p PixelShuffle.

1. High Preserving Block (HPB) §3.1 fig. 3 eq. 5. An ARFB processes F_{n-1}; HFM extracts its high-frequency part P_{high}; the ARFB output is downsampled (×1/2) and passed 5 times through one weight-shared ARFB; P_{high} goes through its own ARFB; the low branch is bilinearly upsampled and concatenated with the high branch; then a 1×1 conv halves channels, channel attention (SE-style) reweights channels, a final ARFB runs, and F_{n-1} is added back (global residual) §3.1.

   F_{n-1}^{''} = [f_{a}(P_{high}),\ \uparrow f_{a}^{\circlearrowright^5}(\downarrow F_{n-1}^{'})] eq. 5

2. High-frequency Filtering Module (HFM) §3.2 fig. 4 eq. 6. T_A = \mathrm{avgpool}(T_L, k) with k=2, upsample T_A back to C\times H\times W to get T_U, and output T_L - T_U §3.2 §4.2. The paper presents it as a differentiable stand-in for a Fourier high-pass filter; it has no learnable weights.

3. Adaptive Residual Feature Block (ARFB) §3.2.1 fig. 6 eq. 7. Two residual units (RUs) in sequence, then concat of both RU outputs → 1×1 conv → 3×3 conv. Each RU is a channel Reduction (to C/2) and Expansion (back to C) with learnable scalars on both paths:

   y_{ru} = \lambda_{res}\cdot f_{ex}(f_{re}(x_{ru})) + \lambda_{x}\cdot x eq. 7

   Both λ initialised to 1 §4.2. (The paper writes x for x_{ru} on the right-hand side.)

4. Pre/post-processing for ET §3.3.1 fig. 7. Unfold the C\times H\times W map with a k×k kernel (k=3) into N = H·W tokens of size k^2 C; no position embedding; after the ET, Fold back to C\times H\times W §3.3.1 §4.2.

5. Efficient Transformer (ET) §3.3.2 fig. 8 eq. 8. Pre-norm encoder block:

   E_{m1} = \mathrm{EMHA}(\mathrm{Norm}(E_i)) + E_i,\quad E_o = \mathrm{MLP}(\mathrm{Norm}(E_{m1})) + E_{m1} eq. 8

6. Efficient Multi-Head Attention (EMHA) §3.3.2 fig. 8. Reduction linear layer C → C/2; one linear layer to Q, K, V; reshape to B\times m\times N\times \frac{C_1}{m} with m=8 heads; Feature Split cuts Q, K, V along the token axis into s=4 equal segments; scaled dot-product attention (no mask) within each segment, so the attention matrix is \frac{N}{s}\times\frac{N}{s}; concat segment outputs; Expansion linear layer back to C §3.3.2 §4.2.

Forward pass as implemented in the official code (which differs from the paper in places; see Paper vs code):

# ESRT forward  [§3] [eq. 1-4] [code: model/esrt.py#L204-L221]
# Input:  lr_image (B, 3, h, w) in [0, 1]; scale in {2, 3, 4}
# Output: sr_image (B, 3, scale*h, scale*w)

def esrt(lr_image):
    # shallow 3x3 conv, 3 -> 32 channels  [eq. 1] [code: model/esrt.py#L181]
    shallow = conv3x3(lr_image, out=32)
    # LCB + LTB with outer learnable-scaled residual  [code: model/esrt.py#L149-L161]
    body = backbone(shallow)
    # extra 3x3 conv (n_blocks=1, so 32 -> 32)  [code: model/esrt.py#L201]
    body = conv3x3(body, out=32)
    # deep branch: conv + PixelShuffle, then conv to RGB  [eq. 4]
    # [code: model/esrt.py#L190-L193] [code: model/common.py#L87-L114]
    deep = conv3x3(pixel_shuffle_upsampler(body, scale), out=3)
    # shallow branch: same upsampler, then conv WITHOUT bias, followed by ReLU
    # [code: model/esrt.py#L196-L197] [code: model/esrt.py#L61-L87]
    skip = relu(conv3x3_nobias(pixel_shuffle_upsampler(shallow, scale), out=3))
    return deep + skip

# Lightweight CNN + Transformer backbone  [§3.1] [§3.3]
# [code: model/esrt.py#L137-L161]
def backbone(features):                         # (B, 32, h, w)
    out1 = hpb_1(features)
    out2 = hpb_2(out1)
    out3 = hpb_3(out2)
    # concat HPB outputs, 3x3 conv 96 -> 32 (conv not in the paper)
    fused = conv3x3(concat([out1, out2, out3]), out=32)
    # one ET over unfolded 3x3 patches  [§3.3.1] [code: util/transformer.py#L150-L165]
    tokens = unfold(fused, kernel=3, stride=1, zero_pad="same")  # (B, 288, h*w)
    tokens = efficient_transformer(transpose(tokens))            # (B, h*w, 288)
    # Fold sums the 9 overlapping contributions per pixel (no averaging)
    # [code: util/tools.py#L59-L76]
    folded = fold(transpose(tokens), out_size=(h, w), kernel=3, padding=1)
    folded = conv3x3(folded, out=32)
    return scale_a * features + scale_b * folded  # learnable scalars, init 1

# High Preserving Block  [§3.1] [eq. 5] [code: model/esrt.py#L112-L135]
def hpb(x):                                     # (B, 32, h, w)
    x1 = arfb_in(x)
    # HFM with k=2  [§3.2] [eq. 6]
    low = avg_pool(x1, kernel=2)                # (B, 32, h/2, w/2)
    high = x1 - bilinear_resize(low, (h, w), align_corners=True)
    # five passes through ONE shared ARFB at half resolution  [eq. 5]
    for repeat in range(5):
        low = arfb_low(low)
    high = arfb_high(high)
    up = bilinear_resize(low, (h, w), align_corners=True)
    # 1x1 fusion conv 64 -> 32 (with ReLU in code), then channel attention
    fused = channel_attention(relu(conv1x1_nobias(concat([up, high]))))
    return arfb_out(fused) + x

# Adaptive Residual Feature Block  [§3.2.1] [eq. 7]
# [code: model/esrt.py#L89-L110]
def arfb(x):
    r1 = residual_unit(x)
    r2 = residual_unit_2(r1)
    # concat of both RUs (each with a learnable scalar), 1x1 conv 64 -> 32
    merged = relu(conv1x1_nobias(concat([w_a * r2, w_b * r1])))
    # code inserts channel attention here; not in the paper
    merged = channel_attention(merged)
    merged = relu(conv3x3_nobias(merged))
    return w_c * x + w_d * merged

# Residual unit with adaptive weights (RSA)  [eq. 7]
# [code: model/esrt.py#L43-L59]
def residual_unit(x):                           # (B, 32, h, w)
    # code uses 3x3 convs, paper says 1x1
    reduced = prelu(conv3x3(x, out=16))
    expanded = conv3x3(reduced, out=32)
    return lambda_x * x + lambda_res * expanded

# Efficient Transformer, pre-norm  [§3.3.2] [eq. 8]
# [code: util/transformer.py#L138-L165]
def efficient_transformer(tokens):              # (B, N, 288)
    tokens = tokens + emha(layer_norm_1(tokens))
    tokens = tokens + mlp(layer_norm_2(tokens))
    return tokens

# MLP: 288 -> 72 -> 288, ReLU, no dropout  [code: util/transformer.py#L60-L76]
def mlp(tokens):
    return linear_2(relu(linear_1(tokens)))

# Efficient multi-head attention  [§3.3.2] [code: util/transformer.py#L81-L132]
def emha(tokens):                               # (B, N, 288)
    # Reduction: 288 -> 144, no bias
    reduced = linear_reduce(tokens)
    # Q, K, V from one linear layer, no bias; 8 heads of 18 dims
    queries, keys, values = split_heads(linear_qkv(reduced), heads=8)
    # Feature Split: contiguous chunks of N // 4 tokens in raster order
    segment = N // 4
    outputs = []
    for q_seg, k_seg, v_seg in zip(
            chunk(queries, segment), chunk(keys, segment),
            chunk(values, segment)):
        # scale uses 288 / 8 = 36, i.e. 1/6, not 1/sqrt(18)
        weights = softmax(q_seg @ transpose(k_seg) * 36 ** -0.5)
        outputs.append(weights @ v_seg)
    merged = merge_heads(concat(outputs, axis=tokens))  # (B, N, 144)
    # Expansion: 144 -> 288, with bias
    return linear_expand(merged)

Hyperparameters

Results

Author-reported. Trained on DIV2K only, bicubic degradation (implied by the code's DIV2K_train_LR_bicubic path DIV2K.py); evaluated on Set5, Set14, BSD100, Urban100 and Manga109 §4.1.

Main comparison, PSNR / SSIM tab. 1:

Cost comparison, input 1280×720 tab. 2:

HPB ablation on Set5 ×4 tab. 3 (RB = plain residual block in place of ARFB):

ET ablation on Set5 ×4; GPU memory measured during training at 48×48 patches, batch 16 tab. 4:

ET added to a half-depth RCAN (5 residual groups), ET placed before reconstruction §4.5.2 tab. 5:

RealSR tab. 6:

Versus SwinIR, caption says ×4 tab. 7; SwinIR was trained with extra Flickr2K data, ESRT was not §4.7:

Headline: at ×4, ESRT has the best Set14, BSD100, Urban100 and Manga109 scores among the listed lightweight models, with the largest margin on Urban100 (26.39 vs 26.25 dB for LatticeNet), which the authors attribute to repeated patches in that dataset tab. 1 §4.3. It is behind LatticeNet and RFDN-L on Set5 tab. 1. It is about 4× slower than CARN/IMDN at inference tab. 2.

Paper vs code

Checked against model/esrt.py, util/transformer.py, util/tools.py and train.py at the recorded commit.

• RU kernel size. Paper: Reduction/Expansion are 1×1 convs §3.2.1 §4.2. Code: both are 3×3 convs with a PReLU between them esrt.py. This changes parameter count and receptive field substantially.
• Extra channel attention inside every ARFB. Paper's ARFB is two RUs → concat → 1×1 → 3×3 §3.2.1 fig. 6. Code inserts a CA layer after the 1×1 conv, adds ReLU after the 1×1 and 3×3 convs (BasicConv defaults to ReLU, no bias), puts learnable scalars on both RU outputs before the concat, and adds an outer learnable-weighted residual around the whole ARFB esrt.py esrt.py.
• HPB fusion conv is followed by ReLU in code (BasicConv) esrt.py; the paper does not mention an activation §3.1.
• Extra convs in the backbone. Code applies a 3×3 conv to the concatenated HPB outputs before the ET (96 → 32 channels), a 3×3 conv after Fold, wraps LCB+LTB in a residual with learnable scalars, and applies another 3×3 conv before reconstruction esrt.py esrt.py esrt.py. Eq. 3 shows the concat going straight into the ET eq. 3.
• Reconstruction. Eq. 4 puts a conv before PixelShuffle only on the deep branch; in code both branches are Upsampler (conv to 4C or 9C channels, PixelShuffle) followed by a 3×3 conv to RGB; the shallow branch's final conv has no bias and is followed by ReLU esrt.py common.py. No mean shift is applied (commented out) esrt.py.
• Attention scale. Code computes the scale from the un-reduced dim: (288/8)^-0.5 = 1/6, although heads have 144/8 = 18 dims transformer.py. The paper does not state the scale.
• Feature Split is torch.split(q, N//4) on the raster-ordered token axis transformer.py: segments are horizontal bands of the feature map, and if N is not divisible by 4 a fifth, smaller segment appears. The paper says "s equal segments" §3.3.2.
• MLP is a bottleneck (288 → 72 → 288, ReLU) although mlp_ratio=4. is passed and ignored transformer.py transformer.py; the paper gives no MLP size.
• Augmentation includes vertical flips in code common.py; the paper lists horizontal flips and 90° rotation §4.2.
• Inference splits the LR input into 4 overlapping quadrants (10-pixel overlap), recursively, until each piece is under 60,000 pixels test.py test.py; not mentioned in the paper. forward_chop in test.py hard-codes scale = 4 test.py, so ×2 and ×3 evaluation with that script needs a fix.

Pitfalls

• Inconsistent parameter counts. ESRT is 770K (×3) and 751K (×4) in tab. 1, 0.68M in tab. 2, 751K in tab. 3 and tab. 4, and 770K in the "×4" SwinIR table tab. 7. Do not use any single number as a target without counting the code's model.
• The SwinIR table is probably ×3, not ×4. ESRT's BSD100 and Manga109 values in tab. 7 (29.15 / 0.8063, 33.95 / 0.9455) and its 770K parameters are exactly the ×3 row of tab. 1, while the caption says ×4. The 4,191M memory figure comes from the ×4 ablation tab. 4.
• No ×2 numbers. The text discusses ×2 visual results and a ×2 trade-off plot §4.3 fig. 10, and the code trains ×2 by default train.py, but tab. 1 only has ×3 and ×4 rows.
• Small cross-table mismatches. Set5 ×4 is 32.19 dB in tab. 1 and tab. 3 but 32.18 dB for "1 ET" in tab. 4. MAFFSRN ×4 Set14 is printed as 26.62 in tab. 1, far below every other ×4 model; likely a transcription error, but that is our inference.
• "Original Transformer" baseline is underspecified. tab. 4 cites a plain Transformer (16,057M memory) but the paper does not say how its tokens, width or depth were configured, so the 4× memory saving is hard to reproduce exactly.
• Attention range depends on input size. Because the split is by contiguous quarters of the token sequence, each token attends to a band of about H/4 rows of the current input. During training (48×48 LR) that is 12 rows; at test time with chopping it is a quarter of each chopped piece. With N not divisible by 4 a fifth, smaller segment appears. Results may change if you remove chopping or use a different tile size.
• Fold sums, it does not average. The 3×3 Fold adds up to 9 overlapping contributions per pixel (fewer at borders) tools.py. A reimplementation that averages will change scale at borders vs interior, and pretrained weights will not match.
• Memory is still O((N/4)²). Splitting into 4 segments only cuts the attention matrix 16×; for a full 1280×720 LR input (921,600 tokens) one segment's attention matrix would be about 5.3×10¹⁰ entries per head, which is why inference must be tiled. The paper's "low memory" figure is measured on 48×48 training patches tab. 4.
• Runtime is not low. ESRT has 163 layers and is about 4× slower than IMDN/CARN on a 1280×720 input despite comparable FLOPs tab. 2; the five shared-weight ARFB passes per HPB add depth without adding parameters.
• "Momentum 0.9" for Adam §4.2 maps to β1 = 0.9; β2 is not stated (code uses the PyTorch default) train.py.
• HFM upsampling method is not stated in the paper (only "upsampled") §3.2; code uses bilinear with align_corners=True esrt.py.
• Code repository hygiene. train.py imports model.architecture, which is not in the repository's file list, and esrt.py imports thop at module level train.py esrt.py; expect to patch imports before training.

Lineage

• Builds on: ViT (arXiv:2010.11929) for the encoder-only pre-norm block and patch tokens, replacing non-overlapping patches with overlapping unfolded 3×3 neighbourhoods §3.3.1 §3.3.2; IMDN (arXiv:1909.11856), whose png-to-npy data preparation the README follows; SE channel attention (arXiv:1709.01507) inside HPB §3.1.
• Compared against: VDSR, MemNet, EDSR-baseline, SRMDNF, CARN, IMDN, RFDN-L, MAFFSRN, LatticeNet tab. 1; RCAN tab. 5; LK-KPN on RealSR tab. 6; SwinIR (arXiv:2108.10257) tab. 7.
• Superseded by: not known from content in this wiki.