1 Introduction↩︎

Vision Transformer (ViT) has become one of the representative architectures in visual recognition [1]–[3]. ViT divides an image into a sequence of patches and processes each patch as a token, enabling it to capture long-range dependencies within an image in a manner different from convolutional neural networks. In particular, self-attention can directly model relationships between distant image regions, and has shown strong performance across various visual recognition tasks, including image classification, object detection, and semantic segmentation [2], [4]–[7]. Thus, patch-wise interactions based on self-attention have become an important component of modern vision models.

However, self-attention in ViT aggregates features based on relative weights normalized by softmax. Specifically, the weight assigned to each key patch for a given query patch is not determined independently by the query–key pair itself, but is instead determined through competition with the other key patches considered at the same time. This property makes it difficult to interpret the relevance between patches as an absolute measure. Moreover, because softmax normalization assigns non-negative weights to all keys, it cannot explicitly reject patches with low relevance. In visual recognition, an input image often contains not only the target object but also many patches corresponding to background regions, textures, and redundant local patterns that are not necessarily required for recognition. Therefore, self-attention, which relatively weights and aggregates all patches, may allow redundant or irrelevant information to be mixed into the feature representation, potentially degrading recognition performance and interpretability.

To address such limitations of softmax attention, a mechanism called Screening has been proposed in the field of language modeling [8]. Screening first normalizes queries and keys, computes bounded query–key similarities, and then applies thresholding to these similarities to independently evaluate the relevance of each key. While softmax attention redistributes probability mass across all keys, Screening defines an absolute relevance value for each query–key pair and assigns zero relevance to pairs below the threshold, thereby removing their contribution during value aggregation. This property enables Screening to explicitly reject irrelevant tokens and selectively aggregate only relevant tokens. Another important difference from self-attention is that Screening does not need to forcefully assign weight to any key when no relevant key exists, allowing it to represent the absence of relevant information.

In this work, we propose VisionScreen, a new vision model that adapts the Screening mechanism to two-dimensional image patch grids for visual recognition. VisionScreen treats image patches as tokens arranged on a two-dimensional grid and computes absolute relevance based on query–key similarity between patches. Furthermore, by extending the distance-based selection mechanism in Screening to two-dimensional space, each patch can selectively aggregate patches that are both content-wise relevant and spatially appropriate. As a result, the proposed method aims to learn feature representations that are less affected by redundant or irrelevant patches by judging the relevance of each patch independently, without relying on competition among other patches. In contrast to the original Screening mechanism, which selectively aggregates past tokens along a causal one-dimensional sequence, visual recognition is a non-causal setting in which all patches are observed simultaneously and each patch has two-dimensional neighborhood relationships. Therefore, we introduce Screening that considers relative positions between image patches, extending selection along the sequence direction to selection in the spatial direction, thereby realizing feature aggregation suited to the spatial structure of image data.

In our experiments, we evaluate the proposed method on image classification benchmarks and show that it achieves effective recognition performance compared with conventional ViT. These results suggest that selective patch aggregation based on Screening can be effective for visual recognition as an alternative to relative information aggregation based on softmax attention.

2 Related Work↩︎

2.1 Vision Transformers and Softmax Attention↩︎

ViT models images as sequences of patches and captures global dependencies via self-attention [1]. Subsequent variants such as DeiT [6], Swin Transformer [5], and MAE [7] have further improved performance, demonstrating the effectiveness of attention-based architectures in vision.

However, these models rely on softmax attention, where patch interactions are computed as relative weights normalized over all keys. As a result, each query–key relevance is not evaluated independently, and low-relevance patches cannot be explicitly rejected due to the non-zero weighting of all keys [8]. In contrast, our method introduces Screening to enable independent relevance evaluation and explicit exclusion of low-relevance patches.

2.2 Vision Models Replacing Attention↩︎

Although self-attention is a powerful token mixing mechanism, several vision models have explored alternatives to attention for aggregating information across patches. MLP-Mixer performs token mixing and channel mixing using only multilayer perceptrons, without convolution or attention, and achieves competitive performance in image recognition [9], [10]. PoolFormer further argues that the general MetaFormer structure, rather than attention itself, is a key factor for strong visual representation learning, and shows that even a simple pooling operation can serve as an effective token mixer [11]. These studies suggest that self-attention is not necessarily an indispensable component for visual recognition, and that alternative token mixing mechanisms can also learn effective visual features.

However, attention-replacement methods such as MLP-Mixer and PoolFormer mainly focus on simplifying or improving token mixing, rather than independently evaluating the relevance of each patch based on query–key relationships. In contrast, our method computes absolute relevance based on query–key similarity using Screening and explicitly removes low-relevance patches through thresholding. Therefore, our work differs from existing attention-replacement approaches in that it does not merely replace attention with a simpler operation, but addresses the relative-weighting limitation of softmax attention and introduces a selective patch aggregation mechanism based on absolute relevance.

