gSwin: Gated MLP Vision Model with Hierarchical Structure of Shifted Window

Although the mechanism of attention that focuses directly on a specific region in an image has actually been studied in the image domain for quite some time (Mnih et al., 2014; Xu et al., 2015; Parmar et al., 2018; Woo et al., 2018; Cordonnier, Loukas, and Jaggi, 2019; Hu, Shen, and Sun, 2018; Zhang et al., 2020), it is only recently that vision models built principally on self-attention mechanisms have been proposed (Carion et al., 2020; Wu et al., 2020; Dosovitskiy et al., 2021). Of these, ViT (Dosovitskiy et al., 2021) (Figure 1) produced competitive results and marked the beginning of research on vision models based on the Transformer architecture. Although it has only been a short time since they were proposed, Transformer for vision has become popular, and it has also been applied in related domains including video (Liu et al., 2021c; Arnab et al., 2021) and 3D point cloud (Guo et al., 2021b).

In the 2D vision domain, a vast amount of research to improve ViT has been actively conducted. One trend is oriented toward improving learning strategies, such as DeiT (Touvron et al., 2021b) and MoCo v3 (Chen, Xie, and He, 2021). Another direction is the improvement of model architectures, including CaiT (Touvron et al., 2021c), DeepViT (Zhou et al., 2021), HaloNet (Vaswani et al., 2021), etc. A major stream of research on improving model structure is the incorporation of spatial structure in images, which has not been explicitly used in Transformers, except for the customary use of positional encoding and the patch partitioning before the initial layer. An approach of incorporating spatial structure is the consideration of hybrid architectures with the auxiliary use of convolution, e.g. BoTNet (Srinivas et al., 2021), ConViT (d’Ascoli et al., 2021), LeViT (Graham et al., 2021), CvT (Wu et al., 2021), LocalViT (Li et al., 2021), etc. Others adopted hierarchical structure that lacks in vanilla Transformer models despite its importance in vision domain (See Section 2.3), e.g.  T2T-ViT (local token aggregation) (Yuan et al., 2021), HVT (pooling-based aggregation) (Pan et al., 2021), TNT (alternating local/global self-attention) (Han et al., 2021), MViT (fine-to-coarse aggregation using multi-head pooling attention) (Fan et al., 2021), PVT (pyramid structure for dense prediction) (Wang et al., 2021), ViL (efficient transformer on hierarchical structure) (Zhang et al., 2021), Swin Transformer (restricted self-attention in shifted windows) (Liu et al., 2021b), etc.

Looking back in history, the neural models of image recognition have been the multi-layer perceptron (MLP) since the early days when Rosenblatt (1958) proposed the perceptron in 1950s. Nevertheless, in the modern context, MLP-Mixer (Tolstikhin et al., 2021), ResMLP (Touvron et al., 2021a) and gMLP (Liu et al., 2021a), which were proposed almost simultaneously, would be the pioneers of the recent trend of using only axis-wise MLPs, or more precisely, axis-wise dense affine projections with non-linear activation functions.

Of these, MLP-Mixer and ResMLP simply apply axis-wise MLPs in channel and spatial axes alternately. Following these studies, a number of vision models based on axis-wise MLPs have been proposed recently. One of the current trends is the approach of constructing complex receptive fields despite being MLPs, by reshaping the tensors in the intermediate layers, e.g. Hire-MLP (Guo et al., 2021a), Vision Permutator (Hou et al., 2022). Some methods in this line even excluded the spatial-mixing MLPs, e.g. CycleMLP (Chen et al., 2021), S${}^2$-MLP (Yu et al., 2022), AS-MLP (Lian et al., 2021). In these methods, only the channel-mixing MLPs are used, and the exchange of information between tokens is taking place by spatially shifting each channel separately.

A method that deserves special mention in a different line is the gMLP (Liu et al., 2021a) (Figure 1). Its noteworthy characteristic is the use of the spatial gating unit (SGU), a gating mechanism for spatial mixing, which mimics an attention mechanism. Another feature of this method is its simplicity: only one linear (fully-connected projection) layer is used where the stacking of two or more dense layers are required in other related methods. The gMLP model has achieved competitive performance both in language and vision domains with standard Transformer models BERT (Devlin et al., 2018) and ViT (Dosovitskiy et al., 2021). It is also reported that the gMLP could learn shift-invariant features.

