# A Manifold Representation of the Key in Vision Transformers

arXiv: https://arxiv.org/abs/2402.00534 · HTML: https://mlsanity.com/paper/2402.00534/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 40 references that are in the map.
- First posted on arXiv 1 Feb 2024.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Mask R-CNN](https://mlsanity.com/paper/1703.06870/) (2017)
- [Feature Pyramid Networks for Object Detection](https://mlsanity.com/paper/1612.03144/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [MMDetection: Open MMLab Detection Toolbox and Benchmark](https://mlsanity.com/paper/1906.07155/) (2019)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [Transformer in Transformer](https://mlsanity.com/paper/2103.00112/) (2021)
- [DeiT III: Revenge of the ViT](https://mlsanity.com/paper/2204.07118/) (2022)
- [Vision Transformer with Deformable Attention](https://mlsanity.com/paper/2201.00520/) (2022)
- [FLatten Transformer: Vision Transformer using Focused Linear Attention](https://mlsanity.com/paper/2308.00442/) (2023)
- [The geometry of hidden representations of large transformer models](https://mlsanity.com/paper/2302.00294/) (2023)
- [RMT: Retentive Networks Meet Vision Transformers](https://mlsanity.com/paper/2309.11523/) (2023)
- [KVT: k-NN Attention for Boosting Vision Transformers](https://mlsanity.com/paper/2106.00515/) (2021)
- [Learning to Merge Tokens in Vision Transformers](https://mlsanity.com/paper/2202.12015/) (2022)
- [TransNeXt: Robust Foveal Visual Perception for Vision Transformers](https://mlsanity.com/paper/2311.17132/) (2023)
- [Learning Efficient Vision Transformers via Fine-Grained Manifold Distillation](https://mlsanity.com/paper/2107.01378/) (2021)
- [Refiner: Refining Self-attention for Vision Transformers](https://mlsanity.com/paper/2106.03714/) (2021)
- [Improving Transformers with Probabilistic Attention Keys](https://mlsanity.com/paper/2110.08678/) (2021)
- [Learned Queries for Efficient Local Attention](https://mlsanity.com/paper/2112.11435/) (2021)
- [Robustifying Token Attention for Vision Transformers](https://mlsanity.com/paper/2303.11126/) (2023)
- [Scratching Visual Transformer's Back with Uniform Attention](https://mlsanity.com/paper/2210.08457/) (2022)
- [Multi-manifold Attention for Vision Transformers](https://mlsanity.com/paper/2207.08569/) (2022)
- [State Representation Learning Using an Unbalanced Atlas](https://mlsanity.com/paper/2305.10267/) (2023)

## Built on by (most cited first, 2 of 2)

- [MusRec: Zero-Shot Text-to-Music Editing via Rectified Flow and Diffusion Transformers](https://mlsanity.com/paper/2511.04376/) (2025)
- [FlowSonic: Stable Zero-Shot Music Editing via High-Order Trajectory Integration](https://mlsanity.com/paper/2607.17526/) (2026)

## Abstract (arXiv)

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.