3 Method↩︎

3.1 Overview↩︎

We adapt the Screening mechanism to visual recognition. Our goal is not to redesign the Screening operation itself, but to make it suitable for non-causal two-dimensional image patch grids. As shown in Fig. 1, we replace the self-attention module in a ViT block with a two-dimensional Screening module.

Given an input image, we first divide it into non-overlapping patches and embed them into patch tokens following the standard ViT pipeline. Let \(\mathbf{X} = [\mathbf{x}_1, \ldots, \mathbf{x}_N] \in \mathbb{R}^{N \times d}\) denote the patch tokens, where each patch token has a two-dimensional grid coordinate \(\mathbf{p}_n = (p_n^x, p_n^y)\). The proposed module follows the original Screening mechanism for query, key, value, and gate projections, unit-length normalization, Trim transform, TanhNorm, and gated aggregation. The main modifications are the introduction of axial two-dimensional RoPE and the extension of the distance-aware softmask to two-dimensional space.

Unlike the causal processing of the original Screening mechanism, VisionScreen allows each patch to compute relevance over all patches in the image grid. This non-causal setting is suitable for image recognition, where all patches are available simultaneously.

3.2 Extensions for 2D images↩︎

3.2.1 2D RoPE↩︎

The original Screening mechanism is designed for one-dimensional token sequences. In contrast, image patches are arranged on a two-dimensional grid, and their spatial relationships are important for visual recognition. To incorporate such spatial structure into query–key similarity, we adopt axial two-dimensional Rotary Position Embedding (2D RoPE), following RoPE for ViT [12], which extends RoPE [13] to vision models.

Let \(\mathbf{p}_n=(p_n^x,p_n^y)\) denote the two-dimensional coordinate of the \(n\)-th patch, where \(p_n^x\) and \(p_n^y\) correspond to the horizontal and vertical positions, respectively. For each head, we first obtain normalized query and key vectors \(\bar{\mathbf{q}}_n\) and \(\bar{\mathbf{k}}_n\) as in the original Screening mechanism. We divide their feature dimensions into two equal groups corresponding to the horizontal and vertical axes. Within each group, every pair of feature dimensions is rotated according to the patch coordinate along the corresponding axis.

For an axis \(a \in \{x,y\}\), the rotation applied to the \(t\)-th pair of feature dimensions is defined as \[\mathbf{R}(p_n^a,t) = \begin{bmatrix} \cos(\theta_t p_n^a) & -\sin(\theta_t p_n^a) \\ \sin(\theta_t p_n^a) & \cos(\theta_t p_n^a) \end{bmatrix}, \label{eq:axial95rope}\tag{1}\] where \(p_n^a\) denotes the coordinate of the \(n\)-th patch along axis \(a\). The frequency is defined as \[\theta_t = 10000^{-t/(d_h/4)}, \quad t \in \{0,1,\ldots,d_h/4-1\}, \label{eq:rope95frequency}\tag{2}\] where \(d_h\) denotes the head dimension of the query and key vectors.

Let \(\bar{\mathbf{q}}_{n,a,t}\) and \(\bar{\mathbf{k}}_{n,a,t}\) denote the \(t\)-th pair of query and key dimensions assigned to axis \(a\). The position-encoded components are obtained as \[\tilde{\mathbf{q}}_{n,a,t} = \bar{\mathbf{q}}_{n,a,t} \mathbf{R}(p_n^a,t)^\top, \quad \tilde{\mathbf{k}}_{n,a,t} = \bar{\mathbf{k}}_{n,a,t} \mathbf{R}(p_n^a,t)^\top.\] The rotated components for the two axes together form the position-encoded query and key vectors \(\tilde{\mathbf{q}}_n\) and \(\tilde{\mathbf{k}}_n\). The similarity between query patch \(n\) and key patch \(m\) is then computed as \[s_{nm} = \tilde{\mathbf{q}}_n \tilde{\mathbf{k}}_m^\top.\] Thus, the query–key similarity incorporates relative spatial information along both horizontal and vertical axes.

3.2.2 Spatial Softmask↩︎

