F2T2-HiT: FFT & Hierarchical Transformer SIRR

• F2T2-HiT is a transformer-based architecture that unites FFT attention and hierarchical Transformer blocks in a U-shaped encoder–decoder for single image reflection removal.
• It employs a dual-branch design, splitting computation into spatial and frequency branches with FFT and multi-scale window self-attention for robust feature extraction.
• Experimental results show improved PSNR and SSIM scores across benchmarks, demonstrating effective removal of complex reflections and preservation of fine details.

F2T2-HiT refers to a Transformer-based neural architecture for the Single Image Reflection Removal (SIRR) problem, specifically the "U-Shaped FFT Transformer and Hierarchical Transformer" introduced in (Cai et al., 5 Jun 2025). This design unites two fundamental attention mechanisms—Fast Fourier Transform (FFT) attention and multi-scale Hierarchical Transformer (HiT) blocks—within a U-shaped encoder–decoder (UNet) framework to enable state-of-the-art performance in removing unwanted reflections from single images. F2T2-HiT addresses the core SIRR challenge: reflections in natural images exhibit wide heterogeneity in frequency, spatial extent, coverage, and detail, making both global context and multi-scale local understanding critical for high-quality restoration.

1. Architectural Composition

The backbone of F2T2-HiT is a single-stage U-shaped encoder–decoder closely related to NAFNet, but with original blocks replaced by three "plain" modules at each resolution stage: the NAF block, Hierarchical Transformer (HiT) block, and FFT-Transformer (F2T2) block. The encoder path consists of four down-sampling stages that halve spatial resolution and double the channel dimension at each step; the decoder is symmetric, with corresponding up-sampling stages. Skip-connections directly connect encoder and decoder features at matching resolutions, preserving spatial detail at each level. Layer Normalization is applied prior to each block, and GELU nonlinearities replace standard ReLU activations throughout all feed-forward networks.

1.1 FFT-Transformer (F2T2) Block

Each F2T2 block bifurcates computation into a spatial branch for local pattern modeling and a frequency branch for global feature extraction. Given input $X\in\mathbb{R}^{C\times H\times W}$, the process involves:

• Spatial branch: A compact spatial FFN and depthwise convolution yield $X_s$.
• Frequency branch: Compute the 2D FFT,

$F = \mathcal{F}(X)$

followed by linear projections in the frequency domain

$Q = W_Q F,\quad K = W_K F,\quad V = W_V F$

and frequency self-attention:

$$A_{\mathrm{freq}} = \mathrm{softmax}\left(\frac{Q K^H}{\sqrt{d}}\right) V$$

The output is returned to the spatial domain via the inverse FFT:

$X_f = \mathcal{F}^{-1}(A_{\mathrm{freq}})$

• Fusion: Spatial and frequency branches are combined, typically through summation or a $1\times1$ convolution, generating the final block output $X'$.

This design ensures extensive receptive fields even at shallow network layers, enabling early access to global context.

1.2 Hierarchical Transformer (HiT) Block

The HiT block generalizes windowed self-attention across three spatial scales: window sizes $s\in\{4,8,16\}$. For each $s$:

• The feature map is partitioned into disjoint $X_s$0 windows.
• Each window undergoes a Spatial-Channel Correlation (S-SC and C-SC) self-attention mechanism of linear cost in $X_s$1.
• Outputs $X_s$2 are produced for each scale and fused via softmax-weighted sum, where the weights $X_s$3 are learned:

$X_s$4

• A channel-wise FFN (Gated Linear Unit and depthwise conv) further processes $X_s$5.

The HiT block captures both fine, localized patterns (small windows) and broad, long-range dependencies (large windows) within linear computational complexity per window.

2. Mathematical Operations

Key differentiable operations governing F2T2-HiT include:

• Layer Normalization:

$X_s$6

• GELU Activation:

$X_s$7

• Depthwise Convolution in FFN:

$X_s$8

• Channel Attention (SCA): Spatial dimensions are squeezed, processed via a small MLP, then used to re-weight channel activations.

The learning objective is the standard pixel-wise L1 loss:

$X_s$9

where $F = \mathcal{F}(X)$0 is the network reconstruction and $F = \mathcal{F}(X)$1 is the ground-truth background image.

3. Training Procedure and Implementation

