\section{Introduction}
\label{sec:intro}

Image processing tasks, such as restoration and enhancement, are important computer vision problems, which aim to produce a desired output from a degraded input.
Various types of degradations may require different image enhancement treatments, such as denoising, deblurring, super-resolution, dehazing, low-light enhancement, and so on.
Given the increased availability of curated large-scale training datasets, recent high-performing approaches~\cite{zamir2020learning,zamir2021multi,zhu2012deconvolving,yli_comisr_iccv2021,cho2021rethinking,chen2021pre,matlin2012removal,chen2020proxiqa,liang2021swinir,delbracio2020projected} based on highly designed convolutional neural network (CNN) have demonstrated state-of-the-art (SOTA) performance on many tasks.

\begin{figure}[!t]
\center
\footnotesize
\setlength{\tabcolsep}{0pt}
\def0.185{0.185}
\begin{tabular}{@{}c!{}c!{}c!{}c!{}c@{}}
\includegraphics[height=0.185\textwidth]{figures/teaser_denoise.pdf}  &  
\includegraphics[height=0.185\textwidth]{figures/teaser_deblur_1.pdf} &
\includegraphics[height=0.185\textwidth]{figures/teaser_derain.pdf} &
\includegraphics[height=0.185\textwidth]{figures/teaser_dehaze.pdf} &
\includegraphics[height=0.185\textwidth]{figures/teaser_enhance.pdf} \\
\end{tabular}
\vspace{-3mm}
\caption{Our proposed MAXIM model significantly advances state-of-the-art performance on five image processing tasks in terms of PSNR: 1) Denoising ($+$\textbf{0.24} dB on SIDD~\cite{abdelhamed2018high}), 2) Deblurring ($+$\textbf{0.15} dB on GoPro~\cite{nah2017deep}) 3) Deraining ($+$\textbf{0.86} dB on Rain100L~\cite{yang2017deep}), 4) Dehazing ($+$\textbf{0.94} dB on RESIDE~\cite{li2019benchmarking}), and 5) Retouching (Enhancement) ($+$\textbf{1.15} dB on FiveK~\cite{bychkovsky2011learning}).}
\label{fig:teaser-figure}
\vspace{-5mm}
\end{figure}

Improving the architectural design of the underlying model is one of the keys to improving the performance of most computer vision tasks, including image restoration. 
Numerous researchers have invented or borrowed individual modules or building blocks and implemented them into low-level vision tasks, including residual learning~\cite{zhang2017beyond,ledig2017photo,wang2018esrgan}, dense connections~\cite{zhang2018residual,wang2018esrgan}, hierarchical structures~\cite{lai2017deep,kupyn2019deblurgan,jiang2021enlightengan}, multi-stage frameworks~\cite{zamir2021multi,chen2021hinet,jiang2020multi,zhang2019deep}, and attention mechanisms ~\cite{zamir2020learning,zamir2021multi,niu2020single,wang2019spatial}.

Recent research explorations on Vision Transformers (ViT)~\cite{dosovitskiy2021an,carion2020end,liu2021Swin} have exemplified their great potential as alternatives to the go-to CNN models.
The elegance of ViT~\cite{dosovitskiy2021an} has also motivated similar model designs with simpler global operators such as MLP-Mixer~\cite{tolstikhin2021mlp}, gMLP~\cite{liu2021pay}, GFNet~\cite{rao2021global}, and FNet~\cite{lee2021fnet}, to name a few.
Despite successful applications to many high-level tasks~\cite{dosovitskiy2021an,liu2021Swin,vaswani2021scaling,arnab2021vivit,xie2021segformer,sun2020transtrack}, the efficacy of these \textit{global} models on low-level enhancement and restoration problems has not been studied extensively.
The pioneering works on Transformers for low-level vision~\cite{chen2021pre,cao2021video} directly applied full self-attention, which only accepts relatively small patches of fixed sizes (e.g., 48$\times$48).
Such a strategy will inevitably cause patch boundary artifacts when applied on larger images using cropping~\cite{chen2021pre}.
Local-attention based Transformers~\cite{liang2021swinir,wang2021uformer} ameliorate this issue, but they are also constrained to have limited sizes of receptive field, or to lose non-locality~\cite{wang2018non,dosovitskiy2021an}, which is a compelling property of Transformers and MLP models relative to hierarchical CNNs.

