Brain-Inspired Stepwise Patch Merging for Vision Transformers

The Vision Transformer (ViT) (Dosovitskiy et al. 2020) revolutionized visual tasks by introducing the transformer architecture to computer vision. ViT segments images into non-overlapping patches, projects these patches linearly into token sequences, and processes them using a transformer encoder. ViT models have demonstrated superior performance in image classification and other downstream tasks, surpassing CNNs (He et al. 2016; Krizhevsky, Sutskever, and Hinton 2012; LeCun, Bengio et al. 1995) when trained with large-scale pretraining datasets and advanced training methodologies. Motivated by the success of CNNs and the necessity to address dense prediction tasks, researchers have incorporated the feature pyramid structure within transformers. This innovation has led to the development and widespread adoption of HVTs (Dong et al. 2022; Ren et al. 2022; Fan et al. 2023; Guo et al. 2022; Hou et al. 2022).

Hierarchical feature representation plays a pivotal role in dense prediction tasks, prompting extensive research in this domain. Existing approaches can be broadly categorized into fixed-grid and dynamic feature-based methods. Fixed-grid methods, exemplified by works such as PVT (Wang et al. 2021) and Swin (Liu et al. 2021), merge patches within adjacent windows using 2D convolution. In contrast, dynamic methods, such as DynamicViT (Rao et al. 2021) adaptively extract features by eliminating redundant patches and retaining essential ones, thereby forming hierarchical feature maps. EviT (Liang et al. 2022) enhances this approach by selecting the top K tokens with the highest average values across all heads for the next stage, merging the remaining tokens. PS-ViT (Yue et al. 2021) further refines the process by iteratively adjusting patch centers towards the object to enrich object information within the hierarchical feature maps. Token Merging (Bolya et al. 2022) employs cosine similarity to progressively merge similar tokens, thereby increasing model throughput.

Fixed-grid methods are constrained by their singular and relatively small receptive fields, and excessively enlarging the grid size leads to increased computational overhead. Dynamic feature-based methods, while adaptive, may discard low-scoring tokens that contain valuable information and often lack end-to-end training capabilities. Our proposed Stepwise Patch Merging approach distinguishes itself from both fixed-grid and dynamic feature-based methods. The Multi-Scale Aggregation module in Stepwise Patch Merging provides an expanded and enriched receptive field, which is advantageous for long-distance modeling. Additionally, the Guided Local Enhancement module enhances the extraction of local discriminative features and supports end-to-end training. Moreover, Stepwise Patch Merging can be directly applied to dense prediction tasks, resulting in improved performance.

Our proposed Multi-Scale Aggregation (MSA) module draws inspiration from the brain’s remarkable ability to effectively model long-range dependencies when processing visual information. In the brain, the visual system achieves precise modeling of long-range dependencies through multi-level and multi-scale information processing. Neurons with small receptive fields process local features, and this information is progressively integrated over larger areas by neurons with larger receptive fields, capturing complex patterns and objects. Additionally, the brain’s extensive network of long-range neural connections allows for the exchange and integration of data from various parts of the visual field, facilitating a comprehensive understanding of the scene. Furthermore, neurons within the same level possess receptive fields of varying sizes, enabling the brain to simultaneously process local details and global features. This sophisticated mechanism of combining local and global information processing in the brain inspired the design of our MSA module, which aims to enhance feature diversity and capture long-range dependencies effectively.

Inspired by these mechanisms, the MSA module first divides the input channels $C$ into $N$ distinct heads, each undergoing depth-wise convolutions with varying receptive fields. This method not only reduces the parameter count and computational cost but also facilitates the extraction of multi-granularity information, akin to how different neurons in the brain handle information processing. Subsequently, the MSA module employs larger convolutional kernels to further expand the receptive field, thereby enhancing its capability to model long-range dependencies. Following this, Channel Shuffle (Zhang et al. 2018) is used to interleave channels containing features of different scales, followed by a series of linear projections to fuse these multi-scale features. The number of linear projections is $\frac{C}{N}$, with each projection having unique parameters. Finally, the $N$ heads are concatenated, and a final linear projection adjusts the number of channels to the specified $C^{\prime}$.

By leveraging the brain’s mechanism for effective long-range dependency modeling, the MSA module better captures and integrates key features, significantly enhancing the model’s performance in complex visual recognition tasks.