F2T2-HiT is implemented in PyTorch, trained on 8 NVIDIA A100 GPUs with a batch size of 8, using random $F = \mathcal{F}(X)$2 image crops. Inputs undergo data augmentation through random horizontal flips and random $F = \mathcal{F}(X)$3 rotations. Training utilizes the Adam optimizer (initial learning rate $F = \mathcal{F}(X)$4) and a cosine-annealing scheduler with restarts: three successive $F = \mathcal{F}(X)$5-iteration periods, with restart multipliers $F = \mathcal{F}(X)$6, $F = \mathcal{F}(X)$7, and $F = \mathcal{F}(X)$8 (total $F = \mathcal{F}(X)$9 iterations). The total training time is approximately 24 hours.

4. Experimental Results and Benchmarking

F2T2-HiT is trained on a composite dataset comprising 89 real image pairs from Zhang et al. (CVPR 2018), 200 real pairs from Li et al. (CVPR 2020), over 22,000 synthetic pairs from PASCAL VOC 2007/2012, as well as the RRW dataset synthesized via Hu et al.'s method (CVPR 2024). Evaluation occurs on three real-image benchmarks: Nature (20 images), Real (20 images), and SIR² (454 images).

Performance is assessed by PSNR [$Q = W_Q F,\quad K = W_K F,\quad V = W_V F$0] and SSIM [$Q = W_Q F,\quad K = W_K F,\quad V = W_V F$1]. F2T2-HiT achieves the following (Table 1 from (Cai et al., 5 Jun 2025)):

Method	Nature (PSNR/SSIM)	Real (PSNR/SSIM)	SIR² (PSNR/SSIM)	Average (PSNR/SSIM)
Input Image	20.44 / 0.785	18.96 / 0.733	22.76 / 0.885	22.51 / 0.878
Zhu et al. (24)	25.96 / 0.843	23.82 / 0.817	25.45 / 0.910	25.40 / 0.904
F2T2-HiT	26.08 / 0.837	21.64 / 0.766	25.72 / 0.903	25.57 / 0.894

PSNR gains are highest on Nature and SIR². The Real set remains particularly challenging, likely due to extreme diversity in reflection and image content.

Ablation studies reveal that the hierarchical Transformer (HiT) block confers a +0.22 dB PSNR improvement over the Restormer block baseline, while adding the F2T2 branch provides a further +0.57 dB, resulting in an aggregate quantitative improvement.

Qualitative evaluation shows F2T2-HiT effectively removes large, diffuse reflections (such as skylight), while reliably preserving fine image details (e.g., foliage). Residual artifacts, typically “ghosting,” emerge only when the background and reflection layers share highly similar frequency content.

5. Core Insights and Mechanistic Analysis

F2T2-HiT's design yields robust SIRR performance by tightly integrating global and multi-scale local cues:

• FFT-based attention ensures global context awareness and effective extraction of recurring low-frequency reflection patterns at all network depths.
• Hierarchical multi-scale windowing enables concurrent modeling of both small-scale details and broad structural patterns.
• Skip-connections in the U-shaped backbone facilitate high-fidelity propagation of spatial detail, critical for image restoration tasks.

The architecture’s FFT-based modeling is particularly advantageous because real-world reflections often manifest as distinct, strong frequency components. Windowed and hierarchical attention allows the system to adapt to the extensive spatial and frequency variability exhibited by reflections in natural photographs.

6. Limitations and Prospective Developments

While F2T2-HiT establishes new quantitative baselines, it is subject to several limitations:

• FFT operations, although efficient, introduce computational overhead compared to purely spatial alternatives. Sparse or learned frequency bases could further reduce cost and improve adaptability.
• The model is trained exclusively in a single-image setting; extension to multi-view or polarization inputs is a prospective direction for increasing robustness to more challenging reflection scenarios.
• The exclusive use of an L1 loss, without adversarial (GAN) or perceptual losses, may limit the system's ability to fully optimize perceptual image quality, especially in scenes with subtle reflections or ambiguous boundaries.

This suggests that the F2T2-HiT framework can be extended to broader image restoration contexts where global frequency patterns and local flexibility are important, subject to investigation of frequency domain representation learning and multi-instance or time-resolved image inputs.

7. Significance within Reflection Removal Research

F2T2-HiT represents an overview of global FFT-based attention with multi-scale hierarchical Transformer techniques in a UNet-style architecture, resulting in improved SIRR capabilities, especially for images with complex or large-scale reflection patterns. The explicit mathematical formulations, practical ablation benchmarks, and qualitative evidence collectively demonstrate that frequency and spatial attention, when unified, provide strong inductive biases for ill-posed image separation tasks (Cai et al., 5 Jun 2025). The framework sets a new reference point for future exploration of transformer-based architectures in image decomposition and restoration problems.