To overcome these issues, we propose a generic image processing network, dubbed \textbf{{MAXIM}}, for low-level vision tasks.
A key design element of MAXIM is the use of \textit{multi-axis} approach (\cref{ssec:multi-axis-gating-mlp}) that captures both local and global interactions in parallel.
By mixing information on \textit{a single axis} for each branch, this MLP-based operator becomes `\textit{fully-convolutional}' and scales linearly with respect to image size, which significantly increases its flexibility for dense image processing tasks.
We also define and build a pure MLP-based cross-gating module, which adaptively \textit{gate} the skip-connections in the neck of MAXIM using the same multi-axis approach, and which further boosts performance.
Inspired by recent restoration models, we develop a simple but effective multi-stage, multi-scale architecture consisting of a stack of MAXIM backbones. MAXIM achieves strong performance on a range of image processing tasks, while requiring very few number of parameters and FLOPs.
Our contributions are:

\vspace{-2mm}
\begin{itemize}[leftmargin=*]
\itemsep0em 
    \item A novel and generic architecture for image processing, dubbed MAXIM, using a stack of encoder-decoder backbones, supervised by a multi-scale, multi-stage loss.
    \vspace{-0.6mm}
    \item A multi-axis gated MLP module tailored for low-level vision tasks, which always enjoys a global receptive field, with linear complexity relative to image size.
    \vspace{-0.6mm}
    \item A cross gating block that cross-conditions two separate features, which is also global and fully-convolutional.
    \vspace{-0.6mm}
    \item Extensive experiments show that MAXIM achieves SOTA results on more than 10 datasets including denoising, deblurring, deraining, dehazing, and enhancement.
   
\end{itemize}

\section{Related Work}
\label{sec:related-work}

\noindent\textbf{Restoration models.}
Driven by recent enormous efforts on building vision benchmarks, learning-based models, especially CNN models, have been developed that attain state-of-the-art performance on a wide variety of image enhancement tasks~\cite{zamir2021multi,chen2021hinet,jiang2021enlightengan,yli_comisr_iccv2021,shen2019human,chen2020proxiqa,chen2021pre,liang2021swinir}. These increased performance gains can be mainly attributed to novel architecture designs, and/or task-specific modules and units. For instance, UNet~\cite{ronneberger2015u} has incubated many successful encoder-decoder designs~\cite{zamir2021multi,jiang2021enlightengan,cho2021rethinking} for image restoration that improve on earlier single-scale feature processing models~\cite{zhang2017beyond,li2017aod}. Advanced components developed for high-level vision tasks have been brought into low-level vision tasks as well. Residual and dense connections~\cite{zhang2017beyond,ledig2017photo,wang2018esrgan, zhang2018residual,wang2018esrgan}, the multi-scale feature learning~\cite{kupyn2019deblurgan,wang2021uformer,cho2021rethinking}, attention mechanisms~\cite{zamir2020learning,zamir2021multi,niu2020single,wang2019spatial,zhang2018residual}, and non-local networks~\cite{wang2018non,liu2018non,zhang2018residual} are such good examples.
Recently, \textit{multi-stage} networks~\cite{zamir2021multi,chen2021hinet,jiang2020multi,zhang2019deep} have attained promising results relative to the aforementioned \textit{single-stage} models on the challenging deblurring and deraining tasks~\cite{dong2020multi,jiang2020multi,zamir2021multi}.
These multi-stage frameworks are generally inspired by their success on higher-level problems such as pose estimation~\cite{chen2018cascaded, li2019rethinking}, action segmentation~\cite{farha2019ms,li2020ms}, and image generation~\cite{zhang2017stackgan,zhang2018stackgan++}.

\noindent\textbf{Low-level vision Transformers.}
Transformers were originally proposed for NLP tasks~\cite{vaswani2017attention}, where multi-head self-attention and feed-forward MLP layers are stacked to capture non-local interactions between words.
Dosovitskiy \etal coined the term Vision Transformer (ViT)~\cite{dosovitskiy2021an}, and demonstrated the first pure Transformer model for image recognition. Several recent studies explored Transformers for low-level vision problems, \eg, the pioneering pre-trained image processing Transformer (IPT)~\cite{chen2021pre}.
Similar to ViT, IPT directly applies vanilla Transformers to image patches. The authors of~\cite{cao2021video} presented a spatial-temporal convolutional self-attention network that exploits local information for video super-resolution. More recently, Swin-IR~\cite{liang2021swinir} and UFormer~\cite{wang2021uformer} apply efficient window-based local attention models on a range of image restoration tasks. 

\noindent\textbf{MLP vision models.}
More recently, several authors have argued that when using a patch-based architecture as in ViT, the necessity of complex self-attention mechanisms becomes questionable. For instance, MLP-Mixer~\cite{tolstikhin2021mlp} adopts a simple token-mixing MLP to replace self-attention in ViT, resulting in an all-MLP architecture. The authors of~\cite{liu2021pay} proposed the gMLP, which applies a spatial gating unit on visual tokens. ResMLP~\cite{touvron2021resmlp} adopts an Affine transformation as a substitute to Layer Normalization for acceleration. Very recent techniques such as FNet~\cite{lee2021fnet} and GFNet~\cite{rao2021global} demonstrate the simple Fourier Transform can be used as a competitive alternative to either self-attention or MLPs. 



\begin{figure*}[!htb]
\centering
\includegraphics[width=0.95\linewidth]{figures/MAX-MLP-New.pdf}
\vspace{-2mm}
\caption{\textbf{MAXIM architecture.} We take (a) an encoder-decoder backbone with each (b) encoder, decoder, and bottleneck containing a multi-axis gated MLP block (\cref{fig:multi-axis-gated-mlp-block}) as well as a residual channel attention block. The model is further boosted by (c) a cross gating block which allows global contextual features to gate the skip-connections. More detailed description can be found in Appendix.}
\vspace{-3mm}
\label{fig:MAXIM-backbone}
\end{figure*}

\section{Our Approach: MAXIM}
\label{sec:approach}

We present, to the best of our knowledge, the first effective general-purpose MLP architecture for low-level vision, which we call \textbf{M}ulti-\textbf{AXI}s \textbf{M}LP for image processing (\textbf{MAXIM}). 
Unlike previous low-level Transformers~\cite{chen2021pre,cao2021video,liang2021swinir,wang2021uformer}, MAXIM has several desired properties, making it intriguing for image processing tasks.
First, MAXIM expresses global receptive fields on arbitrarily large images with linear complexity; Second, it directly supports arbitrary input resolutions, \ie, being fully-convolutional; Lastly, it provides a balanced design of local (\texttt{Conv}) and global (\texttt{MLP}) blocks, outperforming SOTA methods without the necessity for large-scale pre-training~\cite{chen2021pre}. 

\subsection{Main Backbone}
\label{ssec:main-architecture}

The MAXIM backbone (\cref{fig:MAXIM-backbone}a) follows the encoder-decoder design principles that originated with UNet~\cite{ronneberger2015u}.
We have observed that operators having small footprints such as \texttt{Conv3x3} are essential to the performance of UNet-like networks. Thus, we rely on a hybrid model design for each block (\cref{fig:MAXIM-backbone}b) -- \texttt{Conv} for local, and \texttt{MLP} for long-range interactions -- to make the most of them.


To allow long-range spatial mixing at different scales, we insert the multi-axis gated MLP block (MAB) into each encoder, decoder, and bottleneck (\cref{fig:MAXIM-backbone}b), with a residual channel attention block (RCAB)~\cite{woo2018cbam,zamir2021multi} (\texttt{LayerNorm}-\texttt{Conv}-\texttt{LeakyReLU}-\texttt{Conv}-\texttt{SE}~\cite{hu2018squeeze}) stacked subsequently.
Inspired by the gated filtering of skip connections~\cite{oktay2018attention, petit2021u}, we extend the gated MLP (gMLP) to build a cross gating block (CGB, \cref{fig:MAXIM-backbone}c), which is an efficient 2nd-order alternative to cross-attention (3rd-order correlations), to interact, or condition two distinct features.
We leverage the global features from \textcolor{brickred}{Bottleneck} (\cref{fig:MAXIM-backbone}a) to gate the skip connections, while propagating the refined global features upwards to the next CGB.
Multi-scale feature fusion~\cite{sun2019deep,zamir2020learning,cho2021rethinking} (\textcolor{brickred}{red} and \textcolor{blueish}{blue} lines) is utilized to aggregate multi-level information in the Encoder$\rightarrow$CGB and CGB$\rightarrow$Decoder dataflow.

\subsection{Multi-Axis Gated MLP}
\label{ssec:multi-axis-gating-mlp}

\begin{figure*}[!htb]
\centering
 \includegraphics[width=0.95\linewidth]{figures/Taller_Multi-Axis_gMLP_Block.pdf}
 \caption{\textbf{Multi-axis gated MLP block} (best viewed in color). The input is first projected to a $[6,4,C]$ feature, then split into two heads. In the \textcolor{brickred}{local branch}, the half head is \textit{blocked} into $3\times 2$ non-overlapping $[\textcolor{brickred}{2},\textcolor{brickred}{2},C/2]$ patches, while we \textit{grid} the other half using a $\textcolor{dark_green}{2}\times \textcolor{dark_green}{2}$ grid in the \textcolor{dark_green}{global branch}. We only apply the gMLP block~\cite{liu2021pay} (illustrated in the right \textcolor{blueish}{gMLP Block}) on \textit{a single axis} of each branch - the \textcolor{brickred}{2nd} axis for the local branch and the \textcolor{dark_green}{1st} axis for the global branch, while shared along the other spatial dimensions. The gMLP operators, which run in parallel, correspond to local and global (dilated) attended regions, as illustrated with different colors (\ie, the same color are spatially mixed using the gMLP operator). Our proposed block expresses both global and local receptive fields on arbitrary input resolutions.}
 \label{fig:multi-axis-gated-mlp-block}
 \vspace{-3mm}
\end{figure*}

Our work is inspired by the multi-axis blocked self-attention proposed in~\cite{zhao2021improved}, which performs attention on more than a single axis.
The attentions performed on two axes on blocked images correspond to two forms of sparse self-attention, namely regional and dilated attention.
Despite capturing local and global information in parallel, this module cannot accommodate image restoration or enhancement tasks where the test images are often of arbitrary sizes.

We improve the `\textit{multi-axis}' concept for image processing tasks, by building a (split-head) multi-axis gated MLP block (MAB), as shown in \cref{fig:multi-axis-gated-mlp-block}.
Instead of applying multi-axis attention in a single layer~\cite{zhao2021improved}, we split in half the heads first, each being partitioned independently.
In the \textcolor{brickred}{local branch}, the half head of a feature of size $(H,W,C/2)$ is \textit{blocked} into a tensor of shape $(\frac{H}{b}\times\frac{W}{b},b\times b,C/2)$, representing partitioning into non-overlapping windows each with size of $(b\times b)$; in the \textcolor{dark_green}{global branch}, the other half head is \textit{gridded} into the shape $(d\times d,\frac{H}{{d}}\times\frac{W}{d},C/2)$ using a fixed $(d\times d)$ grid, with each window having size $(\frac{H}{{d}}\times\frac{W}{{d}})$.
For visualization, we set $b=\textcolor{brickred}{2},d=\textcolor{dark_green}{2}$ in \cref{fig:multi-axis-gated-mlp-block}. 
To make it \textit{fully-convolutional}, we only apply the gated MLP (gMLP) block~\cite{liu2021pay} on \textit{a single axis} of each branch -- the \textcolor{brickred}{$2$nd axis} for the local branch and the \textcolor{dark_green}{$1$st axis} for the global branch -- while sharing parameters on the other spatial axes.
Intuively, applying multi-axis gMLPs in parallel correspond to local and global (dilated) mixing of spatial information, respectively.
Finally, the processed heads are concatenated and projected to reduce the number of channels, which are further combined using the long skip-connection from the input.
It is worth noting that this approach provides an advantage for our model over methods that process fixed-size image patches~\cite{chen2021pre} by avoiding patch boundary artifacts.


