Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Elliptical Attention

Stefan K. Nielsen, Laziz U. Abdullaev, Rachel S. Y. Teo, Tan M. Nguyen

TL;DR

The advantages of Elliptical Attention over the baseline dot-product attention and state-of-the-art attention methods on various practical tasks, including object classification, image segmentation, and language modeling across different data modalities are empirically demonstrated.

Abstract

Pairwise dot-product self-attention is key to the success of transformers that achieve state-of-the-art performance across a variety of applications in language and vision. This dot-product self-attention computes attention weights among the input tokens using Euclidean distance, which makes the model prone to representation collapse and vulnerable to contaminated samples. In this paper, we propose using a Mahalanobis distance metric for computing the attention weights to stretch the underlying feature space in directions of high contextual relevance. In particular, we define a hyper-ellipsoidal neighborhood around each query to increase the attention weights of the tokens lying in the contextually important directions. We term this novel class of attention Elliptical Attention. Our Elliptical Attention provides two benefits: 1) reducing representation collapse and 2) enhancing the model's robustness as Elliptical Attention pays more attention to contextually relevant information rather than focusing on some small subset of informative features. We empirically demonstrate the advantages of Elliptical Attention over the baseline dot-product attention and state-of-the-art attention methods on various practical tasks, including object classification, image segmentation, and language modeling across different data modalities.

Elliptical Attention

TL;DR

The advantages of Elliptical Attention over the baseline dot-product attention and state-of-the-art attention methods on various practical tasks, including object classification, image segmentation, and language modeling across different data modalities are empirically demonstrated.

Abstract

Pairwise dot-product self-attention is key to the success of transformers that achieve state-of-the-art performance across a variety of applications in language and vision. This dot-product self-attention computes attention weights among the input tokens using Euclidean distance, which makes the model prone to representation collapse and vulnerable to contaminated samples. In this paper, we propose using a Mahalanobis distance metric for computing the attention weights to stretch the underlying feature space in directions of high contextual relevance. In particular, we define a hyper-ellipsoidal neighborhood around each query to increase the attention weights of the tokens lying in the contextually important directions. We term this novel class of attention Elliptical Attention. Our Elliptical Attention provides two benefits: 1) reducing representation collapse and 2) enhancing the model's robustness as Elliptical Attention pays more attention to contextually relevant information rather than focusing on some small subset of informative features. We empirically demonstrate the advantages of Elliptical Attention over the baseline dot-product attention and state-of-the-art attention methods on various practical tasks, including object classification, image segmentation, and language modeling across different data modalities.
Paper Structure (62 sections, 13 theorems, 74 equations, 7 figures, 21 tables, 1 algorithm)

This paper contains 62 sections, 13 theorems, 74 equations, 7 figures, 21 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background: Self-Attention and Non-Parametric Regression
  3. Self-Attention Mechanism
  4. A Non-Parametric Regression Perspective of Self-Attention
  5. Elliptical Attention: Leveraging Hyper-Ellipsoids to Pay More Attention Without Losing Focus
  6. Improving NW Regression with Hyper-Ellipsoids
  7. Capturing Coordinate-wise Variability and Building the Mahalanobis Transformation
  8. Promoting Robustness and Avoiding Representation Collapse
  9. An Efficient Estimator of the Coordinate-wise Variability
  10. Full Technical Formulation of Elliptical Attention
  11. Experimental Results
  12. Robust Language Modelling
  13. Image Classification under Adversarial Attack
  14. Long Sequence Modelling on the LRA Benchmark
  15. Image Segmentation on ADE20K
  16. ...and 47 more sections

Key Result

Lemma 1

Let $\mathcal{M} : \mathbb R^D \to \mathbb R^N$ denote the transformed Elliptical $\mathrm{softmax}$ operator for a given set of keys as $\mathcal{M}(\boldsymbol{x}) := \frac{1}{\sum_{j\in [N]} \exp(\boldsymbol{x}^\top \boldsymbol{M} \boldsymbol{k}_j)} \left[ \exp(\boldsymbol{x}^\top \boldsymbol{M}

Figures (7)

  • Figure 1: Comparison of Attention Heatmaps. Elliptical pays attention to more relevant information. DeiT focuses on just a subset of informative features while Elliptical considers a wider set of contextually relevant information, helping to produce more accurate and robust predictions. Attention scores are min-max scaled for visualization purposes.
  • Figure 2: Left: The function does not vary in the $x_2$ axis so we stretch the neighborhood in that direction. Right: The stretched ellipsoidal neighborhood includes 4 more keys.
  • Figure 3: Representation Collapse on WikiText-103. Elliptical Attention learns more diverse representations.
  • Figure 4: ImageNet Efficiency: Comparison of throughput and max memory allocated for DeiT, Elliptical, RVT, RKDE, MoM on Tiny, Small, and Base sizes. Elliptical is the most efficient robust model. Numerical analysis in Table \ref{['table: efficiency analysis']} of Appendix \ref{['sec: appendix-experimental details']}.
  • Figure 5: Left: Evolution of mean values of key perturbations over successive layers. Right: Mean key perturbations at different layers after 300 epochs. The figures show that as the number of layers increases, mean key perturbations over layers stabilize around a constant value.
  • ...and 2 more figures

Theorems & Definitions (24)

  • Definition 1: Coordinate-wise Variability of $f: \mathbb{R}^D \rightarrow \mathbb{R}^{D_v}$
  • Remark 1
  • Remark 2
  • Remark 3
  • Lemma 1
  • Proposition 1: Robustness of Elliptical Attention
  • Proposition 2: Elliptical Attention maintains expressive power by reducing noise
  • Proposition 3: Coordinate-wise Variability Estimator
  • Remark 4
  • Definition 2: Elliptical Attention Computation
  • ...and 14 more