PaperlayerView exact revision on arXiv

Source-linked AI summary

MAXIM: Multi-Axis MLP for Image Processing

Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, Yinxiao Li

arXiv:2201.02973v2eess.IVcs.CV

TL;DR

Existing Transformer and MLP vision models face high-resolution and locality challenges in low-level image processing. MAXIM uses a UNet-shaped architecture with multi-axis gated MLPs and cross-gating for global, fully-convolutional feature mixing, achieving state-of-the-art results across more than ten benchmarks. Its generalization to real-world examples remains relatively inadequate, and it can overfit certain benchmarks.

  • Problem

    Existing global vision models have limited evidence in low-level restoration and may require fixed patches or restricted receptive fields, hindering flexible high-resolution processing.

  • Method

    MAXIM combines a hierarchical encoder-decoder backbone with multi-axis gated MLPs for local-global mixing and a pure MLP cross-gating block for cross-feature conditioning.

  • Results

    MAXIM achieves state-of-the-art results on more than ten datasets spanning denoising, deblurring, deraining, dehazing, and enhancement.

  • Takeaways & Limitations

    MAXIM provides a general-purpose low-level vision architecture with global receptive fields, arbitrary-resolution support, and linear complexity relative to image size.

  • Takeaways & Limitations

    MAXIM generalizes relatively inadequately to real-world examples and tends to slightly overfit certain benchmarks without strong regularization.

Abstract

from arXiv · show