\noindent\textbf{Complexity analysis.}
The computational complexity of our proposed Multi-Axis gMLP block (MAB) is:
\vspace{-1mm}
\begin{equation}
\label{eq:complexity}
\small
\Omega(\mathrm{MAB})=\underbrace{b^2HWC}_{\textcolor{dark_green}{\mathrm{Global\  gMLP}}}+\underbrace{d^2HWC}_{\textcolor{brickred}{\mathrm{Local\  gMLP}}}+\underbrace{10HWC^2}_{\textcolor{blueish}{\mathrm{Dense\ layers}}},
\end{equation}
which is \textit{linear} with respect to image size $HW$, while other global models like ViT, Mixer, and gMLP are \textit{quadratic}.

\noindent\textbf{Universality of the multi-axis approach.} Our proposed parallel multi-axis module (\cref{fig:multi-axis-gated-mlp-block}) presents a principled way to apply 1D operators on 2D images in a scalable manner.
It also allows for significant flexibility and universality.
For example, a straightforward replacement of a gMLP with a spatial MLP~\cite{tolstikhin2021mlp}, self-attention~\cite{dosovitskiy2021an}, or even Fourier Transform~\cite{rao2021global,lee2021fnet} leads to a family of MAXIM variants (see~\cref{ssec:ablation}D), all sharing globality and fully-convolutionality.
It is also easily extensible to \textit{any} future 1D operator that may be defined on, \eg, Language models.


\subsection{Cross Gating MLP Block}
\label{ssec:cross-gating-block-mlp}

A common improvement over UNet is to leverage contextual features to selectively \textit{gate} feature propagation in skip-connections~\cite{oktay2018attention,petit2021u}, which is often achieved by using cross-attention~\cite{vaswani2017attention,chen2021crossvit}.
Here we build an effective alternative, namely cross-gating block (CGB, \cref{fig:MAXIM-backbone}c), as an extension of MAB (\cref{ssec:multi-axis-gating-mlp}) which can only process a single feature.
CGB can be regarded as a more general conditioning layer that interacts with multiple features~\cite{vaswani2017attention,perez2018film,chen2021crossvit}.
We follow similar design patterns as those used in MAB.

To be more specific, let $\mathbf{X},\mathbf{Y}$ be two input features, and $\mathbf{X}_1,\mathbf{Y}_1\in\mathbb{R}^{H\times W\times C}$ be the features projected after the first \texttt{Dense} layers in \cref{fig:MAXIM-backbone}c. Input projections are then applied:
\begin{equation}
\label{eq:cross-gating-block-input-projection}
\small
\vspace{-1mm}
\mathbf{X}_2=\sigma(\mathbf{W}_1 \mathsf{LN}(\mathbf{X}_{1}))\ ,\ \ \mathbf{Y}_2=\sigma(\mathbf{W}_2 \mathsf{LN}(\mathbf{Y}_{1}))
\end{equation}
where $\sigma$ is the $\mathsf{GELU}$ activation~\cite{hendrycks2016gaussian}, $\mathsf{LN}$ is Layer Normalization~\cite{ba2016layer}, and $\mathbf{W}_1,\mathbf{W}_2$ are \texttt{MLP} projection matrices.
The multi-axis blocked gating weights are computed from $\mathbf{X}_2,\mathbf{Y}_2$, respectively, but applied \textit{reciprocally}:
\begin{equation}
\label{eq:cross-gating-block-cross-gating}
\small
\hat{\mathbf{X}}=\mathbf{X}_2\odot \mathrm{G}(\mathbf{Y}_2)\ ,\ \  \hat{\mathbf{Y}}=\mathbf{Y}_2\odot \mathrm{G}(\mathbf{X}_2)
\end{equation}
where $\odot$ represents element-wise multiplication, and the function $\mathrm{G}(\cdot)$ extracts multi-axis cross gating weights from the input using our proposed multi-axis approach (\cref{ssec:multi-axis-gating-mlp}):
\begin{equation}
\label{eq:cross-gating-block-gating-function}
\small
\mathrm{G}(\mathbf{x})=\mathbf{W}_5([\mathbf{W}_3\mathsf{Block}_{b}(\mathbf{z_1}), \mathbf{W}_4\mathsf{Grid}_{d}(\mathbf{z_2})]) 
\end{equation}
where $[\cdot,\cdot]$ denotes concatenation.
Here $(\mathbf{z_1},\mathbf{z_2})$ are two independent heads split from $\mathbf{z}$ along the channel dimension, where $\mathbf{z}$ represents the projected features $\mathbf{x}$ after activation:
\begin{equation}
\small
\label{eq:cross-gating-block-gating-function-input-proj}
[\mathbf{z_1},\mathbf{z_2}]=\mathbf{z}=\sigma(\mathbf{W}_6 \mathsf{LN}(\mathbf{x})),
\end{equation}
and $\mathbf{W}_3,\mathbf{W}_4$ are spatial projection matrices applied on the \textcolor{brickred}{2nd} and \textcolor{dark_green}{1st} axis of the blocked/gridded features having fixed window size $b\times b$ ($\mathsf{Block}_{b}$), and fixed grid size of $d\times d$ ($\mathsf{Grid}_{d}$), respectively.
Finally, we adopt residual connection from the inputs, following an output channel-projection that maintains the same channel dimensions as the inputs ($\mathbf{X}_1,\mathbf{Y}_1$), using projection matrices $\mathbf{W}_7$, $\mathbf{W}_8$, denoted by
\begin{equation}
\label{eq:cross-gating-block-output-proj}
\small
\mathbf{X}_{3}=\mathbf{X}_{1}+\mathbf{W}_7\hat{\mathbf{X}}\ ,\ \  \mathbf{Y}_{3}=\mathbf{Y}_{1}+\mathbf{W}_8\hat{\mathbf{Y}}.
\end{equation}
The complexity of CGB is also tightly-bounded by \cref{eq:complexity}.

\subsection{Multi-Stage Multi-Scale Framework}
\label{ssec:multi-stage-multi-scale-framework}

We further adopt a multi-stage framework because we find it more effective, as compared to scaling up the model width or height (see ablation \cref{ssec:ablation}A).
We deem full resolution processing~\cite{park2020multi,ren2019progressive,chen2021hinet} a better approach than a multi-patch hierarchy~\cite{suin2020spatially,zamir2021multi,zhang2019deep}, since the latter would potentially induce boundary effects across patches. 
To impose stronger supervision, we apply a multi-scale approach~\cite{li2019rethinking,chen2018cascaded,cho2021rethinking} at each stage to help the network learn. We leverage the supervised attention module~\cite{zamir2021multi} to propagate attentive feature maps progressively along the stages.
We leverage our proposed cross-gating block (\cref{ssec:cross-gating-block-mlp}) for any cross-stage feature fusion. See Appendix for details.


