MB-TaylorFormer: Multi-branch Efficient Transformer Expanded by Taylor Formula for Image Dehazing

In plain words

Dehazing removes the haze from a photo so the scene looks clear. Transformers handle this well because every pixel can look at the whole image, but the standard attention step compares every pixel with every other one, so its cost grows very fast as the image gets bigger. This paper rewrites that step so the work grows in a straight line with the number of pixels, and adds a small learned correction for the error the rewrite introduces.

The first stage of the network also looks at the image through several paths at once. Each path uses a different window size and a different, flexible sampling shape, so the model gets both coarse and fine views of the same region. The authors report that this design dehazes several benchmarks with strong scores while staying light.

TL;DR

MB-TaylorFormer is a U-Net-shaped Transformer for single-image dehazing. Its attention uses a first-order Taylor expansion of softmax, so cost is linear in the number of pixels, and a per-head gate corrects the error of that expansion. Patch embedding stacks depthwise deformable convolutions in several branches. Author-reported: 40.71 dB PSNR on SOTS-Indoor for the base model and 42.64 dB for the large model tab. 1. The paper alone leaves several details open; the official code fills some of them and adds a learned per-head scale that the equations omit.

Problem

• Softmax attention costs time quadratic in the number of positions, which limits its use on high-resolution dehazing inputs §1 eq. 5 §3.3.
• Window-restricted attention lowers cost but shortens the receptive field §1.
• Fixed-kernel patch embedding produces tokens of one scale only §1 §3.2.
• Earlier linear-attention methods approximate the softmax kernel with feature maps. This paper instead expands softmax itself, to keep the value approximation close to vanilla attention §2 §1.

Method

The network takes a hazy image and predicts a residual that is added back to the input §3.1. Its parts, in order:

1. Backbone §3.1 fig. 2. A 3×3 convolution extracts shallow features. A four-level encoder–decoder follows. Each level applies multi-scale patch embedding, then several Transformer branches, then SKFF fusion of the branches §3.1. Downsampling uses pixel-unshuffle and upsampling uses pixel-shuffle §3.1. Skip connections are concatenated and reduced with a 1×1 convolution, except at the first level §3.1. A refinement stage follows the decoder; the paper describes it as a residual block, but the code uses a Transformer stage (see Paper vs code) §3.1. A final 3×3 convolution outputs a residual R, and the restored image is I' = I + R §3.1. Depthwise separable convolutions are used throughout §3.1.
2. Multi-scale patch embedding §3.2 fig. 3. Each path is a depthwise separable deformable convolution (DSDCN): a depthwise deformable 3×3 layer followed by a pointwise layer. Offsets come from a depthwise 3×3 and a pointwise convolution. Layers are stacked, not placed side by side, so each deeper path adds semantic depth §3.2. Each deformable layer is followed by Hardswish §3.2. The offsets are truncated to [−3, 3], which gives a receptive field from 1×1 up to 9×9 §3.2 fig. 4. Cost and parameter counts of DSDCN versus a standard DCN are given in eq. 1–eq. 4.
3. Taylor-expanded multi-head self-attention §3.3. The paper expands exp(QᵀK/√D) to first order and writes the attention as 1 + Q̃ᵀK̃ eq. 7–eq. 8. Every query and key vector is scaled to norm 0.5 so that the expansion stays in the range [−0.25, 0.25] where the first-order form is accurate §3.3 fig. 5. Matrix-product associativity then gives a form whose cost is linear in the number of positions eq. 9. Q, K and V come from a 1×1 convolution followed by a depthwise 3×3 convolution §3.3. The number of heads increases from the top level to the bottom §3.3. Cost of standard MSA is eq. 10; cost of T-MSA is eq. 11, and the paper assumes h×w is much larger than D §3.3.
4. Multi-scale attention refinement (MSAR) §3.4. The remainder of the first-order expansion is not modelled. MSAR corrects it with a gate per head: the spatially concatenated normalised Q and K of each head pass through a grouped convolution whose kernel sizes depend on the head count, then a sigmoid eq. 12. The head output is multiplied by this gate, projected with W^P, and added to the block input eq. 13 §3.4. The gate has one channel per head rather than one per channel, which the authors chose for cost tab. 3.

Pseudocode for the attention with the MSAR gate, written from the equations and checked against the official code. The per-head scale on the last steps is in the code but not in the paper.

# Taylor-expanded attention with MSAR gate  [eq. 9] [eq. 12] [eq. 13]
# Input:  x (B, N, D), N = H * W, D = heads * d
# Output: y (B, N, D)