In neuroscience, a hierarchical structure of visual cortex has already been argued in 1960s by Hubel and Wiesel (1962). It has been a common understanding that local patches of images are matched with Gabor-like patterns in the early stage viz. primary visual cortex (V1), and the results are integrated at later stages for more global and intelligent information processing (Marĉelja, 1980). Neocognitron (Fukushima, 1980), one of the earliest CNN models, was based on the insights from physiological hierarchy model above. Subsequent CNN models in the deep learning era also naturally utilize hierarchical structures of images.

In contrast, neither Vision Transformers nor Vision MLPs had structures that inherently incorporates the ‘spatial inductive bias’ of visual hierarchy. Therefore, from such motivations it is natural to consider integrating the vision hierarchy model into Transformers, as shown in Section 2.1. In addition to the physiological motivations, such a built-in hierarchical structure has also an advantage that the output of intermediate layers can be used in downstream tasks that require multiple resolutions, including the semantic segmentation as were in e.g. UPerNet (Xiao et al., 2018). Furthermore, divide-and-conquer-like hierarchical architectures that often logarithmically reduce the computational complexity will naturally improve the efficiency.

One such proposal is the Swin Transformer (Liu et al., 2021b). It first divides the image into many small patches and applies self-attention within those patches, which are shifted by half for each layer. The structure is then gradually integrated from local to global in a bottom-up manner. (Figure 1.) Particular essentials in the structure of Swin Transformer are the bottom-up hierarchical structure for each resolution and the shifting of windows. Indeed, it has been reported that as long as this structure is maintained, the attention can actually be replaced by average pooling (PoolFormer), or even an identity map, although performance is slightly reduced (Yu et al., 2021). This fact would be a evidence that the hierarchical structure of the Swin Transformer is advantageous as a vision model.

In this paper, we aim to achieve high performance while the computational and parameter efficiency is on par with the existing methods, by incorporating both merits of the two methods that have evolved in different directions from ViT (Dosovitskiy et al., 2021) (Figure 1). In particular, we consider the two daughters of ViT, viz. gMLP (Liu et al., 2021a) (Figure 1) and Swin Transformer (Liu et al., 2021b) (Figure 1). From Figure 1, we may observe that gMLP’s overall structure is largely the same as ViT, but the structure of individual blocks has been significantly modified. On the other hand, Swin Transformer’s overall structure has been significantly modified in terms of the introduction of a hierarchical structure, but the internal structure of the individual blocks is largely the same as in ViT, except for the introduction of shifted windows. Thus, there is little overlap and a high degree of independence in the differences from ViT in each of these methods. Therefore, by incorporating both at the same time, higher improvements can be expected.

Based on these observations, the authors propose gSwin as a new vision model that inherits features of both gMLP and Swin Transformer. Figure 2 shows the architecture of gSwin. In the following subsection, we will discuss the methods of Window-SGU.

Note that this is not the same as nor a variant of MetaFormers (Yu et al., 2021), which is a general form for a series of models in which the self-attention module is replaced with another module in Swin Transformer. In this study, they considered using gMLP module instead of self-attention, but the architecture is not the same as our gSwin. For example, gSwin does not have subsequent LayerNorm, MLP and residual connections after each gMLP block. Our method is more parameter efficient, which will be shown experimentally later in this paper.

Let us first recall the gMLP (Liu et al., 2021a) model. Figure 1 shows the overall architecture of gMLP, which consists of stacking of axis-wise fully-connected layers, layer normalization (Ba, Kiros, and Hinton, 2016), GELU (Hendrycks and Gimpel, 2016) activation function and the spatial gating unit (SGU). Here, let us briefly introduce the SGU. Let $Z\in R^{H\times W\times 2C}$ be the input features, and let $Z_1,Z_2\in R^{\left(HW\right)\times C}$ be the reshaped and split features $[Z_1,Z_2]\leftarrow split\left(reshape\right(Z\left)\right)$. Then the SGU performs the following operation,

