Multi-manifold Attention for Vision Transformers

TL;DR

A novel attention mechanism, called multi-manifold multi-head attention, is proposed in this work to substitute the vanilla self-attention of a Transformer, and can guide a Vision Transformer to become more attentive towards important appearance, color and texture features of an image, leading to improved classification and segmentation results.

Abstract

Vision Transformers are very popular nowadays due to their state-of-the-art performance in several computer vision tasks, such as image classification and action recognition. Although their performance has been greatly enhanced through highly descriptive patch embeddings and hierarchical structures, there is still limited research on utilizing additional data representations so as to refine the selfattention map of a Transformer. To address this problem, a novel attention mechanism, called multi-manifold multihead attention, is proposed in this work to substitute the vanilla self-attention of a Transformer. The proposed mechanism models the input space in three distinct manifolds, namely Euclidean, Symmetric Positive Definite and Grassmann, thus leveraging different statistical and geometrical properties of the input for the computation of a highly descriptive attention map. In this way, the proposed attention mechanism can guide a Vision Transformer to become more attentive towards important appearance, color and texture features of an image, leading to improved classification and segmentation results, as shown by the experimental results on well-known datasets.

Figures (8)

• Figure 1: With the modelling of input in the Euclidean $E^d$, Grassmann $\mathcal{G}(p,d)$ and SPD $\mathcal{S}_{++}^d$ manifolds and the computation of the respective distance maps, the proposed MMA produces a fused attention map with high discriminative power.
• Figure 2: Proposed MMA architecture. The symbol $\otimes$ denotes matrix multiplication, while $BN$ is batch normalization.
• Figure 3: Class activations in C-100 using MMA-ViT-Lite-6/4. (a) Original images and class activation maps using (b) Euclidean, (c) SPD, (d) Grassmann and (e) all three manifolds.
• Figure 4: Class activations in T-ImageNet. (a) Original images and class activation maps using (b) CVT-6/4, (c) MMA-CVT-6/4, (d) Swin-T, (e) MMA-Swin-T, (f) CCT-7-3x2 and (g) MMA-CCT-7-3x2 models.
• Figure 5: Impact of MMA on tested ViT variants trained on C-100 and T-ImageNet in terms of model performance (MP) and generalization ability (GA).