\begin{table}[!t]
\centering
\scriptsize
\setlength{\tabcolsep}{3pt}
\renewcommand{1.12}{1.1}
\begin{tabular}{l|cc|cc||cc}
& \multicolumn{2}{c}{SIDD~\cite{abdelhamed2018high}} & \multicolumn{2}{c}{DND~\cite{plotz2017benchmarking}} & \multicolumn{2}{c}{Average} \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$ &  PSNR$\uparrow$ & SSIM$\uparrow$ \\
\toprule
DnCNN~\cite{zhang2017beyond} & 23.66 & 0.583 & 32.43 & 0.790 & 28.04 & 0.686 \\
MLP~\cite{burger2012image} & 24.71 & 0.641 & 34.23 & 0.833 & 29.47 & 0.737 \\
BM3D~\cite{dabov2007image} & 35.65 & 0.685 & 34.51 & 0.851 & 35.08 & 0.768 \\
CBDNet\textcolor{red}{*}~\cite{guo2019toward} & 30.78 & 0.801 & 38.06 & 0.942 & 34.42 & 0.872  \\
RIDNet\textcolor{red}{*}~\cite{anwar2019real} & 38.71 & 0.951 & 39.26 & 0.953 & 38.99 & 0.952 \\
AINDNet\textcolor{red}{*}~\cite{kim2020transfer} & 38.95 & 0.952 & 39.37 & 0.951 & 39.16 & 0.952 \\
VDN~\cite{yue2019variational} & 39.28 & 0.956 & 39.38 & 0.952  & 39.33 & 0.954 \\
SADNet\textcolor{red}{*}~\cite{chang2020spatial} & 39.46 & 0.957 & 39.59 & 0.952 & 39.53 & 0.955 \\
CycleISP\textcolor{red}{*}~\cite{zamir2020cycleisp} & 39.52 & 0.957 & 39.56 & \textbf{0.956} & 39.54 & \underline{0.957} \\
MIRNet~\cite{zamir2020learning} & \underline{39.72} & \underline{0.959} & \textbf{39.88} & \textbf{0.956} & \underline{39.80} & \textbf{0.958} \\
MPRNet~\cite{zamir2021multi} & {39.71} & {0.958} & {39.80} & \underline{0.954} & 39.76 & 0.956 \\
\midrule
MAXIM-3S & \textbf{39.96} & \textbf{0.960} & \underline{39.84} & \underline{0.954} & \textbf{39.90} & \underline{0.957} \\
\bottomrule
\end{tabular}
\caption{Denoising results. Our model is only trained on SIDD~\cite{abdelhamed2018high} and evaluated on SIDD~\cite{abdelhamed2018high} and DND~\cite{plotz2017benchmarking}, where \textcolor{red}{*} denotes methods using additional training data.}
\label{tab:denoising}
\end{table}

Formally, given an input image $\mathbf{I}\in \mathbb{R}^{H\times W\times3}$, we first extract its multi-scale variants by downscaling: $\mathbf{I}_n,\ n=1,...,N$. MAXIM predicts multi-scale restored outputs at each stage $s$ of $S$ stages, yielding a total of $S\times N$ outputs: $\mathbf{R}_{s,n}$. Despite being multi-stage, MAXIM is trained \textit{end-to-end} with losses accumulating across stages and scales:
\begin{equation}
\label{eq:loss}
\small
\mathcal{L}=\sum_{s=1}^S\sum_{n=1}^N [\mathcal{L}_{char}(\mathbf{R}_{s,n}, \mathbf{T}_{n})+\lambda\mathcal{L}_{freq}(\mathbf{R}_{s,n}, \mathbf{T}_{n})],
\end{equation}
where $\mathbf{T}_n$ denotes (bilinearly-rescaled) multi-scale target images, and $\mathcal{L}_{char}$ is the Charbonnier loss~\cite{zamir2021multi}:
\begin{equation}
\label{eq:char-loss}
\small
\mathcal{L}_{char}(\mathbf{R},\mathbf{T})=\sqrt{\|\mathbf{R}-\mathbf{T}\|^2+\epsilon^2},
\end{equation}
where we set $\epsilon=10^{-3}$. $\mathcal{L}_{freq}$ is the frequency reconstruction loss that enforces high-frequency details~\cite{jiang2021focal,cho2021rethinking}:
\begin{equation}
\label{eq:freq-loss}
\small
\mathcal{L}_{freq}(\mathbf{R},\mathbf{T})=\|\mathcal{F}(\mathbf{R})-\mathcal{F}(\mathbf{T})\|_1
\end{equation}
where $\mathcal{F}(\cdot)$ represents the 2D Fast Fourier Transform. We used $\lambda=0.1$ as the weighting factor in all experiments.


\section{Experiments}
\label{sec:experiments}

We aim at building a generic backbone for a broad spectrum of image processing tasks. 
Thus, we evaluated MAXIM on five different tasks: (1) denoising, (2) deblurring, (3) deraining, (4) dehazing, and (5) enhancement (retouching) on \textbf{17} different datasets.
More comprehensive results and visualizations can be found in Appendix.

\subsection{Experimental Setup}
\label{ssec:experiment-setup}

\noindent\textbf{Datasets and metrics.} We measured PSNR and SSIM~\cite{wang2004image} metrics between ground truth and predicted images to make quantitative comparisons. We used SIDD~\cite{abdelhamed2018high} and DND~\cite{plotz2017benchmarking} for denoising, GoPro~\cite{nah2017deep}, HIDE~\cite{shen2019human}, and RealBlur~\cite{rim2020real} for debluring, a combined dataset Rain13k used in~\cite{zamir2021multi} for deraining. The RESIDE~\cite{li2019benchmarking} is used for dehazing, while Five-K\cite{bychkovsky2011learning} and LOL~\cite{wei2018deep} are evaluated for enhancement.

\begin{figure*}[!ht]
\centering
\footnotesize
\def-2.5pt{1pt}
\def0.99{0.99}
\def0.12\textwidth{0.12\textwidth}
\setlength{\tabcolsep}{1pt}
\begin{tabular}{C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}}
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Denoising_Visuals_SIDD.pdf}
} \\
Input & Target & VDN~\cite{yue2019variational} & DANet~\cite{yue2020dual} & MIRNet~\cite{zamir2020learning} & CycleISP~\cite{zamir2020cycleisp} & MPRNet~\cite{zamir2021multi} & MAXIM-3S \\
\end{tabular}
\caption{Denoising comparisons. The example from SIDD~\cite{abdelhamed2018high} shows that our method produces cleaner denoising results.}
\label{fig:denoising-visuals}
\end{figure*}

