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 adopts the recently popular MLP-mixer architecture to the task of image reconstruction. The paper claims that the proposed approach therein(I2I-mixer) gives better or comparable performance as other SOTA approaches(UNet and ViT) but provides two important advantages. It requires much fewer parameters and also learns with fewer training examples. To demonstrate this, experiments have been performed on three image reconstruction tasks- Gaussian denoising, Real-world camera noise denoising and MRI reconstruction.