where $W\in R^{\left(HW\right)\times \left(HW\right)}$ and $b\in R^{\left(HW\right)}$ are trainable parameters, and $\odot$ denotes the element-wise multiplication. Such a gating mechanism allows us to explicitly incorporate the second-order terms $Z_1[h_1,w_1,:]\odot Z_2[h_2,w_2,:]$ associated with the binomial relationship between all the site pairs $(h_1,w_1)$ and $(h_2,w_2)$. The gMLP paper does not go so far as to say that this is the primary reason for the effectiveness of Transformers (in which up to ternary relations are incorporated), but the SGU could be an alternative mechanism to attention to explicitly incorporate higher-order relations, implicitly assuming that it could be a contributing factor. According to (Liu et al., 2021a, Figures 3), the resulting learned weight parameters $W$ were almost-sparse Gabor-like patterns. Furthermore, (Liu et al., 2021a, Figure 9) suggests that Toeplitz-like (i.e., shift invariant) structures are automatically acquired. These facts imply that it is not necessary to model $W$ as a matrix with full degrees of freedom. Instead, it would be sufficient to consider only local interactions by partitioning the features into small patches.

Based on this motivation, in this paper, we propose the Window-SGU module, which has basically the same structure as the SGU module of gMLP, but instead of the entire image, it is divided into small patches and token mixing is performed independently in each of them. That is, we first split the input feature $\text{Z}\in R^{H\times W\times 2C}$ into small patches ${\text{Z}}^{\left(i\right)}\in R^{h\times w\times 2C}$ where $1<h\ll H,1<w\ll W$. (The areas that extend beyond the $H\times W$ region during this patch partitioning are handled with zero padding as usual.) After the patch partitioning, each patch $Z^{\left(i\right)}\in R^{h\times w\times 2C}$ is reshaped and split into $Z_1^{\left(i\right)},Z_2^{\left(i\right)}\in R^{\left(hw\right)\times C}$, and then the SGU is applied to each patch independently as follows,

where $W_{\text{win}}\in R^{\left(hw\right)\times \left(hw\right)},b_{\text{win}}\in R^{\left(hw\right)}$. The resulting $Y^{\left(i\right)}$-s are again tiled and the boundaries are processed appropriately to obtain $Y\in R^{H\times W\times C}$.

If the patch partitioning is always fixed, there will always be elements that are in the neighborhood but do not interact, resulting in block noise. To prevent this, it is necessary to change the patch partitioning as necessary so that interaction between neighboring elements occurs. For layers with shifts, we cannot use cyclic-shifting as proposed in (Liu et al., 2021b) because SGUs lack the ability of masking, which self-attention has. Then naive approach is to add zero-padding so that we have uniform-shaped windows, which costs FLOPs. Here, we propose another approach equivalent to zero-padding but more efficient; the Padding-free shift. Let the shape of windows be $(h,w)=(7,7)$, then we have 9 groups of windows of shape $(h,w)=(3,3),(3,7),(3,4),(7,3),(7,7)\dots$, located at upper-left, upper, upper-right, left, center, etc. Calculate SGUs independently for each group and then concatenate all windows. The weight and bias of each group are copied from those of the main group at the center of shape $(7,7)$ properly; e.g., $W_{\text{upper}}\in R^{(3\times 7)\times (3\times 7)}$ is defined as

The Padding-free shift reduces FLOPs of gSwin-T from 3.8G (zero-padding) to 3.6G.

The experimental condition was as follows. The dataset was ImageNet-1K (Deng et al., 2009), divided into training and validation sets consisting of $\sim$1.3M and 50k images, respectively.

