Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Manifold Representation of the Key in Vision Transformers

Li Meng, Morten Goodwin, Anis Yazidi, Paal Engelstad

TL;DR

This work investigates decoupling the key from the query and value in vision transformers by endowing the key with a manifold representation across multiple charts. The authors expand the key to H charts using a trainable scale, process per-chart features with group-wise 1D convolutions, and then aggregate via several strategies (SpatialK, KUA, SimpleK, VanillaK), optionally incorporating context broadcasting. Across ViT and Swin Transformer families on ImageNet-1K and COCO, the approach yields consistent gains, notably ViT-B (+0.87 percentage points) and Swin-T (+0.52 percentage points) on ImageNet-1K, with VanillaK excelling for ViTs and SpatialK for Swin; object detection and instance segmentation also benefit. The results demonstrate that improved key representations can enhance performance beyond simply increasing parameters or FLOPs, while also highlighting the need for budget-aware designs and further exploration of efficient manifold-based attention mechanisms. This work paves the way for more nuanced key representations in vision transformers and motivates future work on reducing computational cost while maintaining accuracy gains.

Abstract

Vision Transformers implement multi-head self-attention via stacking multiple attention blocks. The query, key, and value are often intertwined and generated within those blocks via a single, shared linear transformation. This paper explores the concept of disentangling the key from the query and value, and adopting a manifold representation for the key. Our experiments reveal that decoupling and endowing the key with a manifold structure can enhance the model's performance. Specifically, ViT-B exhibits a 0.87% increase in top-1 accuracy, while Swin-T sees a boost of 0.52% in top-1 accuracy on the ImageNet-1K dataset, with eight charts in the manifold key. Our approach also yields positive results in object detection and instance segmentation tasks on the COCO dataset. We establish that these performance gains are not merely due to the simplicity of adding more parameters and computations. Future research may investigate strategies for cutting the budget of such representations and aim for further performance improvements based on our findings.

A Manifold Representation of the Key in Vision Transformers

TL;DR

This work investigates decoupling the key from the query and value in vision transformers by endowing the key with a manifold representation across multiple charts. The authors expand the key to H charts using a trainable scale, process per-chart features with group-wise 1D convolutions, and then aggregate via several strategies (SpatialK, KUA, SimpleK, VanillaK), optionally incorporating context broadcasting. Across ViT and Swin Transformer families on ImageNet-1K and COCO, the approach yields consistent gains, notably ViT-B (+0.87 percentage points) and Swin-T (+0.52 percentage points) on ImageNet-1K, with VanillaK excelling for ViTs and SpatialK for Swin; object detection and instance segmentation also benefit. The results demonstrate that improved key representations can enhance performance beyond simply increasing parameters or FLOPs, while also highlighting the need for budget-aware designs and further exploration of efficient manifold-based attention mechanisms. This work paves the way for more nuanced key representations in vision transformers and motivates future work on reducing computational cost while maintaining accuracy gains.

Abstract

Vision Transformers implement multi-head self-attention via stacking multiple attention blocks. The query, key, and value are often intertwined and generated within those blocks via a single, shared linear transformation. This paper explores the concept of disentangling the key from the query and value, and adopting a manifold representation for the key. Our experiments reveal that decoupling and endowing the key with a manifold structure can enhance the model's performance. Specifically, ViT-B exhibits a 0.87% increase in top-1 accuracy, while Swin-T sees a boost of 0.52% in top-1 accuracy on the ImageNet-1K dataset, with eight charts in the manifold key. Our approach also yields positive results in object detection and instance segmentation tasks on the COCO dataset. We establish that these performance gains are not merely due to the simplicity of adding more parameters and computations. Future research may investigate strategies for cutting the budget of such representations and aim for further performance improvements based on our findings.
Paper Structure (16 sections, 5 equations, 2 figures, 3 tables, 1 algorithm)

This paper contains 16 sections, 5 equations, 2 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Related Work
  4. Tokens
  5. Self-attention
  6. Queries
  7. Keys
  8. Manifolds
  9. Method
  10. Experimental Details
  11. Results
  12. Charts and CB
  13. Visualization
  14. Object Detection
  15. Discussion
  16. ...and 1 more sections

Figures (2)

  • Figure 1: Overview of the proposed architecture. Instead of using a linear layer to generate the key from the input, we expand the input into multiple charts and process these through multilayer perceptrons (MLPs) or MLP-like structures. These representations are then fused to form a meta-representation that serves as the key.
  • Figure 2: Attention maps of the last layer for each head. We compute the joint attention for each head separately using ViT-S/16 models pre-trained on ImageNet-1K. The first row for each original image contains the attention maps of the baseline ViT-S model, and the second row contains the attention maps of the ViT-S model with VanillaK.