Salient Mask-Guided Vision Transformer for Fine-Grained Classification

TL;DR

This paper tackles fine-grained visual classification (FGVC) by addressing the tendency of Vision Transformers (ViT) to attend to background regions. It introduces SM-ViT, which integrates a salient object detector to generate a mask and crops foreground content, then uses a Salient Mask-Guided Encoder (SMGE) to bias the class-token attention toward salient patches via a controlled augmentation: the class-token attention scores are increased by $ (x_{max} \cdot d_{\theta})$ for masked tokens, implemented as $x_{scor_{cls}}^{i} \leftarrow x_{scor_{cls}}^{i} + (x_{max} \cdot d_{\theta})$ when $m_i=1$, with $Y = \text{softmax}(X'_{scor})$ propagated across layers. This approach requires no extra trainable parameters and improves FGVC accuracy while using lower input resolutions. Empirically, SM-ViT yields up to about 1 percentage point gains on Stanford Dogs, CUB, and NABirds over vanilla ViT, with additional benefits such as fewer trainable parameters and faster inference when SMGE is disabled. The method demonstrates strong generalization across datasets and suggests that saliency-guided attention is a promising direction for improving FGVC with transformer-based architectures, while remaining compatible with existing backbones and saliency detectors.

Abstract

Fine-grained visual classification (FGVC) is a challenging computer vision problem, where the task is to automatically recognise objects from subordinate categories. One of its main difficulties is capturing the most discriminative inter-class variances among visually similar classes. Recently, methods with Vision Transformer (ViT) have demonstrated noticeable achievements in FGVC, generally by employing the self-attention mechanism with additional resource-consuming techniques to distinguish potentially discriminative regions while disregarding the rest. However, such approaches may struggle to effectively focus on truly discriminative regions due to only relying on the inherent self-attention mechanism, resulting in the classification token likely aggregating global information from less-important background patches. Moreover, due to the immense lack of the datapoints, classifiers may fail to find the most helpful inter-class distinguishing features, since other unrelated but distinctive background regions may be falsely recognised as being valuable. To this end, we introduce a simple yet effective Salient Mask-Guided Vision Transformer (SM-ViT), where the discriminability of the standard ViT`s attention maps is boosted through salient masking of potentially discriminative foreground regions. Extensive experiments demonstrate that with the standard training procedure our SM-ViT achieves state-of-the-art performance on popular FGVC benchmarks among existing ViT-based approaches while requiring fewer resources and lower input image resolution.

Figures (4)

• Figure 1: Visualised attention performance comparison of vanilla ViT (first row) and our SM-ViT (second row). For ViT we demonstrate the averaged attention map for the final class token. While for SM-ViT we first show the extracted saliency mask from the salient object detection module, and then the final class token's averaged attention map augmented according to this mask.
• Figure 2: The overall architecture of our proposed SM-ViT. An image is first fed into the salient object detection module to extract its saliency mask and automatically generate a bounding box, which are then used to prepare a binary mask and to crop the image respectively. Further, the cropped image is fed into the ViT-like architecture, where it is first split into patches, projected into the embedding space, the positional embedding is added to the patches, and a class token is prepended. Next, the resulted sequence of tokens is passed through each layer of our Salient Mask-Guided Encoder (SMGE), where inside the multi-head self-attention mechanism the flatten binary mask is used to augment attention scores of the class token accordingly. Lastly, the class token values from the last SMGE layer are passed to a classification head to perform categorisation.
• Figure 3: Visualisation of vanilla ViT and our SM-ViT results on different datasets. The first row shows original images, while the second and third rows demonstrate averaged by all heads attention maps generated by the class token at the final encoder layer of vanilla ViT and SM-ViT, respectively. Brightness intensity represents the total amount of attention, where the more attention the class token pays to a region, the brighter it is, and the other way around.
• Figure 4: Visual comparison of class token averaged attention maps at the last encoder layer of our SM-ViT, first fine-tuned with SMGE, and then used with disabled (first row) and enabled (second row) SMGE during inference.