\begin{table}[!t]
\centering
\scriptsize
\setlength{\tabcolsep}{1.6pt}
\renewcommand{1.12}{1.1}
\begin{tabular}{l|cc|cc||cc}
& \multicolumn{2}{c}{GoPro~\cite{nah2017deep}} & \multicolumn{2}{c}{HIDE~\cite{shen2019human}} & \multicolumn{2}{c}{Average} \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$ \\
\toprule
DeblurGAN~\cite{kupyn2018deblurgan} & 28.70 & 0.858 & 24.51 & 0.871 & 26.61 & 0.865 \\
Nah \etal~\cite{nah2017deep} & 29.08 & 0.914 & 25.73 & 0.874  & 27.41 & 0.894 \\
Zhang \etal~\cite{zhang2018dynamic} & 29.19 & 0.931 & - & -  & - & - \\
DeblurGAN-v2~\cite{kupyn2019deblurgan} & 29.55 & 0.934 & 26.61 & 0.875 & 28.08  & 0.905 \\
SRN~\cite{tao2018scale} & 30.26 & 0.934 & 28.36 & 0.915  & 29.31 &  0.925 \\
Shen \etal~\cite{shen2019human} & - & - & 28.89 & 0.930 & - & - \\
Gao \etal~\cite{gao2019dynamic} & 30.90 & 0.935 & 29.11 & 0.913 &  30.01 & 0.924 \\
DBGAN~\cite{zhang2020deblurring} & 31.10 & 0.942 & 28.94 & 0.915 &  30.02 & 0.929 \\
MT-RNN~\cite{park2020multi} & 31.15 & 0.945 & 29.15 & 0.918 & 30.15 & 0.932\\
DMPHN~\cite{zhang2019deep} & 31.20 & 0.940 & 29.09 & 0.924  & 30.15 & 0.932 \\
Suin \etal~\cite{suin2020spatially} & 31.85 & 0.948 & 29.98 & 0.930 & 30.92 & 0.939 \\
MPRNet~\cite{zamir2021multi} & {32.66} & \underline{0.959} & \underline{30.96} & \underline{0.939} & \underline{31.81} & \underline{0.949} \\
Pretrained-IPT~\cite{chen2021pre} & 32.58 & - & - & -  & - & - \\
MIMO-UNet+~\cite{cho2021rethinking} & 32.45 & 0.957 & 29.99 & 0.930 & 31.22 & 0.944 \\
HINet~\cite{chen2021hinet} & \underline{32.71} & \underline{0.959} & {30.32} & {0.932} & 31.52 &  0.946 \\
\midrule
MAXIM-3S & \textbf{32.86} & \textbf{0.961} & \textbf{32.83} & \textbf{0.956} & \textbf{32.85} & \textbf{0.959} \\
\bottomrule
\end{tabular}
\caption{Deblurring results. Our model is trained on GoPro~\cite{nah2017deep} and evaluated on the GoPro and the HIDE dataset~\cite{shen2019human}.}
\label{tab:deblurring-gopro-hide}
\end{table}

\begin{table}[!t]
\centering
\scriptsize
\setlength{\tabcolsep}{1.6pt}
\renewcommand{1.12}{1.1}
\begin{tabular}{l|cc|cc||cc}
& \multicolumn{2}{c}{RealBlur-R~\cite{rim2020real}} & \multicolumn{2}{c}{RealBlur-J~\cite{rim2020real}} & \multicolumn{2}{c}{Average} \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$\\
\toprule
Hu \etal~\cite{hu2014deblurring} & 33.67 & 0.916 & 26.41 & 0.803 & 30.04  & 0.860 \\
Nah \etal~\cite{nah2017deep} & 32.51 & 0.841 & 27.87 & 0.827 & 30.19 & 0.834 \\
DeblurGAN~\cite{kupyn2018deblurgan} & 33.79 & 0.903 & 27.97 & 0.834 & 30.88 & 0.869 \\
Pan \etal~\cite{pan2016blind} & 34.01 & 0.916 & 27.22 & 0.790 & 30.62 & 0.853 \\
Xu \etal~\cite{xu2013unnatural} & 34.46 & 0.937 & 27.14 & 0.830 & 30.8 & 0.884 \\
DeblurGAN-v2~\cite{kupyn2019deblurgan} & 35.26 & 0.944 & 28.70 & 0.866 & 31.98 & 0.905 \\
Zhang \etal~\cite{zhang2018dynamic} & 35.48 & 0.947 & 27.80 & 0.847 & 31.64  & 0.897\\
SRN~\cite{tao2018scale} & 35.66 & 0.947 & 28.56 & 0.867 & 32.11 & 0.907 \\
DMPHN~\cite{zhang2019deep} &  35.70 & \underline{0.948} & 28.42 & 0.860 & 32.06 & 0.904 \\
MPRNet~\cite{zamir2021multi} & \textbf{35.99} & \textbf{0.952} & \underline{28.70} & \underline{0.873} & \textbf{32.35} & \textbf{0.913} \\
\hline
MAXIM-3S & \underline{35.78} & 0.947 & \textbf{28.83} & \textbf{0.875} & \underline{32.31} &  \underline{0.911} \\
\midrule[0.05pt]\midrule
\textcolor{red}{$^\dagger$}DeblurGAN-v2 & 36.44 & 0.935 & 29.69 & 0.870 & 33.07 & 0.903 \\
\textcolor{red}{$^\dagger$}SRN~\cite{tao2018scale} & 38.65 & \underline{0.965} & 31.38 & 0.909 & 35.02 & 0.937 \\
\textcolor{red}{$^\dagger$}MPRNet~\cite{zamir2021multi} & \underline{39.31} & \textbf{0.972} & 31.76 & \underline{0.922} & \underline{35.54} & \underline{0.947} \\
\textcolor{red}{$^\dagger$}MIMO-UNet+~\cite{cho2021rethinking} & - & - & \underline{32.05} & 0.921 & - & - \\
\hline
\textcolor{red}{$^\dagger$}MAXIM-3S & \textbf{39.45} & 0.962 & \textbf{32.84} & \textbf{0.935} & \textbf{36.15} & \textbf{0.949} \\
\bottomrule
\end{tabular}
\caption{Deblurring results on RealBlur~\cite{rim2020real}. \textcolor{red}{$^\dagger$} denotes methods that are trained on RealBlur, while those without \textcolor{red}{$^\dagger$} indicate methods trained only on GoPro.}
\label{tab:deblurring-realblur}
\end{table}

\begin{figure*}[!t]
\centering
\footnotesize
\def-2.5pt{-2.5pt}
\def0.99{0.99}
\def0.12\textwidth{0.12\textwidth}
\setlength{\tabcolsep}{1pt}
\begin{tabular}{C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}}
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Deblurring_Visuals.pdf}
} \\
Input & Target & DMPHN~\cite{zhang2019deep} & Suin \etal~\cite{suin2020spatially} & MPRNet~\cite{zamir2021multi} & HINet~\cite{chen2021hinet} & MIMO-UNet~\cite{cho2021rethinking} & MAXIM-3S
\end{tabular}
\caption{Deblurring comparisons. The top row shows an example from GoPro~\cite{nah2017deep} while the second row shows one from HIDE~\cite{shen2019human}.}
\label{fig:deblurring-visuals}
\vspace{-3mm}
\end{figure*}

\noindent\textbf{Training details.} Our proposed MAXIM model is end-to-end trainable and requires neither large-scale pretraining nor progressive training. The network is trained on $256\!\times\!256$ random-cropped patches. We train different iterations for each task. We used random horizontal and vertical flips, $90^{\circ}$ rotation, and MixUp~\cite{zhang2017mixup} with probability $0.5$ for data augmentation. We used the Adam optimizer~\cite{kingma2014adam} with an initial learning rate of $2\!\times\!10^{-4}$, which are steadily decreased to $10^{-7}$ with the cosine annealing decay~\cite{loshchilov2016sgdr}. When testing, we padded the input images to be a multiplier of $64\!\times\!64$ using symmetric padding on both sides. After inference, we cropped the padded image back to original size. 

\noindent\textbf{Architectural configuration.}
We designed two MAXIM variants: a two-stage model called MAXIM-\textcolor{r}{2S}, and a three-stage model, MAXIM-\textcolor{b}{3S}, for different tasks. We start with $32$ initial channels for feature extraction, with 3 downsampling layers, where the features contract from $256^2\times 32$, $128^2\times64$, $64^2\times 128$, to $32^2\times 256$ processed by \textit{two} \textcolor{brickred}{Bottlenecks} (\cref{fig:MAXIM-backbone}a), then symmetrically expanded back to full resolution. The number of parameters and required FLOPs of MAXIM-\textcolor{r}{2S} and MAXIM-\textcolor{b}{3S}, when applied on a $256\times 256$ image are shown in the last two rows of \cref{tab:ablation}A.

\subsection{Main Results}
\label{ssec:main-results}