Our proposed MSA can be formulated as follows:

		$\displaystyle H_{n}=DWConv_{k_{n}\times k_{n}}(x_{n})$		(1)
		$\displaystyle G^{c}=\mathbf{W}^{c}([H_{1}^{c};H_{2}^{c};...;H_{N}^{c}])$
		$\displaystyle\bm{MSA}(X)=\mathbf{W}([G^{1};G^{2};...;G^{\frac{C}{N}}])$

where $X=[x_{1},x_{2},\ldots,x_{N}]$ represents the input $X$ split into multiple heads along the channel dimension, and $x_{n}\in\mathbb{R}^{B\times H\times W\times\frac{C}{N}}$ denotes the $n$-th head. The kernel size of the depth-wise convolution for the $n$-th head is denoted by $k_{n}\in{k_{1},k_{2},\ldots,k_{N}}$. Here, $H_{n}\in\mathbb{R}^{B\times H\times W\times\frac{C}{N}}$ represents the $n$-th head after being processed by the depth-wise convolution with $out_{c}hannels=in_{c}hannels$, and $H_{n}^{c}$ represents the $c$-th channel in the $n$-th head. $\mathbf{W}^{c}\in\mathbb{R}^{N\times N}$ is the weight matrix of the linear projection. $\mathbf{G}^{c}\in\mathbb{R}^{B\times H\times W\times N}$. Finally, $\mathbf{W}\in\mathbb{R}^{C\times C^{\prime}}$ is the weight matrix of the linear projection that adjusts the number of channels to the specified $C^{\prime}$.

Inspired by the brain’s ability to enhance local features through context-aware processing, we developed the Guided Local Enhancement (GLE) module. In the brain, local feature enhancement is achieved by integrating information from both local and global contexts. Higher-level cortical areas provide contextual feedback that refines the processing of local features, ensuring that details are interpreted within the broader visual context. This hierarchical processing involves neurons that respond specifically to local stimuli but are influenced by surrounding contextual information, allowing for more nuanced and precise feature extraction.

Following this principle, the GLE module acts as a local feature enhancer utilizing context-aware guide tokens, as illustrated in Fig. 3. Specifically, we implement self-attention within a local window and introduce a $[Guide]$ token into the input space. This $[Guide]$ token undergoes the same self-attention operations as the patch tokens, thereby providing high-level semantic features to the local window during pooling. By mimicking the brain’s method of using contextual information to refine local feature extraction, the GLE module ensures that the extracted local features are both precise and contextually relevant.

Formally, given an input $X\in\mathbb{R}^{C\times H\times W}$, the $[Guide]$ tokens are generated by a large-kernel depth-wise convolution, referred to as the Guide Token Generator (GTG), which can be described as follows:

$\displaystyle\bm{GTG}(X)=DWConv(GELU(BatchNorm(X))).$

(2)

We focus on the operations performed on a single pixel within the input feature map. We define a set of pixels within a local window centered at pixel $(i,j)$ as $\rho(i,j)$. For a fixed window size of $k\times k$, $\|\rho(i,j)\|=k^{2}$. In our setup, $k$ is equal to the stride of the GTG, both being 2, meaning $\#Windows=\|GTG(X)\|=\frac{H\times W}{4}$. Tokens within a window containing a $[Guide]$ token can be represented by the sequence $\mathbf{z}$:

$\displaystyle\mathbf{z}=\left[S_{(i,j)\sim GTG(X)};S^{1}_{(i,j)\sim\rho(i,j)};\cdots;S^{k^{2}}_{(i,j)\sim\rho(i,j)}\right],$

(3)

then perform the standard self-attention operation on $\mathbf{z}$:

	$\displaystyle\left[\mathbf{q},\mathbf{k},\mathbf{v}\right]$	$\displaystyle=\mathbf{z}\mathbf{U}_{qkv},$		(4)
	$\displaystyle SA(\mathbf{z})$	$\displaystyle=softmax(\frac{\mathbf{q}\mathbf{k}^{\top}}{\sqrt{C}})\mathbf{v},$

where $\mathbf{U}_{qkv}\in\mathbb{R}^{C\times 3C}$, and we ignore the relative positional relationships between tokens within a window, so positional encoding is not used. Finally, we select the $[Guide]$ token as the output of the GLE:

		$\displaystyle SA(\mathbf{z})=\left[A_{(i,j)\sim GTG(X)};A_{(i,j)\sim\rho(i,j)}\right],$		(5)
		$\displaystyle\bm{GLE}(X_{(i,j)})=A_{(i,j)\sim GTG(X)}.$