The original Screening mechanism uses a distance-aware softmask defined along a one-dimensional causal sequence. We extend this softmask to two-dimensional image space by using the Euclidean distance between patch coordinates. For two patches at \(\mathbf{p}_i\) and \(\mathbf{p}_j\), we define \[d_{ij} = |\mathbf{p}_i - \mathbf{p}_j|_2 .\] Each head has a learnable screening window \[w = \exp(s_w) + 1, \label{eq:screening95window}\tag{3}\] where \(s_w\) is a learnable scalar. Following the original Screening mechanism [8], we use a cosine function to smoothly decrease the spatial weight as the distance from the query patch increases. The weight is equal to one when the two patches are at the same position and gradually decreases to zero at the boundary of the screening window. The two-dimensional spatial softmask is then defined as \[m_{ij} = \begin{cases} \frac{1}{2} \left( \cos \left( \frac{\pi d_{ij}}{w} \right) + 1 \right), & d_{ij} < w, \\ 0, & \mathrm{otherwise}. \end{cases} \label{eq:2d95softmask}\tag{4}\]

This formulation allows each head to learn a spatial screening range over the image plane. To compute the final relevance score, we first apply the same Trim transform as in the original Screening mechanism to the query–key similarity \(s_{ij}\) \[\alpha_{ij} = \operatorname{Trim}(s_{ij}).\] We then apply the spatial softmask to the resulting content-based relevance score \[\alpha_{ij}^s = \alpha_{ij} m_{ij}.\] Unlike softmax attention, \(\alpha_{ij}^s\) is not normalized over keys. Therefore, each patch is selected based on its absolute relevance to the query patch rather than through competition with other patches.

4 Experiments↩︎

4.1 Experimental Setup↩︎

We evaluate VisionScreen on image classification benchmarks. As a baseline, we use ViT-Tiny/16 with an input resolution of \(224 \times 224\). Both ViT-Tiny/16 and VisionScreen follow the same overall model scale: 12 layers, 3 heads, an embedding dimension of 192, and a patch size of \(16 \times 16\). We conduct experiments on ImageNet-1k [14] and CIFAR-100 [15]. All models are trained from scratch on each dataset. For ImageNet-1k, we train models with a global batch size of 1024, while for CIFAR-100 we use a global batch size of 512. All models are trained for 300 epochs, with the first 5 epochs used for learning-rate warmup followed by cosine learning-rate scheduling. We use AdamW as the optimizer. To isolate the effect of the architecture, we avoid advanced data augmentation and regularization techniques. Following simple training settings used in prior work [1], [16], [17], we only apply random resized cropping and random horizontal flipping during training.

4.2 Quantitative Comparison with ViT↩︎

We compare VisionScreen with the ViT-Tiny/16 baseline on ImageNet-1k and CIFAR-100. Table 1 shows the top-1 classification accuracy and the number of model parameters. VisionScreen achieves higher accuracy than ViT on both datasets while using fewer parameters. On ImageNet-1k, VisionScreen improves top-1 accuracy from 68.1% to 72.5%, corresponding to a gain of 4.4 percentage points. On CIFAR-100, VisionScreen improves top-1 accuracy from 50.3% to 52.4%, showing a gain of 2.1 percentage points. These results indicate that Screening-based selective patch aggregation is effective for visual recognition and can provide better recognition performance than softmax attention under a comparable model scale.

In addition to classification accuracy, we also compare the validation loss curves during training. Figure 2 shows the validation loss on ImageNet-1k and CIFAR-100. On both datasets, VisionScreen achieves a lower validation loss than the ViT baseline throughout the later stages of training. This trend is consistent with the top-1 accuracy results in Table 1, suggesting that the improvement of VisionScreen is not limited to the final classification accuracy but is also reflected in the optimization behavior. These results indicate that the proposed Screening-based patch aggregation can learn more effective visual representations than softmax attention under the same training setting.

VisionScreen also uses fewer parameters than the ViT-Tiny/16 baseline. This reduction mainly comes from two design choices. First, VisionScreen uses axial 2D RoPE instead of learnable absolute position embeddings. Second, while ViT-Tiny/16 uses a query and key dimension of 64 for each head, VisionScreen uses a dimension of 16, following the original Screening configuration [8]. Since Screening computes relevance from unit-normalized query and key vectors, the relevance score is based on bounded cosine similarity rather than a high-dimensional dot product followed by softmax normalization. Therefore, we consider that VisionScreen can effectively estimate patch relevance even with lower-dimensional query and key representations.

4.3 Visualization of Patch Interactions↩︎

To further analyze the behavior of VisionScreen, we visualize token-to-token interaction maps and compare them with those of ViT. For ViT, we visualize the softmax attention weights averaged over all heads in the final layer. For VisionScreen, we visualize the relevance scores computed from query–key similarity before value aggregation, also averaged over all heads in the final layer.

We consider two types of visualizations. First, Fig. 3 visualizes the interaction from the classification token (CLS token) to all patch tokens. This visualization shows which image regions contribute to the global representation used for classification. Second, Fig. 4 visualizes the interaction from a center patch to all patch tokens. The selected center patch is marked with a light-blue cross in the figure. This visualization allows us to observe which regions the selected patch refers to when aggregating information from other patches.