Table 1 shows the parameter settings for three variants of the proposed method (we will explain about the #heads column in Section 4.4.1). We set window size to $7$ for all gSwin models.

The training conditions were as follows (most of them align with those of Swin Transformer (Liu et al., 2021b) for fair comparison). The optimizer was AdamW (Kingma and Ba, 2014). Data augmentation and regularization techniques were adopted including randaugment (Cubuk et al., 2020; Wightman, 2019), and the drop-path regularization (Larsson, Maire, and Shakhnarovich, 2016), whose rates are shown in Table 2 (tuned by grid search). The initial learning rate was 0.001 and the weight decay was 0.05. The input image size was ${224}^2$, the batch size was 1024 and the number of epochs was 300 with cosine scheduler (first 20 epochs for linear warm-up). Of these, the epoch with the best top 1 accuracy in the validation dataset was selected (in case of a tie, the selection was based on the top 5 accuracy). The training was performed twice with different seeds. We report the average (and unbiased variance) values of the two trials.

We compared the image recognition performance of our method with that of DeiT (Touvron et al., 2021b), Swin Transformer (Liu et al., 2021b), gMLP (Liu et al., 2021a) and Hire-MLP (Guo et al., 2021a). Table 3 shows the results. Compared to Swin Transformers, gSwin-T achieves +0.4% top 1 accuracy with 21% less parameters, and gSwin-S achieves the same accuracy with 20% less. We may observe that gSwin (proposed) is more efficient in terms of the number of parameters and floating point operations to achieve the same level of accuracy with others.

We performed an evaluation experiment with object detection and instance segmentation tasks, to test performance of gSwin on downstream tasks. The dataset we used in this experiment was COCO (Lin et al., 2014). The evaluation metric was box and mask AP (average precision). We used Cascade Mask R-CNN (He et al., 2017; Cai and Vasconcelos, 2018) as object detection framework, as Swin Transformer (Liu et al., 2021b) does.

When we train the object detection, we reused checkpoints of the ImageNet-1K experiment rather than training a new model in a full-scratch fashion. The specific procedure was as follows. Firstly, the checkpoint model that achieved the highest accuracy on ImageNet-1K validation dataset was used as the starting point. Then, we continued training for 36 epochs on 4 V100 GPUs, where the batch size was 4 for each GPU. The optimizer was AdamW (Kingma and Ba, 2014), the initial learning rate was 0.0001, and the weight decay was 0.05. We used multi-scale training (the input image was resized so that the shorter is between 480 and 800 pixels and the longer is at most 1333). During the training, we evaluated the AP scores for every epoch and selected the best one with the highest box AP. This transfer learning was done three times using different random seeds, which resulted in 6 different models. We report the average score of these 6 models.

From results in Table 4, we observe that gSwin-T achieves +0.45/+0.38 box/mask AP with 8% less parameters, and gSwin-S achieves similar APs with 9% less parameters.

We also performed an evaluation experiment with an image segmentation task. The dataset we used in this experiment was ADE20K (Zhou et al., 2019). The evaluation metric was mIoU (mean intersection over union). We used UPerNet (Xiao et al., 2018) as base framework, as in (Liu et al., 2021b).

When we train the semantic segmentation, we reused checkpoints of the ImageNet-1K experiment. The specific procedure was as follows. Firstly, the checkpoint model that achieved the highest accuracy on ImageNet-1K validation dataset was used as the starting point. Then, we continued training for 160k steps with linear scheduler (first 1.5k steps for linear warm-up) on 4 GPUs, where the batch size was 4 for each GPU. The optimizer was AdamW (Kingma and Ba, 2014), the initial learning rate was $6\times {10}^{-5}$, the weight decay was 0.01, and the input image size was ${512}^2$. The augmentations we adopted was random horizontal flip, random rescaling (the resolution ratio within $[0.5,2.0]$) and random photometric distortion, as in Swin Transformer (Liu et al., 2021b). During the training, we evaluated the mIoU score for every 4k steps and selected the best one with the highest validation score. This transfer learning was done three times using different random seeds, which resulted in 6 different models. We report the average score of these 6 models.

Results on ADE20-K are shown in Table 5, where the mIoU(aug) column is the score with test-time multi-scale ($[0.5,0.75,1.0,1.25,1.5,1.75]\times$ resolution) horizontal-flip augmentation. These show that gSwin-T achieves +1.8/+1.85% mIoU with/without augmentation with 13% less parameters, and gSwin-S achieves +0.46/+0.57% mIoU with 14% less parameters, compared to Swin Transformer.