CubeFormer: A Simple Yet Effective Lightweight Image Super‑Resolution Baseline

Introduction

Single‑image super‑resolution (SR) aims to reconstruct a high‑resolution (HR) image from a degraded low‑resolution (LR) input. Recent Vision‑Transformer (ViT) based SR methods achieve strong performance but demand heavy computation, making deployment difficult. Existing lightweight SR approaches fall into two categories: homogeneous structures that use only spatial or channel attention, and heterogeneous structures that simply stack both, often failing to fully exploit low‑level features and resulting in blurry textures.

To address these limitations, the authors propose CubeFormer, a simple yet effective baseline that enhances feature diversity through a novel cube attention mechanism and two specialized transformer blocks (Intra‑CTB and Inter‑CTB). This design enables comprehensive information interaction in three‑dimensional space, leading to richer texture recovery.

Related Work

Early CNN‑based SR models such as SRCNN, DRCN, CARN, and IMDN progressively improved feature aggregation. More recent ViT‑based methods (SwinIR, ESRT, HAT) achieve higher quality but remain computationally expensive. Lightweight visual Transformers (LVT, MobileViT) attempt to reduce parameters via convolutional self‑attention or matrix factorization, yet still struggle with fine‑detail reconstruction.

Method

Overall Pipeline

CubeFormer takes an LR image, extracts shallow features with a convolutional layer, passes them through a backbone composed of cascaded Cube Transformer Groups (CTGs), and finally reconstructs the HR image using a pixel‑shuffle module.

Cube Attention

The cube attention extends conventional 2‑D attention to 3‑D space by partitioning the feature map into non‑overlapping cubes. It follows four stages: (1) QKV generation via convolutional projections, (2) cube embedding that reshapes each cube into a vector, (3) affinity matrix computation using softmax on the dot‑product of Q and K, and (4) cube merging to restore spatial resolution.

Two sampling strategies are employed:

Block sampling creates intra‑cubes for local self‑attention, capturing fine‑grained voxel relationships within each cube.

Grid sampling forms inter‑cubes for sparse global attention, enabling long‑range feature interaction across cubes.

These mechanisms are instantiated in two transformer blocks:

Intra‑CTB focuses on local detail extraction via intra‑cube attention.

Inter‑CTB aggregates global context through inter‑cube attention.

Both blocks follow a standard Transformer two‑stage architecture (LayerNorm → Attention → LayerNorm → Feed‑Forward Network) but replace the attention module with the respective cube attention variant.

CubeFormer‑lite

To further reduce parameters, CubeFormer‑lite replaces the full CTG with a lightweight CTG‑lite. After channel shuffle, a channel‑split operation divides the feature map into two halves; one half proceeds through Intra‑CTB and Inter‑CTB, while the other bypasses them and is later concatenated. This selective processing cuts the number of parameters in the transformer blocks without sacrificing performance.

Learning Strategy

The training objective combines a spatial reconstruction loss (L2 distance between predicted and ground‑truth HR images) and a frequency reconstruction loss that emphasizes high‑frequency components via Fast Fourier Transform. The total loss is a weighted sum: Loss_total = Loss_spatial + 0.01 * Loss_frequency.

Experiments

Experimental Setup

Models are trained on the DIV2K dataset and evaluated on BSD100, Urban100, and Manga109 using PSNR and SSIM on the Y channel of YCbCr. LR images are generated by bicubic down‑sampling with scale factors 2×, 3×, and 4×. Data augmentation includes random horizontal flips, 90°/270° rotations, and channel shuffle. CubeFormer uses six CAG layers with 64 channels each; both intra‑ and inter‑cube attentions employ four‑head multi‑head attention. Cube dimensions and grid sizes are set to 8. Training runs for 800 k iterations with batch size 32, initial learning rate 1e‑4 (halved every 200 k iterations), using the Adam optimizer. All experiments are conducted on an NVIDIA RTX 4090 GPU with PyTorch implementation.

Comparison with State‑of‑the‑Art

CubeFormer is compared against 14 recent lightweight and efficient SR methods, including IMDN, RFDN, LatticeNet, SwinIR, RLFN, ESRT, Shuffemixer, GASSL, MLRN, SAFMN, OSFFNet, SeemoRe, SRConvNet, and OmniSR. Quantitative results (Table 1) show that CubeFormer consistently achieves the highest PSNR/SSIM across all scales. CubeFormer‑lite also sets new records on BSD100, Urban100, and Manga109 while using fewer parameters than competing efficient models such as ShuffleMixer and GASSL‑S.

Ablation Study

Replacing Cube Attention with only spatial attention, only channel attention, or their simple combination leads to lower PSNR, confirming the superiority of the heterogeneous cube attention design. Further ablations on the CTG modules show that removing either Intra‑CTG or Inter‑CTG degrades performance, highlighting the complementary roles of local detail extraction and global context aggregation.

Conclusion

CubeFormer demonstrates that a simple baseline equipped with cube attention and dual transformer blocks can achieve state‑of‑the‑art performance in lightweight image super‑resolution. The proposed mechanisms effectively increase feature diversity, enable both local and global modeling, and maintain a low parameter count, making CubeFormer suitable for resource‑constrained deployment.