def taylor_msar_attention(x, heads, d):
    # Q, K, V: 1x1 conv, then depthwise 3x3 conv  [§3.3]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L208]
    queries, keys, values = project_qkv(x)  # each (B, heads, N, d)

    # Scale every query and key vector to norm 0.5  [§3.3] [eq. 8]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L215-L219]
    queries = 0.5 * queries / norm(queries, over=d)
    keys = 0.5 * keys / norm(keys, over=d)

    # MSAR gate: per-head grouped conv over [queries, keys], then sigmoid
    # [eq. 12] [code: basicsr/models/archs/MB_TaylorFormer.py#L222-L224]
    gate = sigmoid(grouped_conv(concat(queries, keys)))  # (B, heads, N, 1)

    # Sums shared by all queries; no N x N matrix is formed  [eq. 9]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L225-L231]
    value_sum = sum(values, over=N)                 # (B, heads, 1, d)
    key_value = transpose(keys) @ values            # (B, heads, d, d)
    key_sum = sum(keys, over=N)                     # (B, heads, 1, d)

    # First-order Taylor form: numerator over denominator  [eq. 9]
    numerator = value_sum + queries @ key_value     # (B, heads, N, d)
    denominator = N + queries @ transpose(key_sum)  # (B, heads, N, 1)
    attended = numerator / (denominator + 1e-6)

    # Learned per-head scale, initialised to 1; not in the equations
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L176] [#L234]
    attended = attended * temperature

    # Apply the gate, merge heads, project  [eq. 13]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L235-L238]
    attended = attended * gate
    return project_out(merge_heads(attended))

The block adds this output to its input, then applies a feed-forward network with a residual connection §3.3 MB_TaylorFormer.py. The paper does not describe the feed-forward network. The code uses a GELU-gated design MB_TaylorFormer.py.

# One encoder or decoder stage  [§3.1]
# Input:  x (B, C, H, W)
# Output: y (B, C, H, W)

def stage(x, num_blocks):
    # Stacked DSDCN layers; each branch takes one layer's output  [§3.2]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L644-L651]
    branch_inputs = multi_scale_patch_embed(x)  # list of (B, C, H, W)

    # Each branch runs its own stack of Transformer blocks  [§3.1]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L287-L298]
    branch_outputs = []
    for tokens in branch_inputs:
        for block_index in range(num_blocks):
            tokens = tokens + attention(norm(tokens))  # see above
            tokens = tokens + feed_forward(norm(tokens))
        branch_outputs.append(tokens)

    # Fuse branches with selective-kernel attention (SKFF)  [§3.1]
    # [code: basicsr/models/archs/MB_TaylorFormer.py#L467-L502]
    fused = skff(branch_outputs)

    # Stage residual  [code: basicsr/models/archs/MB_TaylorFormer.py#L828]
    return fused + x

Hyperparameters

Paper values come first; code values are from the two shipped configs (Dehazing/Options/MB-TaylorFormer-B.yml and -L.yml), which are identical where not noted.

Datasets §4.1: synthetic RESIDE subsets (ITS, OTS; evaluated on SOTS) and the real O-HAZE and Dense-Haze sets. The real sets have 45 and 55 pairs; the paper uses the last 5 of each as test images and the rest for training §4.1.

Results

Author-reported. Setup: the caption of fig. 1 says all models are trained on SOTS-Indoor. The shipped configs train on the ITS training split and validate on ITS/SOTS MB-TaylorFormer-B.yml MB-TaylorFormer-B.yml. The paper does not state the training set for each row of tab. 1. Baselines are the methods in tab. 1; MACs are reported without an input resolution.

Table 1 tab. 1, full, as in the source (best per column in bold and second best underlined in the paper; the formatting is dropped here):

Headline results, from the text §4.2 §4.3:

• SOTS-Indoor: B gains 2.19 dB over SGID-PFF; L gains 4.12 dB tab. 1.
• SOTS-Outdoor: L is 2.91 dB above Dehamer tab. 1.
• O-HAZE: best PSNR is 0.20 dB above Dehamer; Dense-Haze, 0.04 dB above it tab. 1.

Ablations, all on SOTS-Indoor with the B model §4.4:

Table 2 (multi-scale patch embedding, tab. 2):

Table 3 (MSAR, tab. 4):

Table 4 (replacing T-MSA with other linear attention, tab. 5):

Table 5 (approximation order, tab. 7):

Generalisation §4.5, trained and tested as in the cited works:

Table 6, snow removal on CSD tab. 6:

Table 7, rain removal on RainCityscapes tab. 7:

Qualitative results are in fig. 6 (synthetic) and fig. 7 (real). The feature-map comparison of single-scale and multi-scale tokens is fig. 8.

Paper vs code

Official code: FVL2020/ICCV-2023-MB-TaylorFormer, commit c504250a45a91cefb20c9ee706df0a834e62beb0. The paper's abstract links it. The owner is a user account, not an organisation; the README lists the six authors and their university contact addresses. The repository has no LICENSE file, so its code is cited by path and line only, with no quoted snippets. The README points to a later V2 (arXiv:2501.04486); that paper is not covered here.

• Learning-rate schedule. Paper: start 2e-4, cosine to 1e-6 §4.1. Config values: 2e-4 held for 200k iterations, then cosine from 2e-4 to 1e-5 over 300k iterations MB-TaylorFormer-B.yml. The comments in the same file say 3e-4 with a 92k/208k split and a final value of 1e-6 MB-TaylorFormer-B.yml; the numbers do not match the comments.
• Training length. Paper: "epoch is set as 500" for ablations §4.4. Config: 500,000 iterations MB-TaylorFormer-B.yml.
• Branch count. Paper: three branches, with receptive-field ranges ordered x ≤ y ≤ z §3.2. B config uses two per level; L uses three at levels 2–4 and two at level 1 MB-TaylorFormer-B.yml MB-TaylorFormer-L.yml. The code has no x ≤ y ≤ z constraint; it only clamps the offsets MB_TaylorFormer.py.
• Per-head temperature. Not in equations 9 or 13. The code multiplies each head's output by a learned scalar initialised to 1 MB_TaylorFormer.py MB_TaylorFormer.py.
• Refinement stage. Paper: "a residual block" after the encoder–decoder §3.1. Code: a Transformer refinement stage with its own patch embedding MB_TaylorFormer.py MB_TaylorFormer.py. A ResBlock class exists but is not used by the model MB_TaylorFormer.py (its only other reference is commented out, #L394).
• Normalisation type. Config sets LayerNorm_type: WithBias MB-TaylorFormer-B.yml. The model passes 'BiasFree' to every stage, so the config value is ignored MB_TaylorFormer.py (same pattern at other levels).
• Augmentation. Paper: random cropping and flipping §4.1. Code: one of eight flip/rotation variants, applied to both images transforms.py paired_image_dataset.py.
• Gradient clipping. Not in the paper. Code clips the global norm at 0.01 image_restoration_model.py.
• Feed-forward network. Not described in the paper. Code: GELU-gated channel-expansion design MB_TaylorFormer.py.
• Training data. The paper evaluates on OTS, O-HAZE and Dense-Haze §4.1. The repository ships only the two ITS configs ; it has no config or weights for the other datasets as far as the tree shows.
• Pretrained model. B config initialises from an author-local path MB-TaylorFormer-B.yml; L config sets none MB-TaylorFormer-L.yml. Whether B trains from scratch depends on that file existing; not checked.
• Supplementary. The paper defers the T-MSA pseudo-code and the MSAR details to supplementary material §3.3 §3.4, which is not in the LaTeX source and was not fetched.

Pitfalls

• Only heads 1, 2, 4 and 8 work. The MSAR kernel sizes are a fixed table keyed on head count MB_TaylorFormer.py. Any other count leaves the variable undefined, so the model fails at build time. The paper does not mention this.
• Denominators stay positive. After scaling, each Q̃ᵀK̃ lies in [−0.25, 0.25] (from the paper's stated range), so each term 1 + Q̃ᵀK̃ is at least 0.75 and the denominator is at least 0.75N. The code adds 1e-6 as a guard MB_TaylorFormer.py.
• Scaling is per token. Queries are normalised across channels per token; keys are normalised across channels per token too. Normalising keys over the token axis instead would change the attention MB_TaylorFormer.py.
• Gradient through clamp. Offsets are clamped with torch.clamp MB_TaylorFormer.py. Clamped offsets receive zero gradient, so a branch stuck at the limit stops learning its offset.
• Level-1 skip concatenation. Level 1 concatenates an upsampled 24-channel map with the encoder's 24-channel map to feed a 48-channel patch embedding. This works only because the second level's channel count is twice the first level's MB_TaylorFormer.py MB_TaylorFormer.py. The code comments that no 1×1 reduction is used there.
• Text versus table. The text says the B model reaches 0.994 SSIM on SOTS-Indoor §4.2; tab. 1 gives 0.992 for B and 0.994 for L.
• Efficiency claim. The text says B uses "merely 10%" of SGID-PFF's parameters and "14%" of its computational cost §4.2. The table gives 2.68M against 13.87M (about 19%) and 38.5G against 156.4G (about 25%). These ratios are computed from tab. 1; the paper does not state them.
• Approximation order. tab. 5 ranks the first-order variant lowest of the three, while the paper keeps it for cost §4.4. The table's "Swin" row has no softmax-attention reference, so the ordering does not show how far the first-order form is from softmax.
• MACs have no stated input size. tab. 1 and the ablations report MACs without an input resolution, so cost comparisons are only as good as the common input the authors used, which is not stated.
• Restart schedule. The second cosine phase starts from 2e-4 again with a lower floor. Use the config values, not the paper's 1e-6 floor, if you aim to reproduce the released runs.
• Generalisation §4.5 uses CSD and RainCityscapes; the table shows the model only, not the training recipe for those tasks.

Lineage

• Builds on: Restormer (arXiv:2111.09881), whose U-shaped layout the paper adopts §3 and whose channel-wise attention (MDTA) it compares against §2 tab. 4. Deformable convolution §3.2 and SKFF §3.1 are used as cited.
• Compared against: the methods in tab. 1, including linear-attention variants MDTA, Swin, TAR, PVTv2, Cswin and LinFormer tab. 4, and kernel-based linear attention §2.