Axial Attention MLP-Mixer: A New Architecture for Image Segmentation

Abstract

Recently, the MLP-Mixer model has received much attention in vision problems. The advantage of this model is that by using only multi-layer perceptron (MLP) blocks, the model could build well the long-range dependencies of the input patches when pre-trained on huge data sets. Recognizing the importance of information positions in patch processing and the advantage of using MLPs, in this study, we proposed an Axial Attention MLP-Mixer model, shorted as AxialAtt-MLP-Mixer for image segmentation problem. In particular, inspired by advanced attention mechanisms along with position embedding, we proposed a new token layer that replaces the token mixing in the MLP-Mixer model to make the model more aware of global information. In addition, we propose a new model using MLP-Mixer architecture and an axial attention token layer. Through evaluation on two datasets: GlaS and Data Science Bowl 2018, we indicate the superiority of the proposed method along with the ability to get good results right on small datasets without pre-training.