Figures 3 and 4 both visualize an image from the Tench class, the first class in ImageNet-1k, which represents a type of freshwater fish. Figure 3 shows the CLS-to-patch interaction maps. Compared with ViT, VisionScreen produces more selective interaction patterns, reflecting the explicit rejection of low-relevance patches by Screening. This suggests that VisionScreen can focus on visually relevant regions without distributing weights across all patches through softmax normalization. Figure 4 shows the interaction maps from the center patch to all patch tokens. In this visualization, we examine how the selected center token attends to other regions during token aggregation. ViT tends to distribute attention over a wide spatial range, sometimes assigning strong weights to distant regions such as image corners or to objects unrelated to the target class. In the visualized example, ViT assigns relatively strong weights to fishing equipment, such as a fishing rod or reel, rather than to the fish itself. In contrast, VisionScreen focuses more consistently on the fish region, indicating that it aggregates information from semantically relevant patches. This behavior suggests that VisionScreen is less prone to relying on spurious correlations and instead captures more class-relevant visual evidence.

These visualizations qualitatively support the quantitative results, showing that the performance improvement of VisionScreen is accompanied by more selective and semantically meaningful patch aggregation behavior.

4.4 Analysis of Spatial Softmask↩︎

To further analyze the behavior of the proposed spatial softmask, we first visualize the spatial softmasks generated from the learned screening window parameters in the final layer of VisionScreen trained on ImageNet-1k. Figure 5 shows the spatial distribution of the softmask values for each head, with the selected patch marked by a red \(\times\). The screening window defined in Eq. 3 controls the spatial extent of the softmask in Eq. 4 .

As shown in Fig. 5, the spatial extents of the softmasks vary substantially across heads. Some heads assign nonzero weights only to a relatively local region around the selected patch, while others cover a much broader area and capture more global context. This result suggests that the spatial softmask does not impose a uniform spatial range across heads, but instead allows different heads to specialize in different spatial scales. Smaller screening windows can capture fine-grained neighborhood information, whereas larger windows can incorporate long-range dependencies across the image. Therefore, the learned screening windows allow VisionScreen to adaptively balance local and global information through its multi-head design.

We further analyze the learned screening window values across all layers and heads. Fig. 6 shows the learned window size of each head, where the horizontal axis denotes the layer index and the vertical axis denotes the screening window size. The windows in the final layer correspond to the visualization shown in Fig. 5. Considering Figs. 5 and 6 together with the definition of the spatial softmask in Eq. 4 , we observe that when the screening window becomes sufficiently large, approximately above 1000 in our setting, the softmask assigns nearly uniform weights across almost the entire patch grid. In this case, the head behaves closer to global aggregation, similar to the behavior of standard ViT attention without an explicit spatial restriction. However, Fig. 6 shows that only a small number of heads have such large window values. Most heads retain spatially limited windows, indicating that VisionScreen does not simply degenerate into global aggregation.

Another observation is that both shallow and deep layers contain a mixture of local and global heads. While deeper layers in ViT are often considered to mainly aggregate global information, our results suggest that local patch aggregation can remain important even in deeper layers. This indicates that VisionScreen adaptively maintains multiple spatial scales throughout the network, rather than monotonically increasing the receptive range with depth.

4.5 Ablation Study↩︎

We further conduct an ablation study to analyze the effect of the gating mechanism in VisionScreen. The gate consists of a linear projection followed by SiLU and Tanh activations, and its output is multiplied with the screened representation before the final projection. This design follows the original Screening mechanism rather than being a newly introduced component in our method. Nevertheless, since the gate directly modulates the aggregated features, we examine its effect on visual recognition performance.

Table 2 shows the ImageNet-1k top-1 accuracy with and without the gating mechanism. Removing the gate decreases the accuracy from 72.5% to 71.8%, resulting in a 0.7 percentage point drop. This result indicates that the gating mechanism contributes to the performance of VisionScreen. We consider that the gate adaptively modulates the aggregated patch features on a channel-wise basis, suppressing less useful responses while emphasizing informative ones. Since the model still outperforms ViT without the gate, we regard it as an auxiliary mechanism that enhances representational capacity rather than an essential component of Screening.

5 Conclusion↩︎

We proposed VisionScreen, a vision model that extends the Screening mechanism to two-dimensional image patches. By assigning an absolute relevance value to each patch pair, VisionScreen can explicitly reject low-relevance patches, unlike softmax attention, which aggregates features through relative weighting. Experiments on ImageNet-1k and CIFAR-100 show that VisionScreen outperforms ViT-Tiny/16 with fewer parameters, while achieving more stable training and more selective patch interactions.

Future work includes applying VisionScreen to dense prediction tasks and further analyzing its relevance patterns and patch selection behavior.

References↩︎