\noindent\textbf{Denoising.}
We report in \cref{tab:denoising} numerical comparisons on the SIDD~\cite{abdelhamed2018high} and DND~\cite{plotz2017benchmarking} datasets. As may be seen, our method outperformed previous SOTA techniques, \eg, MIRNet~\cite{zamir2020learning} by \textbf{0.24 dB} of PSNR on SIDD while obtaining competitive PSNR (39.84 dB) on DND. \cref{fig:denoising-visuals} shows visual results on SIDD. Our method clearly removes real noise while maintaining fine details, yielding visually pleasant results to the other methods.


\noindent\textbf{Deblurring.} \cref{tab:deblurring-gopro-hide} shows the quantitative comparison of MAXIM-\textcolor{b}{3S} against SOTA deblurring methods on two synthetic blur datasets: GoPro~\cite{nah2017deep} and HIDE~\cite{shen2019human}. Our method achieves \textbf{0.15} dB gain in PSNR over the previous best model HINet~\cite{chen2021hinet}. It is notable that the GoPro-trained MAXIM-\textcolor{b}{3S} model generalizes extremely well on the HIDE dataset, setting new SOTA PSNR values: \textbf{32.83} dB. We also evaluated on real-world blurry images from RealBlur~\cite{rim2020real} under two settings: (1) directly applied the GoPro-trained model on RealBlur, and (2) fine-tuned the model on RealBlur. Under setting (1), MAXIM-\textcolor{b}{3S} ranked \textit{first} on RealBlur-J subset while obtaining the top two performance on RealBlur-R. \cref{fig:deblurring-visuals} shows visual comparisons of the evaluated models on GoPro~\cite{nah2017deep}, HIDE~\cite{shen2019human} and RealBlur~\cite{rim2020real}, respectively. It may be observed that our model recovers \textbf{text} extremely well, which may be attributed to the use of multi-axis MLP module within each block that globally aggregates repeated patterns across various scales.

\begin{figure*}[!t]
\centering
\footnotesize
\def-2.5pt{-2.5pt}
\def0.99{0.99}
\def0.12\textwidth{0.12\textwidth}
\setlength{\tabcolsep}{1pt}
\begin{tabular}{C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}}
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Deraining_Visuals.pdf}
} \\
Input & Target & RESCAN~\cite{li2018recurrent} & PreNet~\cite{ren2019progressive} & MSPFN~\cite{jiang2020multi} & MPRNet~\cite{zamir2021multi} & HINet~\cite{chen2021hinet} & MAXIM-2S \\
\end{tabular}
\caption{Deraining comparisons. The top and bottom rows present examples from Rain100L~\cite{yang2017deep} and Test100~\cite{zhang2019image}, respectively, demonstrating the ability of MAXIM to remove rain streaks while recovering more details, hence yielding more visually pleasant results.}
\label{fig:deraining-visuals}
\end{figure*}

\begin{figure*}[!t]
\centering
\footnotesize
\def-2.5pt{-2.5pt}
\def0.99{0.99}
\def0.12\textwidth{0.12\textwidth}
\setlength{\tabcolsep}{1pt}
\begin{tabular}{C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}}
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Dehazing_Visuals_Indoor.pdf}
} \\
Input & Target & GCANet~\cite{chen2019gated} & GridDehaze~\cite{liu2019griddehazenet} & DuRN~\cite{liu2019dual} & MSBDN~\cite{dong2020multi} & FFA-Net~\cite{qin2020ffa} & MAXIM-2S \\
\end{tabular}
\caption{Dehazing comparisons. The top and bottom rows exemplify visual results from the SOTS indoor and outdoor sets~\cite{li2019benchmarking}.}
\label{fig:dehazing-visuals}
\end{figure*}

\begin{table*}[!t]
\scriptsize
\centering
\setlength{\tabcolsep}{1.25pt}
\renewcommand{1.12}{1.1}
\begin{tabular}{@{}lr@{}}
\begin{minipage}[t]{.68\linewidth}
\begin{tabular}{@{}l|cc|cc|cc|cc|cc||cc@{}}
& \multicolumn{2}{c}{Rain100L~\cite{yang2017deep}} & \multicolumn{2}{c}{Rain100H~\cite{yang2017deep}} & 
\multicolumn{2}{c}{Test100~\cite{zhang2019image}} & 
\multicolumn{2}{c}{Test1200~\cite{zhang2018density}} & 
\multicolumn{2}{c}{Test2800~\cite{fu2017removing}} & \multicolumn{2}{c}{Average}  \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$  & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$  & PSNR$\uparrow$ & SSIM$\uparrow$  & PSNR$\uparrow$ & SSIM$\uparrow$ \\
\toprule
DerainNet~\cite{fu2017clearing} & 27.03 & 0.884 & 14.92 & 0.592 & 22.77 & 0.810 & 23.38 & 0.835 & 24.31 & 0.861 & 22.48 & 0.796  \\
SEMI~\cite{wei2019semi} & 25.03 & 0.842 & 16.56 & 0.486 & 22.35 & 0.788 & 26.05 & 0.822 & 24.43 & 0.782 & 22.88 & 0.744 \\
DIDMDN~\cite{zhang2018density} & 25.23 & 0.741 & 17.35 & 0.524 & 22.56 & 0.818 & 29.65 & 0.901 & 28.13 & 0.867 & 24.58 & 0.770 \\
UMRL~\cite{yasarla2019uncertainty} & 29.18 & 0.923 & 26.01 & 0.832 & 24.41 & 0.829 & 30.55 & 0.910 & 29.97 & 0.905 &   28.02 &  0.880\\ 
RESCAN~\cite{li2018recurrent} & 29.80 & 0.881 & 26.36 & 0.786 & 25.00 & 0.835 & 30.51 & 0.882 & 31.29 & 0.904 & 28.59 & 0.857 \\
PreNet~\cite{ren2019progressive} & 32.44 & 0.950 & 26.77 & 0.858 & 24.81 & 0.851 & 31.36 & 0.911 & 31.75 & 0.916 & 29.42 & 0.897 \\
MSPFN~\cite{jiang2020multi} & 32.40 & 0.933 & 28.66 & 0.860 & 27.50 & 0.876 & 32.39 & 0.916 & 32.82 & 0.930 & 30.75 & 0.903 \\
MPRNet~\cite{zamir2021multi} & 36.40 & 0.965 & 30.41 & 0.890 & \underline{30.27} & 0.897 & \underline{32.91} & 0.916 & 33.64 & 0.938 & 32.73 & 0.921 \\
HINet~\cite{chen2021hinet} & \underline{37.20} & \underline{0.969} & \underline{30.63} & \underline{0.893} & 30.26 & \underline{0.905} & \textbf{33.01} & \underline{0.918} & \textbf{33.87} & \underline{0.940} & \underline{33.00} & \underline{0.925} \\
\midrule
MAXIM-2S & \textbf{38.06} &	\textbf{0.977} & \textbf{30.81} & \textbf{0.903} & \textbf{31.17} & \textbf{0.922}	& \underline{32.37} & \textbf{0.922}	& \underline{33.80} & \textbf{0.943} & \textbf{33.24} & \textbf{0.933} \\
\bottomrule
\end{tabular}
\caption{Deraining comparisons. Our method consistently yields better quality metrics with respect to both PSNR or SSIM on all the tested datasets: Rain100L~\cite{yang2017deep}, Rain100H~\cite{yang2017deep}, Test100~\cite{zhang2019image}, Test1200~\cite{zhang2018density}, Test2800~\cite{fu2017removing}}
\label{tab:deraining}
\end{minipage} &

\hspace{6pt}