It is worth noting that when the $[Guide]$ token is stripped of its semantic information, it degrades into the $[CLS]$ token in the vanilla vision transformer (Dosovitskiy et al. 2020). Experiments show that using our $[Guide]$ tokens results in higher performance (see Tab. 6).

Setting. We first evaluate the proposed SPM framework on the ImageNet-1K dataset (Deng et al. 2009), which comprises 1.28 million training images and 50,000 validation images spanning 1,000 categories. To ensure a fair comparison, all models are trained on the training set and report the top-1 error rate on the validation set. For data augmentation, we apply a suite of techniques including random cropping, random horizontal flipping (Szegedy et al. 2015), label smoothing regularization (Szegedy et al. 2016), mixup (Zhang et al. 2017), CutMix (Yun et al. 2019), and random erasing (Zhong et al. 2020). These augmentations are employed to enhance the robustness and generalization ability of the models. During training, we use the AdamW optimizer (Loshchilov and Hutter 2017) with a momentum parameter of 0.9, a mini-batch size of 128, and a weight decay of $5\times 10^{-2}$. The initial learning rate is set to $1\times 10^{-3}$ and follows a cosine annealing schedule (Loshchilov and Hutter 2016) to gradually reduce the learning rate. All models are trained from scratch for 300 epochs on eight NVIDIA A100 GPUs. For evaluation, we adopt the standard center crop strategy on the validation set, where a $224\times 224$ patch is extracted from each image to assess the classification accuracy.

Result. In Tab. 2, we observe that incorporating the SPM framework into the PVT results in significant improvements in classification accuracy, specifically by 4.4%, 1.9%, and 0.7% in the Tiny, Small, and Medium models, respectively, while adding only a minimal number of parameters compared to the original PVT. The final experimental results indicate that the accuracy of models of various sizes has been enhanced, with the most notable improvement observed in the Tiny model. Remarkably, the combination of PVT-Tiny and SPM achieved a top-1 accuracy of 79.5%, which is comparable to the performance of PVT-Small, despite PVT-Small having nearly 70% more parameters than PVT-Tiny. Furthermore, with the integration of SPM, PVT-Small surpassed PVT-Medium by 0.5%.

Setting. Object detection and instance segmentation experiments were conducted on the challenging COCO benchmark (Lin et al. 2014). All models were trained on the training set comprising 118k images and evaluated on the validation set with 5k images. We validated the effectiveness of different backbones using Mask R-CNN (He et al. 2017). Before training, the weights pre-trained on ImageNet-1K were used to initialize the backbone, and the newly added layers were initialized using the Xavier initialization method (Glorot and Bengio 2010). Our models were trained with a batch size of 16 on 8 NVIDIA A100 GPUs and optimized using the AdamW optimizer (Loshchilov and Hutter 2017) with an initial learning rate of $1\times 10^{-4}$.

Result. As shown in Tab. 1, incorporating the SPM framework into the PVT resulted in significant improvements of 4.1%, 2.6%, and 1.3% in the Tiny, Small, and Medium models, respectively, for the object detection task. Notably, the SPM framework also demonstrated substantial improvements in the instance segmentation task. Several observations can be made by analyzing the detection results of models with different sizes. Models integrated with SPM show marked improvements in detecting medium-sized and large objects. This enhancement is attributed to the original patch merging’s relatively singular and small receptive fields, whereas the MSA module integrates features with diverse receptive fields, enabling the model to more accurately capture long-range dependencies. Moreover, there is a significant improvement in detecting small objects. Although larger models are typically better at modeling global relationships, the disruption of local information may hinder small objects from establishing complete semantic information, leading to missed detections. The GLE module addresses this by enhancing the perception of local discriminative information, resulting in consistent improvements in the detection performance of small objects across models of different sizes.

Setting. The ADE20K dataset (Zhou et al. 2017) is a widely utilized benchmark for semantic segmentation, comprising 150 categories with 20,210 images for training, 2,000 images for validation, and 3,352 images for testing. All the methods compared were evaluated using the Semantic FPN framework (Kirillov et al. 2019). The backbone network of our method was initialized with the pre-trained ImageNet-1k model, and the newly added layers were initialized using the Xavier initialization method. The initial learning rate was set to 0.0001, and the model was optimized using the AdamW optimizer. We trained our models for 40,000 iterations with a batch size of 16 on eight NVIDIA A100 GPUs. The learning rate followed a polynomial decay schedule with a power of 0.9. During training, images were randomly resized and cropped to $512\times 512$ pixels. For testing, images were rescaled to have a shorter side of 512 pixels.

