# Improving Transformers with Probabilistic Attention Keys

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 79 of its 110 references that are in the map.
- First posted on arXiv 16 Oct 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [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)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)

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

- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Architecture Agnostic Federated Learning for Neural Networks](https://mlsanity.com/paper/2202.07757/) (2022)
- [Primal-Attention: Self-attention through Asymmetric Kernel SVD in Primal Representation](https://mlsanity.com/paper/2305.19798/) (2023)
- [On Efficient Training of Large-Scale Deep Learning Models: A Literature Review](https://mlsanity.com/paper/2304.03589/) (2023)
- [Designing Robust Transformers using Robust Kernel Density Estimation](https://mlsanity.com/paper/2210.05794/) (2022)
- [OTSeg: Multi-prompt Sinkhorn Attention for Zero-Shot Semantic Segmentation](https://mlsanity.com/paper/2403.14183/) (2024)
- [Efficient Degradation-agnostic Image Restoration via Channel-Wise Functional Decomposition and Manifold Regularization](https://mlsanity.com/paper/2505.18679/) (2025)
- [Unveiling the Hidden Structure of Self-Attention via Kernel Principal Component Analysis](https://mlsanity.com/paper/2406.13762/) (2024)
- [Elliptical Attention](https://mlsanity.com/paper/2406.13770/) (2024)
- [Devil is in the Uniformity: Exploring Diverse Learners within Transformer for Image Restoration](https://mlsanity.com/paper/2503.20174/) (2025)
- [DBA: Efficient Transformer with Dynamic Bilinear Low-Rank Attention](https://mlsanity.com/paper/2211.16368/) (2022)
- [Attention: Marginal Probability is All You Need?](https://mlsanity.com/paper/2304.04556/) (2023)
- [A Primal-Dual Framework for Transformers and Neural Networks](https://mlsanity.com/paper/2406.13781/) (2024)
- [Transformer Meets Twicing: Harnessing Unattended Residual Information](https://mlsanity.com/paper/2503.00687/) (2025)
- [Transformer with Fourier Integral Attentions](https://mlsanity.com/paper/2206.00206/) (2022)
- [A Manifold Representation of the Key in Vision Transformers](https://mlsanity.com/paper/2402.00534/) (2024)
- [Generalization Guarantees for Representation Learning via Data-Dependent Gaussian Mixture Priors](https://mlsanity.com/paper/2502.15540/) (2025)
- [Reasoning is Periodicity? Improving Large Language Models Through Effective Periodicity Modeling](https://mlsanity.com/paper/2502.21309/) (2025)
- [PSA-MIL: A Probabilistic Spatial Attention-Based Multiple Instance Learning for Whole Slide Image Classification](https://mlsanity.com/paper/2503.16284/) (2025)
- [A Semantic and Motion-Aware Spatiotemporal Transformer Network for Action Detection](https://mlsanity.com/paper/2405.08204/) (2024)
- [Multiview Representation Learning via Distributed Joint Latent Space Structuring](https://mlsanity.com/paper/2504.18455/) (2025)
- [Robust Filter Attention: Self-Attention as Precision-Weighted State Estimation](https://mlsanity.com/paper/2509.04154/) (2025)
- [QUEST: A robust attention formulation using query-modulated spherical attention](https://mlsanity.com/paper/2604.00199/) (2026)
- [Stochastic Vision Transformers with Wasserstein Distance-Aware Attention](https://mlsanity.com/paper/2311.18645/) (2023)
- [EventSSEG: Event-driven Self-Supervised Segmentation with Probabilistic Attention](https://mlsanity.com/paper/2508.14856/) (2025)
- [ForecastGAN: A Decomposition-Based Adversarial Framework for Multi-Horizon Time Series Forecasting](https://mlsanity.com/paper/2511.04445/) (2025)
- [Zonkey: A Hierarchical Diffusion Language Model with Differentiable Tokenization and Probabilistic Attention](https://mlsanity.com/paper/2601.21768/) (2026)
- [One for All: A Non-Linear Transformer can Enable Cross-Domain Generalization for In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.09727/) (2026)

## Abstract (arXiv)

Multi-head attention is a driving force behind state-of-the-art transformers, which achieve remarkable performance across a variety of natural language processing (NLP) and computer vision tasks. It has been observed that for many applications, those attention heads learn redundant embedding, and most of them can be removed without degrading the performance of the model. Inspired by this observation, we propose Transformer with a Mixture of Gaussian Keys (Transformer-MGK), a novel transformer architecture that replaces redundant heads in transformers with a mixture of keys at each head. These mixtures of keys follow a Gaussian mixture model and allow each attention head to focus on different parts of the input sequence efficiently. Compared to its conventional transformer counterpart, Transformer-MGK accelerates training and inference, has fewer parameters, and requires fewer FLOPs to compute while achieving comparable or better accuracy across tasks. Transformer-MGK can also be easily extended to use with linear attention. We empirically demonstrate the advantage of Transformer-MGK in a range of practical applications, including language modeling and tasks that involve very long sequences. On the Wikitext-103 and Long Range Arena benchmark, Transformer-MGKs with 4 heads attain comparable or better performance to the baseline transformers with 8 heads.