\begin{minipage}[t]{.29\linewidth}
\begin{tabular}{@{}l|cc|cc@{}}
& \multicolumn{2}{c}{SOTS-Indoor} & \multicolumn{2}{c}{SOTS-Outdoor} \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & PSNR$\uparrow$ & SSIM$\uparrow$ \\
\toprule
DehazeNet~\cite{cai2016dehazenet} & 21.14 & 0.847 & 22.46 & 0.851 \\
GFN~\cite{Ren-CVPR-2018} & 22.30 & 0 880 & 21.55 & 0.844 \\
GCANet~\cite{chen2019gated} & 30.23 & 0.959 & 19.98  & 0.704 \\
GridDehaze~\cite{liu2019griddehazenet} & 32.14 & 0.983 & 30.86 & 0.981 \\
GMAN~\cite{liu2019single} & 27.93 & 0.896 & 28.47 & 0.944 \\
MSBDN~\cite{dong2020multi} & 33.79 & 0.984 & 23.36 & 0.875 \\
DuRN~\cite{liu2019dual} & 32.12 & 0.980 & 24.47 & 0.839 \\
FFA-Net~\cite{qin2020ffa} & 36.39 & 0.989 & \underline{33.57} & \underline{0.984} \\
AECR-Net~\cite{wu2021contrastive} & \underline{37.17} & \underline{0.990} & - & - \\
\midrule
MAXIM-2S & \textbf{38.11} & \textbf{0.991} & \textbf{34.19} & \textbf{0.985}  \\
\bottomrule
\end{tabular}
\caption{Dehazing comparisons. Our model achieved the best results on both indoor and outdoor scenes.}
\label{tab:dehazing}
\end{minipage} \\
\end{tabular}
\vspace{-3mm}
\end{table*}


\begin{figure*}[!t]
\centering
\footnotesize
\def-2.5pt{-2.5pt}
\def0.99{0.99}
\def0.12\textwidth{0.12\textwidth}
\setlength{\tabcolsep}{1pt}
\begin{tabular}{C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}C{0.12\textwidth}}
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Enhancement_Visuals_FiveK.pdf}
} \\
Input & Target & CycleGAN~\cite{zhu2017unpaired} & Exposure~\cite{hu2018exposure} & DPE~\cite{chen2018deep} & EnlightenGAN & UEGAN~\cite{ni2020towards} & MAXIM-2S \\
\multicolumn{8}{c}{
\includegraphics[width=0.99\linewidth]{figures/Enhancement_Visuals_LOL.pdf}
} \\
Input & Target & Retinex~\cite{wei2018deep} & GLAD~\cite{wang2018gladnet} & KinD~\cite{zhang2019kindling} & EnlightenGAN & MIRNet~\cite{zamir2020learning} & MAXIM-2S \\
\end{tabular}
\caption{Retouching and low-light enhancement comparisons. The top row shows an example from the MIT-Adobe FiveK dataset~\cite{bychkovsky2011learning}, while the bottom row exemplifies a comparison from LOL~\cite{wei2018deep}. Our model generated variegated and more naturalistic colors (top) for retouching, while achieving clearer and brighter visual enhancements in the bottom example.}
\label{fig:enhancement-visuals}
\vspace{-3mm}
\end{figure*}



\noindent\textbf{Deraining.} Following previous work~\cite{jiang2020multi,zamir2021multi}, we computed the performance metrics using the Y channel (in YCbCr color space). \cref{tab:deraining} shows quantitative comparisons with previous methods. As may be seen, our model improved over the SOTA performances on all datasets. The average PSNR gain of our model over the previous best model HINet~\cite{chen2021hinet} is \textbf{0.24} dB. We demonstrate some challenging examples in \cref{fig:deraining-visuals}, which demonstrates that our method consistently delivered faithfully recovered images without introducing any noticeable visual  artifacts. 

\noindent\textbf{Dehazing.} We report our comparisons against SOTA models in \cref{tab:dehazing}. Our model surpassed the previous best model by \textbf{0.94} dB and \textbf{0.62} dB of PSNR on the SOTS~\cite{li2019benchmarking} indoor and outdoor sets. \cref{fig:dehazing-visuals} shows that our model recovered images of better quality on both flat regions as well as textures, while achieving a harmonious global tone.

\noindent\textbf{Enhancement / Retouching.} As \cref{tab:enhancement} illustrates, our model achieved the best PSNR and SSIM values on FiveK~\cite{bychkovsky2011learning} and LOL~\cite{wei2018deep}, respectively. As the top row of \cref{fig:enhancement-visuals} suggests, MAXIM recovered diverse naturalistic colors as compared to other techniques. Regarding the bottom example, while MIRNet~\cite{zamir2020learning} obtained a higher PSNR, we consistently observed that our model attains visually better quality with sharper details and less noise. Moreover, the far more perceptually relevant SSIM index indicates a significant advantage of MAXIM-\textcolor{r}{2S} relative to MIRNet.

\noindent\textbf{Other benchmarks.} Due to space limitations, we detail the outcomes of our experiments on the REDS deblurring~\cite{Nah_2021_CVPR} and the Raindrop removal task~\cite{qian2018attentive} in Appendix.

\subsection{Ablation}
\label{ssec:ablation}

We conduct extensive ablation studies to validate the proposed multi-axis gated MLP block, cross-gating block, and multi-stage multi-scale architecture. The evaluations were performed on the GoPro dataset~\cite{nah2017deep} trained on image patches of size $256\times 256$ for $10^6$ iterations. We used the MAXIM-\textcolor{r}{2S} model as the test-bed for Ablation-A and -B.


