Image-to-Image MLP-mixer for Image Reconstruction

1 INTRODUCTION . Deep neural networks have emerged as highly successful tools for image and signal reconstruction , restoration , and manipulation . They achieve state-of-the-art image quality on tasks like denoising , super-resolution , image reconstruction from few and noisy measurements , and image generation . Current state-of-the-art image reconstruction networks are convolutional . Convolutional neural networks ( CNNs ) achieve better denoising image quality than classical methods such as BM3D ( Zhang et al. , 2017 ; Brooks et al. , 2019 ) . They also perform excellent on many other imaging problems including computed tomography ( McCann et al. , 2017 ) and accelerated magnetic resonance imaging ( MRI ) ( Zbontar et al. , 2018 ) . For example , all top-performing methods at the FastMRI competition , a challenge for accelerated magnetic resonance imaging ( Zbontar et al. , 2018 ; Knoll et al. , 2020 ) , are CNNs . For the related problem of image classification CNNs are also state-of-the-art . However , recent work has shown that new non-convolutional networks can perform comparable when trained on huge datasets . For instance , the vision transformer ( Dosovitskiy et al. , 2021 ) is an attention-based architecture without convolutions that achieves excellent classification accuracy when pre-trained on very large datasets . Most recently , networks solely based on multi-layer perceptrons ( MLPs ) were proposed , including the MLP-mixer ( Tolstikhin et al. , 2021 ; Liu et al. , 2021a ; Chen et al. , 2021 ) . Trained on a huge dataset , the MLP-mixer performs almost as well as the best convolutional architectures while having lower computational costs at inference . Non-convolutional architectures such as the ViT and MLP-mixer impose a lower inductive bias than CNNs . This inductive bias enables CNNs to perform well when little to moderate amounts of training data are available , but might limit performance if abundant data is available . Motivated by this development , and by the simplicity of the MLP-mixer , we propose and study a variant of the MLP-mixer for image reconstruction tasks , with the premise that such a network can give better image quality than convolutional networks if trained on sufficient data . The architecture of the image-to-image MLP-mixer is depicted in Figure 1 . The image-to-image MLP-mixer differs to the MLP mixer in that it retains the relative positions of the patches , which leads to significantly better performance for image reconstruction tasks . Our results show the image-to-image mixer can outperform a state-of-the-art image reconstruction architecture , the U-net Ronneberger et al . ( 2015 ) , by a small margin . We show that the gap in performance between the image-to-image mixer and a U-net increases with the number of training images and the model size ( see Figures 2 and 3 ) . We also show that , even in the regime of relatively few training images , the image-to-image MLP-mixer slightly outperforms a U-net of similar size in image quality , both for denoising images perturbed with Gaussian noise , denoising images perturbed by real-world camera noise , and for compressed sensing reconstruction in magnetic resonance imaging . Phrased differently , to achieve the same denoising performance , the image-to-image MLP-mixer requires fewer parameters ( see Figure 2 ) . The MLP-mixer also outperforms a vision transformer tailored to image-to-image tasks , and BM3D , a classical un-trained denoising algorithm at denoising . 2 IMAGE-TO-IMAGE MLP-MIXER NETWORK ARCHITECTURE . In this section , we introduce an image-to-image MLP-mixer architecture that builds on the original MLP-mixer ( Tolstikhin et al. , 2021 ) . The image-to-image MLP-mixer operates on linearly transformed image patches , just like the MLP-mixer , as illustrated in Figure 1 . However , contrary to the MLP-mixer , the image-to-image mixer imposes some structure by retaining the spacial order of image patches , which turns out to be critical for image reconstruction performance . We start by splitting the image into non-overlapping patches of size P × P × 3 ( our default choice is P = 4 ) . Each patch is viewed as a vector of dimension 3P 2 that is linearly transformed with the same trainable matrix to a space of arbitrary embedding dimension C. This patch embedding step thus transforms an image of dimension H ×W × 3 ( or H ×W × 1 for greyscale images ) to a volume of dimension H/P ×W/P × C. The patch embedding step retains the relative positions of the patches in the image . The MLP-mixer and the vision transformer ( Tolstikhin et al. , 2021 ; Dosovitskiy et al. , 2021 ) also split an image into patches and linearly project the patches , and so do several other architectures for example the swin transformer ( Liu et al. , 2021b ) . We then apply an MLP-mixer layer inspired by the original MLP-mixer module . This MLP-mixer layer mixes the tensor in height dimension , then in width dimension , and finally in channel dimension . Mixing in channel dimension means viewing the tensor of dimension H/P ×W/P × C as a collection of H/P ·W/P vectors of dimension C and passing each of them through the same MLP consisting of a linear layer , followed by a GeLU-non-linearity and then another linear layer . The hidden layer dimension is the input dimension of the respective vector multiplied by a factor of f . We also add skip connections and layer norms to help with the optimization . A mixer layer does not alter the dimensions of the input volume . After N many such mixer layers , the volume is transformed back to an image via a patch expansion step . The patch expansion step transforms the volume consisting of flattened patches , each of dimension C , back to an image of dimensionH/P ×W/P ×3 as follows : First , we linearly transform each patch of dimension C to a patch of dimension CP 2 using a shared linear transformation . This maps the volume of shape H/P ×W/P × C to a volume of shape H/P ×W/P × CP 2 . Second , we reshape the volume to a volume of shape H ×W × C , and finally transform this volume to an image of shape H ×W × 3 by linearly combining the layers ( which can be implemented with a 1 × 1 convolution ) . A similar patch expansion step has been used by the Swin U-net Transformer ( Cao et al. , 2021 ) . The main difference between our image-to-image MLP-mixer architecture and the original MLPmixer is that we transform the image to a 3D tensor instead of a 2D tensor , and the mixer layer is modified to act on a 3D volume . This modification retains the relative location of the patches in the 3D volume which induces an inductive bias enabling the image-to-image MLP-mixer to perform 1 very well when trained on relatively few images . As we show later in Section 3.4 , the inductive bias is less than that of a convolutional network , but more than the original MLP-mixer . A further difference of the image-to-image-Mixer over the original MLP-mixer is the scaling of the number of parameters : The trainable parameters of the token mixing in the original mixer are O ( H2W 2 ) , while the height- and width mixing of the image-to-image MLP-mixer are O ( H2 +W 2 ) . The linear scaling in image resolution of the image-to-image MLP-mixer keeps the total number of trainable parameters low and the architecture memory efficient . 3 EXPERIMENTS . We evaluate the performance of the image-to-image mixer for a variety of image reconstruction problems . We focus on image denoising as it is considered to be a fundamental image reconstruction problem , for its practical importance , and since a good denoiser typically serves as a building block for other tasks such as recovering images from few and noisy measurements . For example , a state-ofthe-art approach for reconstructing an image from few and noisy linear measurements is a so-called variational network which uses a denoiser as a building block ( Sriram et al. , 2020 ) . Reconstructing an image from few and noisy linear measurements is an important inverse problem that arises in accelerated magnetic resonance imaging and sparse-view computed tomography . Baseline methods : We compare the denoising performance of the image-to-image MLP mixer to three baselines : BM3D ( Dabov et al. , 2007 ) , a standard and well performing denoising algorithm that does not rely on any training data . The U-net ( Ronneberger et al. , 2015 ) , a standard image- to-image convolutional network that is a go-to for image reconstruction problems . The U-net performs slightly better than a standard multi-layer convolutional network for image denoising ( Brooks et al. , 2019 ) ( for example better than the multi-layer convolutional network proposed by Zhang et al . ( 2017 ) ) . We also compare to the vision transformer ( Dosovitskiy et al. , 2021 ) , which we adapted for image recovery tasks as follows . We disposed the classification token and replaced the classification head by a linear layer that maps each element of the transformer output to a corresponding image patch . All networks ( the image-to-image MLP-mixer , U-net , and ViT ) are all trained in the same fashion as described next . 3.1 GAUSSIAN DENOISING . We first consider the problem of removing Gaussian noise from ImageNet color images ( Deng et al. , 2009 ) . We constructed a dataset as follows : We collected images of different classes from ImageNet and center-cropped them to a size of 256 × 256 × 3 . We then added zero-mean Gaussian noise of standard deviation σ = 30 to each image channel independently , resulting in a data set consisting of pairs of noisy image yi = xi + zi and corresponding clean image xi . Here , zi is the Gaussian noise . The noisy images have a peak signal-to-noise-ratio ( PSNR ) of 19 dB . We trained the image-to-image MLP-mixer fθ with trainable parameters θ ( and the baseline architectures ) to map the noisy image to the noise by minimizing the loss function L ( θ ) = 1 n n∑ i=1 1 2 ‖yi − fθ ( yi ) − xi‖22 . Here , n is the total number of training images . At inference , we are given a noisy image y and estimate a clean image by subtracting the estimated residual from the noisy observation as x̂ = y − fθ ( y ) . This is referred to as residual learning ( Zhang et al. , 2017 ) , because the network learns to predict the residual . Training the network directly to map a noisy image to a clean image also works , but performs worse than residual learning for all architectures considered here . We split the data set into train and test sets and ensured that images from the same ImageNet class do not exist in both sets simultaneously . This guarantees that the network is not just learning to denoise a specific class only . In Figure 2 , we depict the denoising performance of the different architectures as a function of the number of training examples , ranging from 1000 to 100.000 training images , with constant model size , and as a function of the number of parameters , with constant training set size . The plots show that even in the regime of small training data ( left : 4000 images ) and small model size ( middle : 3 million parameters ) , the image-to-image mixer can outperform the Unet . Thus , the image-to-image MLP-Mixer is more parameter effective in that it outperforms the U-net with fewer parameters , i.e . a 3M version of the image-to-image mixer performs slightly better than a 12M version of the Unet . Most importantly , Figure 2 shows that the image-to-image mixer scales better than the U-net when both the dataset size and the size of the models grow . Moreover , the right panel in Figure 2 shows that for large models ( 24M ) the gap in performance between the image-to-image mixer and the U-net increases : The U-net shows a relatively smaller accuracy improvement when increasing the training set size from 10k to 100k . Thus , we expect even larger improvements when moving to even larger datasets . In the experiment , the model parameters of the image-to-image mixer are varied by changing the number of channels and the hidden dimension of the MLPs . The exact hyperparameter configurations are in Table 3 in the appendix . For the U-net , we increased the model size by increasing the number of channels , and for ViT we increased the model size by increasing both its depth and width .

Summary: The paper proposes the use of a modified MLP-mixer for image restoration. It largely follows the original MLP-mixer architecture, except that the token mixing layers (which go from all tokens to all tokens) are replaced with a pair of separable height-wise and width-wise mixing layers. The final token/patch embeddings are then reshaped back to a pixel grid, and used to reconstruct intensities with a 1x1 conv.  The paper compares the performance of using such an architecture for denoising (and compressive reconstruction) vs ViT and U-Nets.