Result. As shown in Tab. 3, the integration of the SPM framework led to a significant enhancement in the semantic segmentation task. Specifically, the performance of the Tiny, Small, and Large models improved by 5.8%, 6.1%, and 3.7%, respectively. It is evident that the improvement achieved by SPM in segmentation tasks surpasses that in classification and detection tasks. Interestingly, in both classification and detection tasks, the relative improvement brought by SPM compared to the base model gradually decreases as the model size increases. However, on the ADE20K dataset, the performance gain of PVT-Small exceeds that of PVT-Tiny, with improvements of 6.1% and 5.8%, respectively. This phenomenon can be attributed to the GLE module within SPM, which is specifically designed to capture local information. Consequently, SPM demonstrates a significant advantage in semantic segmentation tasks, where detailed local feature extraction is crucial.

To further validate the generalizability of the proposed SPM framework, we integrated SPM into various mainstream Transformer backbones and trained them on the ImageNet-1K dataset. We employed consistent training settings to ensure a fair comparison, and the top-1 accuracies are presented in Tab. 4. The results demonstrate that the performance improvements conferred by SPM are universal across different backbones, indicating its robust generalization capability. Specifically, our method significantly enhanced the performance of Swin-T by 1.1%, Shunted-T by 0.8%, and NAT-Mini by 0.4% on ImageNet-1K. These results underscore the effectiveness of SPM in boosting the performance of various Transformer architectures, highlighting its potential as a versatile enhancement technique in the field of computer vision.

We compared the visualization results of the attention maps with and without using the SPM framework, as shown in Fig. 4. Specifically, we replaced the original patch merging block of PVT-Tiny with our SPM and visualized the first block of the second stage for two separate heads. For example, after employing SPM, the bird’s two wings and tail were successfully linked, whereas the vanilla PVT-Tiny failed to capture the distant tail. This demonstrates that SPM facilitates the network’s ability to establish long-range relationships at shallower layers, leading to significant improvements in classification performance.

Additionally, we employed the Effective Receptive Field (ERF) method (Luo et al. 2016) as a visualization tool to compare the changes in ERF before and after using SPM, as depicted in Fig. 5. It is readily observed that after integrating SPM, the size of the ERF not only increases significantly but also changes shape from a regular square to a radial decay pattern, which aligns more closely with biological vision. This pattern can be attributed to our carefully designed MSA and GLE modules, which together achieve an excellent balance between capturing long-range relationships and preserving local detail features. Consequently, this balance leads to improved performance in classification and downstream dense prediction tasks.

To evaluate the effectiveness of the proposed GLE module, we conducted experiments by replacing GLE with two alternative methods, as shown in Tab. 5). Replacing GLE with a $2\times 2$ convolution or a $3\times 3$ average pooling layer results in a performance decrease of 2.3% and 2.4%, respectively.

We conducted comparative experiments to evaluate the effectiveness of the proposed $[Guide]$ token against two mainstream methods: the $[CLS]$ token (Dosovitskiy et al. 2020) and global average pooling (GAP) (Chu et al. 2021), as presented in Tab. 6. The results demonstrate that the $[Guide]$ token improves model performance by approximately 1.7% compared to these methods, without significantly increasing the number of parameters.

Furthermore, an important observation from Tab. 6 and Tab. 5 is that when the local window size of self-attention and the kernel size of convolution are both set to $2\times 2$, the self-attention method achieves higher accuracy than the convolution method. This suggests that the self-attention mechanism has a superior capability in extracting high-frequency features.

To determine the optimal kernel size for GTG, we conducted a series of performance comparison experiments with different kernel sizes. Ultimately, we set the kernel size of GTG to 7, which achieved the best performance without significantly increasing the number of parameters. This indicates that information outside the local window positively influences the attention results within the window, highlighting the effectiveness of the $[Guide]$ token.

Different stages represent varying levels of semantic information. To validate the generalizability of the SPM framework across different levels of semantic information, we conducted experiments by gradually replacing the original patch merging with SPM (see Tab. 8). From the second and third rows, we observe that SPM enhances network performance by 1.3% with low-level features (stage 1) and by 1.6% with high-level features (stage 3), respectively. Furthermore, from the second, fourth, and fifth rows, it is evident that progressively replacing the original patch merging with SPM linearly improves performance by 1.3%, 3.0%, and 4.4%, respectively.