# Compositional Attention: Disentangling Search and Retrieval

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 52 references that are in the map.
- First posted on arXiv 18 Oct 2021.

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

- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [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)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [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)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [A simple neural network module for relational reasoning](https://mlsanity.com/paper/1706.01427/) (2017)
- [Measuring Compositional Generalization: A Comprehensive Method on Realistic Data](https://mlsanity.com/paper/1912.09713/) (2019)
- [Inductive Biases for Deep Learning of Higher-Level Cognition](https://mlsanity.com/paper/2011.15091/) (2020)
- [Compositionality decomposed: how do neural networks generalise?](https://mlsanity.com/paper/1908.08351/) (2019)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Compositional Attention Networks for Machine Reasoning](https://mlsanity.com/paper/1803.03067/) (2018)
- [A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms](https://mlsanity.com/paper/1901.10912/) (2019)
- [Investigating the Limitations of Transformers with Simple Arithmetic Tasks](https://mlsanity.com/paper/2102.13019/) (2021)
- [The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers](https://mlsanity.com/paper/2108.12284/) (2021)
- [The Consciousness Prior](https://mlsanity.com/paper/1709.08568/) (2017)
- [Coordination Among Neural Modules Through a Shared Global Workspace](https://mlsanity.com/paper/2103.01197/) (2021)

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

- [Is a Modular Architecture Enough?](https://mlsanity.com/paper/2206.02713/) (2022)
- [On Neural Architecture Inductive Biases for Relational Tasks](https://mlsanity.com/paper/2206.05056/) (2022)
- [A Complexity-Based Theory of Compositionality](https://mlsanity.com/paper/2410.14817/) (2024)
- [Induced Natural Language Rationales and Interleaved Markup Tokens Enable Extrapolation in Large Language Models](https://mlsanity.com/paper/2208.11445/) (2022)
- [Breaking Neural Network Scaling Laws with Modularity](https://mlsanity.com/paper/2409.05780/) (2024)
- [Associative Transformer](https://mlsanity.com/paper/2309.12862/) (2023)
- [Modularity in Deep Learning: A Survey](https://mlsanity.com/paper/2310.01154/) (2023)
- [Layer-wise Representation Fusion for Compositional Generalization](https://mlsanity.com/paper/2307.10799/) (2023)
- [Attention-based Iterative Decomposition for Tensor Product Representation](https://mlsanity.com/paper/2406.01012/) (2024)
- [Leveraging Synthetic Targets for Machine Translation](https://mlsanity.com/paper/2305.06155/) (2023)
- [A Short Survey of Systematic Generalization](https://mlsanity.com/paper/2211.11956/) (2022)
- [Representing Signs as Signs: One-Shot ISLR to Facilitate Functional Sign Language Technologies](https://mlsanity.com/paper/2502.20171/) (2025)

## Abstract (arXiv)

Multi-head, key-value attention is the backbone of the widely successful Transformer model and its variants. This attention mechanism uses multiple parallel key-value attention blocks (called heads), each performing two fundamental computations: (1) search - selection of a relevant entity from a set via query-key interactions, and (2) retrieval - extraction of relevant features from the selected entity via a value matrix. Importantly, standard attention heads learn a rigid mapping between search and retrieval. In this work, we first highlight how this static nature of the pairing can potentially: (a) lead to learning of redundant parameters in certain tasks, and (b) hinder generalization. To alleviate this problem, we propose a novel attention mechanism, called Compositional Attention, that replaces the standard head structure. The proposed mechanism disentangles search and retrieval and composes them in a dynamic, flexible and context-dependent manner through an additional soft competition stage between the query-key combination and value pairing. Through a series of numerical experiments, we show that it outperforms standard multi-head attention on a variety of tasks, including some out-of-distribution settings. Through our qualitative analysis, we demonstrate that Compositional Attention leads to dynamic specialization based on the type of retrieval needed. Our proposed mechanism generalizes multi-head attention, allows independent scaling of search and retrieval, and can easily be implemented in lieu of standard attention heads in any network architecture.