\noindent\textbf{A. Individual components.}
We conducted an ablation by progressively adding (1) inter-stage cross-gating blocks (CGB\textsubscript{IS}), (2) a supervised attention module (SAM), (3) cross-stage cross-gating blocks (CGB\textsubscript{CS}, and (4) the multi-scale supervision (MS-Sp). \cref{tab:ablation}A indicates a PSNR gain of 0.25, 0.63, 0.36, 0.26 dB for each respective component.

\begin{table}[!h]
\centering
\scriptsize
\setlength{\tabcolsep}{4pt}
\renewcommand{1.12}{1.1}
\begin{tabular}{l|cc||l|cc}
& \multicolumn{2}{c}{FiveK~\cite{bychkovsky2011learning}} & & \multicolumn{2}{c}{LOL~\cite{wei2018deep}} \\ 
Method & PSNR$\uparrow$ & SSIM$\uparrow$ & Method & PSNR$\uparrow$ & SSIM$\uparrow$ \\
\toprule
CycleGAN~\cite{zhu2017unpaired} & 18.23 & 0.835 &  Retinex~\cite{wei2018deep} & 16.77 & 0.559 \\
Exposure~\cite{hu2018exposure} & 22.35 & 0.861 & GLAD~\cite{wang2018gladnet} & 19.71 & 0.703 \\
EnlightenGAN & 17.74 & 0.828 & EnlightenGAN & 17.48 & 0.657 \\
DPE~\cite{chen2018deep} & 24.08 & 0.922 & KinD~\cite{zhang2019kindling} & 20.37 & 0.804 \\
UEGAN~\cite{ni2020towards} & \underline{25.00} & \underline{0.929} & MIRNet~\cite{zamir2020learning} & \textbf{24.14} & \underline{0.830} \\
\midrule
MAXIM-2S & \textbf{26.15} & \textbf{0.945} & MAXIM-2S & \underline{23.43} & \textbf{0.863}    \\
\bottomrule
\end{tabular}
\caption{Enhancement results on FiveK~\cite{bychkovsky2011learning} and LOL~\cite{wei2018deep}.}
\label{tab:enhancement}
\end{table}

\noindent\textbf{B. Effects of multi-axis approach.}
We further examined the necessity of our proposed multi-axis approach, as shown in \cref{tab:ablation}B. We conducted experiments over (1) baseline UNet, (2) by adding the local branch of MAB (MAB\textsubscript{$\ell$}), (3) by adding the global branch of MAB (MAB\textsubscript{$g$}), (4) by adding the local branch of CGB (CGB\textsubscript{$\ell$}), (5) by adding the global branch of CGB (CGB\textsubscript{$g$}). Note that the huge jump (+1.04 dB) of PSNR by adding MAB\textsubscript{$\ell$} can be largely attributed to the addition of input and output channel projection layers, because we also observe a high performance of \textbf{31.42} dB PSNR if only MAB\textsubscript{$g$} is added. Overall, we observed a \textit{major} improvement when including MAB, and a relatively \textit{minor} gain when adding CGB.

\noindent\textbf{C. Why multi-stage?} Towards understanding this, we scaled up MAXIM in terms of width (channels), depth (downscaling steps), and the number of stages. \cref{tab:ablation}C suggests that packing the backbone into multi-stages yields the best performance \vs complexity tradeoff (32.44 dB, 22.2 M, 339.2 G), compared to making it wider or deeper.

\noindent\textbf{D. Beyond gMLP: the MAXIM families.}
As described in \cref{ssec:multi-axis-gating-mlp}, our proposed multi-axis approach (\cref{fig:multi-axis-gated-mlp-block}) offers a scalable way of applying \textit{any} 1D operators on (high-resolution) images, with linear complexity relative to image size while maintaining fully-convolutional. We conducted a pilot study using MAXIM-\textcolor{g}{1S} and -\textcolor{r}{2S} on SIDD~\cite{abdelhamed2018high} to explore the MAXIM families: MAXIM-FFT, -MLP, -gMLP (modeled in this paper), -SA, where we use the Fourier Transform filter~\cite{rao2021global,lee2021fnet}, spatial MLP~\cite{tolstikhin2021mlp}, gMLP~\cite{liu2021pay}, and self-attention~\cite{dosovitskiy2021an} on spatial axes using the same multi-axis approach (\cref{fig:multi-axis-gated-mlp-block}). As \cref{tab:ablation}D shows, the gMLP and self-attention variants achieved the best performance, while the FFT and MLP families were more computationally efficient. We leave deeper explorations to future works.




\begin{table}[!t]
\centering
\scriptsize
\setlength{\tabcolsep}{1.2pt}
\renewcommand{1.12}{1.1}
\begin{tabular}[t]{@{}c@{}}

\scriptsize
\setlength{\tabcolsep}{1.8pt}
\begin{tabular}{cccc|c}
 \scriptsize{CGB\textsubscript{IS}}  &  \scriptsize{SAM} &  \scriptsize{CGB\textsubscript{CS}} &  \scriptsize{MS-Sp} &  \scriptsize{PSNR}   \\
\midrule
  &   &   &  & 30.73 \\ \arrayrulecolor{lightgray}\hline
 \cmark &   &   & & 30.98   \\ \arrayrulecolor{lightgray}\hline
 \cmark & \cmark &   &  & 31.61   \\ \arrayrulecolor{lightgray}\hline
 \cmark & \cmark & \cmark & & 31.97  \\ \arrayrulecolor{lightgray}\hline 
 \cmark & \cmark & \cmark & \cmark & 32.23 \\ 
\arrayrulecolor{black}\bottomrule
\multicolumn{5}{c}{{A. Individual components.}}\\
\end{tabular}

\hfill
\scriptsize
\setlength{\tabcolsep}{1.8pt}
\begin{tabular}{cccc|c}
\scriptsize{MAB\textsubscript{$\ell$}} & \scriptsize{MAB\textsubscript{$g$}} & \scriptsize{CGB\textsubscript{$\ell$}} & \scriptsize{CGB\textsubscript{$g$}} & \scriptsize{PSNR}    \\
\midrule
  &   &   &   & 30.48   \\ \arrayrulecolor{lightgray}\hline
 \cmark &   &   &  & 31.52    \\ \arrayrulecolor{lightgray}\hline
 \cmark & \cmark &   &  & 31.68 \\ \arrayrulecolor{lightgray}\hline
 \cmark & \cmark & \cmark & & 31.84 \\ \arrayrulecolor{lightgray}\hline
 \cmark & \cmark & \cmark & \cmark & 31.91 \\
\arrayrulecolor{black}\bottomrule
\multicolumn{5}{c}{{B. Effects of multi-axis approach.}}\\

\end{tabular} 
\\\\

\def1.12{1.26}
\setlength{\tabcolsep}{1.8pt}
\scriptsize
\begin{tabular}{l|ccc|c|r|r}
 & S & W & D & PSNR & Params & FLOPs \\
\midrule
 Base & 1 & 32 & 3 & 31.08 & 6.1M & 93.6G  \\ 
 \arrayrulecolor{lightgray}\hline
 \multirow{2}{*}{Wider} & 1 & 64 & 3 & 32.09 & 19.4M & 309.9G \\
                        & 1 & 96 & 3 & 32.31 & 41.7M & 648.9G \\
 \arrayrulecolor{lightgray}\hline
 \multirow{2}{*}{Deeper} & 1 & 32 & 4 & 31.17 & 19.8M & 121.6G  \\
& 1 & 32 & 5 & 31.43 & 75.0M & 153.4G \\ \arrayrulecolor{lightgray}\hline

 More          & 2 & 32 & 3 & 31.82 & 14.1M & 216.4G \\
 stages    & 3 & 32 & 3 & 32.44 & 22.2M & 339.2G \\
\arrayrulecolor{black}\bottomrule
\multicolumn{7}{c}{{C. Why multi-stage?}}\\
\end{tabular}
\hfill

\scriptsize
\setlength{\tabcolsep}{1.8pt}
\def1.12{1.12}
\begin{tabular}{l|r|r|r}
Variant & PSNR & Params & FLOPs \\
\midrule

M1-FFT &  39.67 & 4.1M & 71G \\ 
M1-MLP & 39.75 & 5.4M & 83G \\ 
M1-gMLP & 39.80 & 6.1M & 93G \\ 
M1-SA & 39.79 & 5.3M & 111G \\ \arrayrulecolor{lightgray}\hline
M2-FFT &  39.74 & 10.1M & 172G \\ 
M2-MLP & 39.70 & 12.7M & 195G \\ 
M2-gMLP & 39.83 & 14.1M & 216G \\ 
M2-SA & 39.85 & 12.5M & 250G \\

\arrayrulecolor{black}\bottomrule
\multicolumn{4}{c}{{D. Beyond gMLP.}}\\
\end{tabular} 
\\

\end{tabular}
\caption{Ablation studies. Components in subtable A and B are defined in \cref{ssec:ablation}. S, W, and D denote the number of stages, width, and depth, respectively. M1 and M2 in subtable D denote MAXIM-1S and MAXIM-2S models, respectively.}
\label{tab:ablation}
\vspace{-1mm}
\end{table}

\section{Conclusion}
\label{sec:Conclusion}

We have presented a generic network for restoration or enhancement tasks, dubbed MAXIM, inspired by recently popular MLP-based global models. Our work suggests an effective and efficient approach for applying gMLP to low-level vision tasks to gain global attention, a missing attribute of basic CNNs. Our gMLP initialization of the MAXIM family significantly advances state-of-the-arts in several image enhancement and restoration tasks with moderate complexity. We demonstrate a few applications, but there are many more possibilities beyond the scope of this work which could significantly benefit by using MAXIM.
Our future work includes exploring more efficient models for extremely high-resolution image processing, as well as training large models that can adapt on multiple tasks.

\noindent\textbf{Broader impacts.} 
The proposed model can be used as an effective tool to enhance and retouch daily photos. 
However, enhancing techniques such as denoising and deblurring are vulnerable to malicious use for privacy concerns.
The models trained on specific data may express bias.
These issues should be responsibly taken care of by researchers.




\clearpage