Recent progress on Transformers and multi-layer perceptron (MLP) models provide new network architectural designs for computer vision tasks. Although these models proved to be effective in many vision tasks such as image recognition, there remain challenges in adapting them for low-level vision. The inflexibility to support high-resolution images and limitations of local attention are perhaps the main bottlenecks. In this work, we present a multi-axis MLP based architecture called MAXIM, that can serve as an efficient and flexible general-purpose vision backbone for image processing tasks. MAXIM uses a UNet-shaped hierarchical structure and supports long-range interactions enabled by spatially-gated MLPs. Specifically, MAXIM contains two MLP-based building blocks: a multi-axis gated MLP that allows for efficient and scalable spatial mixing of local and global visual cues, and a cross-gating block, an alternative to cross-attention, which accounts for cross-feature conditioning. Both these modules are exclusively based on MLPs, but also benefit from being both global and `fully-convolutional', two properties that are desirable for image processing. Our extensive experimental results show that the proposed MAXIM model achieves state-of-the-art performance on more than ten benchmarks across a range of image processing tasks, including denoising, deblurring, deraining, dehazing, and enhancement while requiring fewer or comparable numbers of parameters and FLOPs than competitive models. The source code and trained models will be available at \url{https://github.com/google-research/maxim}.

1. Introduction

Image restoration and enhancement require models that handle diverse degradations, but existing global vision models struggle with high-resolution inputs and local attention limits. MAXIM addresses these issues with a fully-convolutional, multi-axis MLP architecture that mixes local and global information.

  • Low-level vision tasks seek desired outputs from degraded inputs, including denoising, deblurring, super-resolution, and dehazing.
  • Global Transformer and MLP models have shown promise in high-level vision, but their effectiveness for low-level enhancement and restoration remains insufficiently studied.
  • Full self-attention requires small fixed-size patches, causing boundary artifacts on larger images, while local attention limits receptive fields or non-locality.
  • MAXIM uses parallel multi-axis MLP branches to capture local and global interactions, support arbitrary resolutions, and scale linearly with image size.
  • MAXIM combines encoder-decoder backbones, multi-scale multi-stage supervision, multi-axis gated MLPs, and cross-gating blocks in a generic image-processing architecture.

2. Related Work

Related work in image restoration has progressed through encoder-decoder designs, multi-scale and multi-stage frameworks, attention, non-local networks, and Transformer or MLP alternatives. These approaches motivate MAXIM’s focus on efficient global interactions for low-level vision.

  • CNN-based restoration models achieve state-of-the-art performance across image-enhancement tasks, supported by large-scale benchmarks and specialized architectural designs.
  • UNet-inspired encoder-decoder designs, multi-scale learning, attention, and non-local networks have improved image-restoration models over earlier single-scale approaches.
  • Multi-stage frameworks have shown promising results on challenging deblurring and deraining tasks, building on their success in higher-level vision problems.
  • Low-level vision Transformers apply attention to image patches, extending Transformer architectures from their original language and recognition settings.
  • MLP vision models replace self-attention with token mixing or spatial gating, establishing all-MLP alternatives for visual processing.

3. Our Approach: MAXIM

MAXIM is a general-purpose low-level vision architecture combining an encoder-decoder backbone with multi-axis MLPs for local and global mixing and cross-gating for feature conditioning. Its fully-convolutional design supports arbitrary resolutions and linear scaling with image size, while multi-stage, multi-scale supervision supports image restoration across tasks.

  • 3.2. Multi-Axis Gated MLP: MAXIM's multi-axis operator provides global receptive fields, supports arbitrary input resolutions, and scales linearly with image size HW.The design applies one-dimensional operators along a single axis per branch while sharing parameters across the other spatial dimensions.
  • 3.1. Main Backbone: MAXIM combines an encoder-decoder backbone with multi-axis gated MLP and residual channel attention blocks, plus cross-gating across skip-connections.Global bottleneck features gate skip connections, while refined global features propagate upward through successive cross-gating blocks.
  • 3.2. Multi-Axis Gated MLP: The multi-axis gated MLP splits features into local blocked and global gridded branches that mix spatial information in parallel.The local branch uses non-overlapping windows, whereas the global branch uses a fixed grid and applies gMLP operations along complementary axes.
  • 3.3. Cross Gating MLP Block: The cross-gating block is a pure-MLP alternative to cross-attention that reciprocally conditions two features through multi-axis gating and residual output projections.Element-wise multiplication combines each feature with gating weights extracted from the other feature.
  • 3.4. Multi-Stage Multi-Scale Framework: MAXIM uses a multi-stage, multi-scale framework with end-to-end losses across stages and scales, including Charbonnier and frequency reconstruction losses.The framework predicts S × N restored outputs from multi-scale inputs and applies frequency reconstruction to enforce high-frequency details.

4. Experiments

MAXIM is evaluated across five image-processing tasks and 17 datasets, achieving strong quantitative and visual results with moderate computational complexity. Ablations further examine its components and multi-axis design.

  • MAXIM is evaluated on denoising, deblurring, deraining, dehazing, and enhancement across 17 datasets.
  • Denoising: 0.24 dB PSNR improvement over MIRNet on SIDD accompanies competitive 39.84 dB PSNR on DND.
  • Deblurring: 0.15 dB PSNR gain over HINet and 32.83 dB PSNR on HIDE demonstrate strong synthetic deblurring performance.
  • Deraining: MAXIM improves over prior state-of-the-art deraining results on all datasets, with a 0.24 dB average PSNR gain over HINet.
  • Dehazing: MAXIM surpasses the previous best model by 0.94 dB and 0.62 dB PSNR on the SOTS indoor and outdoor sets, respectively.
  • Enhancement / Retouching: The model achieves the best reported PSNR and SSIM values on FiveK and LOL, while visual examples show more natural colors and clearer enhancements.
  • Ablation: Ablations report PSNR gains of 0.25, 0.63, 0.36, and 0.26 dB for progressively added components, while multi-axis MAB contributes the major improvement.

C. Why multi-stage?

The experiments compare scaling MAXIM by width, depth, and stages to identify the most favorable performance–complexity tradeoff. Multi-stage packing performs best among these scaling choices, while operator families differ in accuracy and efficiency.

  • 32.44 dB PSNR, 22.2 M parameters, and 339.2 G FLOPs make multi-stage packing the best performance–complexity tradeoff.
  • Self-attention variants achieve the best performance, whereas FFT and MLP families are more computationally efficient.

5. Conclusion

The paper presents MAXIM as a generic MLP-based network for image restoration and enhancement, with strong results at moderate complexity. It also identifies broader applications, efficiency improvements, and responsible-use concerns as boundaries for future work.

  • MAXIM is presented as a generic network for restoration and enhancement tasks, using MLP-based global modeling for low-level vision.
  • The MAXIM family advances state-of-the-art results across several image enhancement and restoration tasks with moderate complexity.
  • Future work includes more efficient models for extremely high-resolution processing and large models adaptable to multiple tasks.
  • The authors note privacy risks from malicious use of enhancement techniques and possible bias in models trained on specific data.

A.1. Datasets and Training Details

The appendix summarizes datasets and task-specific training details for MAXIM’s image-processing evaluations. Training uses task-dependent datasets, losses, and iteration schedules.

  • All datasets used in the paper are summarized in Table 8, with separate training details described for each task.
  • Image Denoising: SIDD training uses 320 high-resolution images, while evaluation uses 1,280 SIDD patches and 1,000 DND images.
  • Image Deraining: Deraining training uses 13,712 clean-rain pairs and evaluates on five test sets; raindrop removal uses 861 training image pairs.

A.2. Architecture Details

MAXIM uses a general multi-stage, multi-scale framework in which each stage contains a single-stage MAXIM backbone. Supervised attention, cross-stage fusion, and multi-scale supervision support later-stage learning.

  • Each stage uses a single-stage MAXIM backbone within a general multi-stage framework.
  • Supervised attention and cross-stage feature fusion help later stages learn.
  • Multi-scale input-output processing generates downscaled input variants and bilinearly downscaled ground-truth variants for supervision.

A.2.1 Configurations

A single-stage MAXIM uses encoder-decoder components with MAB and RCAB blocks, while its grid and block sizes vary by feature resolution. The architecture requires padding inputs so both dimensions are divisible by 64.

  • Architecture: The Encoder uses MAB and RCAB as major components across Encoder, Decoder, and Bottleneck blocks.Bottleneck RCAB blocks replace Conv3x3 layers with Conv1x1 layers.
  • Decoder and CGB configuration: The Decoder mirrors the Encoder configuration, while CGB neck depths use b = d = 16 at depths 1–2 and b = d = 8 at depth 3.
  • Resolution-dependent configuration: Block and grid sizes are 16 for feature sizes ≥128 and 8 for feature sizes <128.
  • Input requirements: Input images must be padded so both spatial dimensions are divisible by 64.

A.2.2 Comparison with Other MLPs

MAXIM combines local and global spatial interactions through multi-axis MLPs while retaining global, fully-convolutional operation with linear pixel complexity. Across restoration and enhancement tasks, it delivers strong visual and quantitative performance with moderate complexity, though real-world generalization remains limited.

  • Receptive fields: MAXIM’s multi-axis design provides sparse local and global dilated spatial interactions in parallel.The receptive-field visualization contrasts MAXIM with MLP-Mixer, gMLP, and Swin-Mixer.
  • Comparison with other MLPs: MAXIM is both fully-convolutional and global, with linear complexity relative to the number of pixels N.The comparison is summarized against other MLP models in Table 10.
  • Performance versus complexity: 21% FLOPs and 70% parameters of MIRNet are required by MAXIM-3S for denoising, while deblurring uses 25% of HINet’s parameters and 19% of IPT’s.These comparisons are reported as performance-versus-complexity results across tasks.
  • Deblurring: 0.1 dB of PSNR separates MAXIM-3S above HINet on the REDS JPEG-artifact deblurring benchmark.Scores are averaged over 300 validation images.
  • Raindrop removal: 31.87 dB and 25.74 dB PSNR are achieved by MAXIM-2S on Raindrop testsets A and B, respectively.
  • Visual comparisons: MAXIM produces sharp, conservative outputs without hallucinating many nonexistent details, and visual comparisons cover denoising, deblurring, deraining, dehazing, retouching, and low-light enhancement.The visual results report improved detail recovery or degradation removal in several tasks.
  • Weight visualizations: MAXIM’s learned block and grid gMLP weights exhibit locality and spatial invariance, with often coupled shapes within the same layer.The global grid gMLP also learns local operations on a uniform dilated grid, but no clear stage-wise filter trend is observed.
  • Limitations: MAXIM may generalize inadequately to real-world examples and slightly overfit certain benchmarks because strong regularization was not applied.The authors suggest more realistic large-scale datasets and future regularization as possible directions.
Loading 2201.02973v2…