Compositional Attention: Disentangling Search and Retrieval
Authors: Sarthak Mittal
Sharath Chandra Raparthy
Irina Rish
Yoshua Bengio
Guillaume Lajoie
Venue: ICLR 2022 Spotlight (2022)
Keywords: compositional attention
flexible search and retrieval
better generalization

Abstract: Multi-head, key-value attention is the backbone of transformer-like model architectures which have proven to be widely successful in recent years. 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 interaction, and (2) retrieval - extraction of relevant features from the selected entity via a value matrix. 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 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 is easy to implement in a variety of established network architectures.

Published as a conference paper at ICLR 2022
COMPOSITIONAL ATTENTION: DISENTANGLING SEARCH AND RETRIEVAL
Sarthak Mittall, Sharath Chandra Raparthy, Irina Rish, Yoshua Bengio, Guillaume Lajoie! Mila, Université de Montréal
ABSTRACT
Multi-head, key-value attention is the backbone of the widely successful Trans- former model and its variants. This attention mechanism uses multiple parallel key-value attention blocks (called heads), each performing two fundamental com- putations: (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 map- ping 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 pro- pose novel attention mechanism, called Compositional Attention, that replaces the standard head structure. The proposed mechanism disentangles search and retrieval and composes them in dynamic, flexible and context-dependent manner through an additional soft competition stage between the query-key combination and value pairing. Through series of numerical experiments, outperforms standard multi-head attention on a variety of tasks, including some aat-of-distribbtion 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."
show that it
INTRODUCTION
Attention mechanisms have become integral parts of machine learning models across a variety of domains. The modern notion of soft-attention was first introduced in Bahdanau al. (2015) for machine translation to allow recurrent networks to perform well over long sequences. Since then, attention has taken center stage in several models that forego recurrent networks altogether (i.e. Transformers ( (Vaswani et al., 2017)), and has been leveraged in wide variety of applications, like natural language (Bahdanau et al., 2015; Vaswani et al., 2017: Devlin et al., 2018), computer vision (Dosovitskiy et al., 2020) and physical reasoning (Ding et al., 2020; Locatello et al., 2020).
At the core of this success a simple idea: enable task-driven flexible connections between elements of sequence to extract and merge information. This process is implemented by attention (or alignment) functions which, in their simplest form, take a reference or query entity and "pick" (i.e. attend to) the most relevant input entities in set of other entities. Modern attention systems refine this key principle in two meaningful ways. First, they utilize key-value attention, where the attention function takes "queries" from the reference entity, matches them to "keys" attached to input entities, and returns "values" representing a transformation of the selected entities. Second, they allow multiple attention mechanisms to in parallel, often called attention heads, allowing the model to attend to multiple entities jointly, leading to increased expressivity. Despite these advances, Transformer- like architectures still struggle on certain tasks (Fan et al., 2020; Nogueira et al., 2021). including context-sensitive associations and cnttof-distributiio (OoD) generalization (Lake & Baroni, 2018b; 8b; Liska al., 2018). They are still far from human-level performance on physical reasoning and object-centric tasks (Webb et al., 2020). In an object-oriented world where entities have several attributes, current multi-head attention mechanisms learn rigid search-retrieval associations which lead to various limitations, as illustrated in Figure I and Section 2.3.
'Correspondence authors sarthmit@ gmail.com .com and g.lajoic@umontrcccc Spunnesurreed implementation available http://ghhiicooomattttiiioooomoo

---

Published as a conference paper at ICLR 2022
What shape the object?
What location the
What the location the purple object?
Data Setting
Color: Red
Color: Green
Colar: Purple
Shape: Cube
Shape: Cone
Stape: Cylinder
Top
Location: Middle
Location: Rotiom
Head
Head
Head
Need
Purameess
Multi-Head Attention
Query Key
No flexiilitt,
Query Key
Search accordlng
Query Key
Rerriene
Value 3
Search according
Regrieve
Sworch acvordinng to shape
to color
location Botion
no cofor
whape Cube
location Middle
Mapping doww'' sas
Existe
Eniets
Head
Search
Retrieval
Search 2
Retrieval
Search
Retrieval 2
Don't need new
Retrieve
shape
panunctens!
Retriest shige
Search accarding
shape Cude
Value 1
Value
Compositional Attention
Value
Query Key
Query Key I
Query
Fleublity
Search acconting
Value
Sourch according
shape
Retriew
Retrieve
Retrieve location
Locution Middle
Maysing cauo
lncation Bottow
Existr
(n)
(b)
Search
(c)
Value
Figure I: Motivation behind Compositional Attention. In a visual question answering setting, we see that the "rigid" mapping between search (query-key) and retrieval (value) in multi-head attention leads to redundant parameters being learned (middle row; (c)). In contrast, when the search and retrieval mechanisms are disentangled and have pairing set dynamically, these can be composed efficiently without learning of redundant parameters (bottom row; (c)). For details, refer Section 2.3
Addressing these shortcomings, there several recent attention-enables systems developed allow better decomposition and re-composition of knowledge (Goyal et al., 2019; 2021a:b), some of which discuss in Appendix A. However, most of these efforts hinge on purpose-buil architectural components that remain niche and often are difficult to implement at scale. To complement these efforts and build on the proven efficacy of Transformers, our goal is develop minimal modifications to key-value attention to enable flexible decomposition of computations found in attention heads, and eliminate some parameter redundancy. Crucially, we aim for mechanism that is easily implemented and plug-and-play for existing Transformers (and all the models based on them).
propose Compositional Attention, where the search and retrieval operations can flexibly com- posed: the key-query search mechanism is no longer bound to a fixed value retrieval matrix, instead it dynamically selected from a shared pool of value matrices accessible by several compositional attention heads. This results in increased flexibility and improved performance.
Contributions Summary. (a) We formally describe the shortcomings rigid cearch-and-retriewe coupling in standard multi-head attention and empirically analyze them through experiments on an illustrative synthetic task (Section 2.3 and 4.1). (b) We present Compositional Attention to disentangle search and retrieval, and validate its advantages with number of experiments (Section 3 and ). (c) Through a series of analyses, we demonstrate how our proposed attention mechanism decomposes relational task structure expected, and facilitates OoD generalization (Section 4). (d) We discuss the computational complexity of our proposed method, which can be scaled in either of the components (search and/or retrieval) independently, and is an easy drop-in replacement for multi-head attention in standard Transformer- like architectures (Section 5).
2 LIMITATIONS OF MULTI-HEAD ATTENTION
In this section, we first introduce the standard notation for multi-head attention (Vaswani et al., 2017) terms of search and retrieval mechanisms. We then highlight how the rigidity of the search-retrieval leads to limitations and redundancies in the parametrization of neural networks.
2.1 MULTI-HEAD ATTENTION BASICS
Key- Value Attention: Given a set of queries and key-value pairs, key-value attention computes a scaled cosine similarity metric between each query and the set of keys. This similarity score determines the contribution of each value in the output for the corresponding query.
More formally, given a set of input elements arranged in a matrix X E Nxd, we first obtain queries Q, keys K and values V using linear transformations X with learnable projection matrices
2

---

Published as a conference paper at ICLR 2022
X
W
Trput
Softmaa
OK
X
Logend
W.K
vi
W 6
Vaine Sultmox
Q.K.
va
Matrix Product
Weighted Bum
Irper
Transpoood Product
Cretgut
Cldss
SoCimax
V.
Figure 2: Computation graph for Compositional Attention. We show computations for one search and retrievals. Multiple searches operate in parallel with different search but shared retrieval parameters. The outputs are then fed to linear network to give the final output as in Equation 14
Wg E Rdxds, Wx Rdxds and Wo E pdxd, respectively. This is given by
= W = W4 = W..
(1)
For each query, similarity score is computed with each key using scaled cosine similarity (called scaled dot-product) to give the attention weights which are used to soft-combine the values as
QKT VTg
Attention (Q,nttonnQ, K,) = Softmax
axis *keys')
where Vdg
_ the scaling factor.
A multi-head attention mechanism combines multiple (say, h) independent key-value attention mechanisms in parallel provide the model the ability jointly attend to different positions and hence increase representational capacity. The outputs resulting from these multiple heads are concatenated together and then linearly projected back to the input dimension using a learnable matrix W E R hd, xd,
Multi-Head Attention:
Multi-Head Concat head) headz ... heads) W
where head,
Attention Qi, Ki, Vi).
2.2 SEARCH AND RETRIEVAL COMPONENTS
Here, we take the multi-head attention defined in Section 2.1 and decompose it into its two funda- mental components Search and Retrieval.
Search: A search is parameterized by the query and key matrices, that is, W,, and W. respectively. These parameters define a notion of compatibility metric between pairs of element xy and sa E X:
QKT
Search (Q, () Softmax
Carrch(Q.A))
axis 'keys'
(4)
Vdz
where Q X W and K = X W.. The above computation gives the compatibility between an element x. with other elements æx under the compatibility metric defined by the Search parameters.
Retrieval: A retrieval is parameterized by value matrix W, describing the kind of features in the input elements X that are relevant and need to be accessed for the downstream task:
Retrieval (Search Q, K), v) = Search (2, (,
(5)
where = X W., Note that each Retrieval defines the kind of attributes to access from input x/Ss and can take any Search result its input.
Multi-head Attention as a rigid pairing of Searches and Retrievals: Given the above definitions, one can see how standard multi-hen head attention amounts to a rigid pairing of Searches and Retrievals, such that an end-to-end function of fixed attribute pairs are learned at optimization. Indeed, h heads are composed of h different searche-retrieval pairs the 4th retrieval performed only on the 7th search. Multi-head attention thus amounts to special case of Equation 4 and 5
head, Retrieval Search (eaarc(Q,,45 QQ,,
K.),
(Note: Same subscript
(6)

---

Published as a conference paper at ICLR 2022
8-2
12
Twe Seares herren!
Trac Searhh fierrve
Taxlc Eerth fettoww
10
Dallend Dinces
355
Mylihee
Du DNCANON
1--00 0888001--3500 (35-
Contextual Retrieval Task Illustration. Dynamic search and retrieval based on search, retrieval and retrieval context features. Each element has corresponding output but we show it only for for brevity. Right: Performance on Contextual Retrieval Task. Here, we compare our proposed model against standard Multi-Head attention model (lower better) on various setups of the task. Our proposed model outperforms the baseline in both in-distribution as well as dittt---distibttiinn settings.
Figure 3: Left:
Viewing multi-head attention through these fixed search-retrieval pairings foreshadows a possible generalization of searches and retrievals which propose below. Before doing so, however, we first highlight specific shortcomings of standard multi-head attention.
2.3 SHORTCOMINGS OF RIGID ASSOCIATIONS
As described in Section 2.2, multi-head attention considers fixed pairing between searches and retrievals. While it has been widely successful across a variety of domains, we hypothesize that this rigid mapping is not always ideal and can sometimes lead to reduced capacity and learning of redundant parameters, missing an opportunity for better systematic generalization. We note that the search associated with each head defines feature (defined by query-key matrices W,, and W₂) based on which compatibility between objects is evaluated. Further, each head's retrieval allows the model access particular feature (defined by value matrix W.) from the searched objects. Following this, we showcase two types of redundancies that can arise in multi-head attention: (a) Search Redundancy which leads to learning of redundant query- matrices and (b) Retrieval Redundancy which leads to learning of redundant value matrices.
We highlight these two redundancies jointly using simple example illustrated in Figure where three objects with attributes: shape, color and location, are the subject of different questions. In (a) the model has to learn to search according to color and correspondingly retrieve shape information while in (b) it has to search according shape and retrieve location. On this task, standard multi-head attention (middle row) should learn two heads, one each for (a) and (b). To answer the question in (c), the model has to search according to color and retrieve location. While searching according to color exists in (a) learned by Head and retrieving location exists in (b) learned by Head 2, there is no way combine them. Thus, another head is needed to obtain the search of Head I and retrieval of Head 2. This leads to parameter redundancy and missed opportunity to factorize knowledge more efficiently, since these pieces of learned knowledge individually exist in Head / and Head 2 already.
The scenario in Figure 1 may look highly idealized because multi-head attention might not limit searches/retrievals on single features and is capable of doing more nuanced soft-combinations. While this may be the case for this simple example, what it highlights is the danger of rigid learned associations that limits re-composition of learned pieces of knowledge, leads to redundant parameters and potentially limits OoD generalization, irrespective of what the model learns. We discuss this in more detail in Appendix 1. and empirically explore these principles in a purpose built diagnosis task we call Contextual Retrieval Task, in Section 4.1. Finally, we reiterate that multi-head attention with d heads can only represent up to h unique (Search- Retrieval) pairings. In what follows, we propose to alleviate this fundamental limitation by allowing for S x R such pairings, with S the number search types and R the number of retrieval types.
3 COMPOSITIONAL ATTENTION DISENTANGLING SEARCH AND RETRIEVAL
We propose novel attention mechanism that relaxes static search-retriewe pairing in favour of more flexible and dynamic mapping. To do this, we let go of the concept of "head" altogether and instead propose independent and recombinable Searches and Retrievals, as defined in Section 2.2. The central innovation lies in the way these two components are combined: with query-key attention on retrievals.
4

---

Published as a conference paper at ICLR 2022
All Data
OoD Train Dala
OoD Test Data
Search 1
0.26
0.28
0.24
0.22
0.2
0.21
0.31
0.46
0.49
0.023
0.022
003
Swarch
0.26
0.27
0.23
0.23
0.2
01
Retrieval
Retrieral
Retrieval
Retrieral
Retrieval
Retreval
Retrieval
Retrivvll
Retreval
Retrival
Retiteval
Retrieaal
0.0
Figure Efficient Composition in Contextual Retrieval Task. We plot the average search-retrieval
activation statistics across data with Left: All possible value combinations, Middle: Subset of value combinations used for training. and Right: Remaining subset value combinations used for OoD testing. The activation statistics switch distinctly between OoD training and testing and stay around the average when all possible subsets are shown, thus highlighting good specialization.
Similar to heads, we start by defining S parallel search mechanisms. That is, we have S different query-key parameterizations We, and Wi, respectively. The output of each search is as defined in Equation 4. In essence, for each search i, we get
Q: = X W and
Ki = WK,
QiK?" vd
Search, Softmax
axis = "keys'
(8)
Next, we define different retrieval mechanisms, which correspond to the different W., matrices. These matrices used to obtain different attributes from the input. Formally, this is summarized as
Vj -==WW Way
where Vj highlights accessing different attributes. Then, corresponding to each search, all possible retrievals are done. This is similar to Equation 5 and is defined as
Retrievaly, Search, Vj
(10)
for all i,j. This step gives all the hypothetical retrievals for each search and out of this, one retrieval per search needs to be instantiated. This instantiation done through a secondary attention mechanism computed using retrieval queries Q2 and retrieval keys Kúj, that are obtained as
Q, X W
Ki Retrievaly W
(11)
where the parameter W. 9. E Rdddd is a different matrix for each search indexed by and, together with W &. is tasked with driving the pairing between search and retrieval. We broadcast the matrix Q from Ie N xdr to IR Vx xxxxd, and define K; E R xixd; by
K, = Concat K K.R)
(12)
Thus, through these retrieval queries and keys, the required instantiation per search is done as
Q,Ki Vd,
T
CAtt; Softmax
axis 'retrieval' Retrievalyy
(13)
Value Scores
where the transpose is over the last two axes. Hence, for each search i, the softmax gives attention weights over all possible retrievals and the winning retrieval is instantiated through this soft attention mechanism. Finally, similar to multi-head attention, the outputs these parallel searches are combined by concatenating them and passing them through linear network:
Compositional Attention = Concat CAtt, CAll,,..,, W
(14)
where 0 E R$d xd Note that in this mechanism, the choice of retrieval for each search is not fixed, and instead dynamically modulated by Q, and K; respectively. We refer the readers to Figure 2 for visual depiction of the computation graph.
Compositional Attention allows the model to have (a) Different number of searches and retrievals, i.e. and It respectively, (b) Dynamic selection of shared retrievals for each search and (c) Representation capacity of S x R (Search- - Retrieval) pairings. Thus, we highlight that Compositional Attention disentangles search and retrieval and solves the redundancies of multi-heat attention (Section 2.3).
5

---

Published as a conference paper at ICLR 2022
Algoritio
Searches
Retrievals
Unary Accuracy
Binary Accuracy
Ternary Accurucy
4
98.6
84.4
64.9
Transformer
8
98.5
84.5
65.4
98.7 .0.2
86.8
66.4
2
98.8 00.1
88.2
66.9
Compositional Transformer
3
98.9.102
89.8
67.1
4
98.6.003
84.9
64.8
Table 1: Performance Sort of CLEVR. We highlight that our proposed model outperforms the baseline across the different question types even with lower number of searches and/or retrievals.
EXPERIMENTS
For all experiments, we consider the standard Transformer model (Vaswani et al., 2017) with parameter sharing across layers (Dehghani et al., 2018) as our baseline. For visual tasks, we consider Vision Transformer introduced by Dosovitskiy al. (2020) as our baseline. Our proposed model, Compositional Transformer, uses the Compositional Attention mechanism as a drop-in replacement for the multi- heal--head attention block while keeping the rest the architecture same. We also perform ablations on number of retrievals, model sizes and complexities as discussed in Appendix B.7
Through our experiments, we show that (a) flexible composition search and retrieval leads better in distribution and out of distribution performance, and (b) our proposed model can achieve similar or better performance, often with fewer retrievals.
4.1 CONTEXTUAL RETRIEVAL TASK
As start, consider a purpose built task to better understand how learned attribute mappings may be recombined by attention mechanisms. Our goal in introducing this experiment is to explicitly control search and retrieval attributes that allows for (1) systematic testing for OoD generalization, and (2) direct comparisons between sepected/ground-truue we propose a supervised set-based regression task where objects with several features need to be selectively associated with one another (e.g. find the closest blue object), and distinct feature values must be returned based on inherent contexts (e.g. report shape if position is middle and size otherwise). Our task seeks minimal setting where such associations can be present, and involves a collection of N objects {x, - each with scalar-valued features, as illustrated in Figure 3 (left). Corresponding to each object x,, the output scalar
and learned activation patterns. Thus,
y: Σ a
(15)
where a, are randomly selected fixed weights, and uf are the results of S comparative searches pairing object x; with other objects in the set. To describe these comparative searches, we first outline the composition of each object. Here,
X, = [(1 =()))...,,')) (g)
(16)
where z' S and z R contain S "search" features and R "retrieval" features respectively, and one-hot vectors r E {0, 1 contain the retrieval preference of object x; for search s.
Search Step: Each object x; searches for objects in parallel that are closest it based the search feature. This operation is given by
j = arg min =
(17)
jfi
where s denotes the th search feature, and hence there are parallel independent searches with j denoting the winner of search for object x;.
Retrieval Step: Corresponding to each winner jt, the retrieval context r' dictates which retrieval feature access from the jž object. This is given by
(18)

---

Published as a conference paper at ICLR 2022
Algorithm
Searches
Retrievals
CIFARIO
FashionMNIST
SVHN
Equilateral Triangle
Transformer
4
4
77.2-03
89.4g00
85.0go.
96.8
77.5_05
89.9
85.0
97.010.
2
77.9
89.9
86.0
97.2 10.3
Compositional Transformer
4
3
78.0
89.9..0
85.2
97.1 1:04
77.6-02
90.0
86.0
96.9
Table Performance on Multi-Task Image Classification, We highlight that our proposed model outperforms the baseline across different number of retrievals.
Training is done by minimizing the L1 loss on { targets (see Appendix C.I for details). While seemingly synthetic, this task considers multiple objects with various attributes and tests the model's ability to flexibly compose searches and retrievals. This is tightly associated with real world scenarios but in controlled low dimensional setting. We provide visualization of the task in Figure 3 (left).
OoD Setup: Since all searches share the different values to be retrieved, we construct an out-of- distribution version of the dataset by showing only certain ground truth search-retrieval combinations in training and others during testing. More specifically, for task with S searches and R retrievals there are R possible unique value combinations which define the retrieval preference. For the Ool setup, we use a fraction of them in the training set and the rest in the test set. We stress that only models that factorize search and retrieval efficiently and compose them flexibly can do well on the unseen search- combinations.
Quantitative Results:
We experiment on a range of searches S and retrievals R. Since this task one-layered attention models for experiments.
only requires set of parallel searches, We observe that compositional attention consistently outperforms multi-head attention across a wide range of task hyperparameters in both in-distribution well OoD settings (Figure 3 right). Further, our proposed method outperforms the baseline across various model complexities. Analysis: We visualize the value scores from Equation 13, aggregated over all searches and entities and contrast them with ground-trute tttnhhhhtthhhho retrievals that the task provides. Fig- ure 5 shows that the model specializes on what features to retrieve. We also analyse search-retrieval pairings by taking model trained on subset of all possible value combinations and plotting its activation statistics across (1) all possible value combinations, (2) value combinations seen during training, and (3) value combinations seen dur- ing OoD testing in Figure 4. As expected, we that activation statistics average out in (1), as each search fea- tures called in a uniform i.i.d. manner. In contrast, we see increasing specialization of activation pattern for (2) and (3), respectively, consistent with selective recombi- nation made possible by compositional attention. The key finding from this analysis that compositional attention is dynamically specializing over value retrievals, allowing for better generalization on unseen value combinations in the OoD setup.
Taxt: Search lrvel
Value
0.07
0.93
0.00
0.00
Value
0.61
0.38
0.00
0.00
Value 3
0.00
0.00
0.00
1.00
Value
0.00
0.00
1.00
0.00
Ressivall
Reteeall
Ratnnaal
Retsitval
Figure 5: Retrieval Specialization Contextual Retrieval Task. The pro- posed model learns to specialize its retrieval (X-axis) based on ground truth values (Y-axis).
Appendix C.I contains further details about the task, hyperparameters and further ablations on model complexities and specialization visualizations across different task settings.
4.2 RELATIONAL REASONING IN SORT-OF-CLEVE
Sort-of-CLEVR (Santoro et al., 2017) is Visual Question-Answerine (VQA) task that tests the model's understanding of the scene by posing questions based on the properties of the objects and their relations with each other. The task consists of three types of questions: (a) Unary, which are based on properties of single objects, (b) Binary, which are based on the relationship between two objects, and (c) Ternary, which are based relationship between three objects.
7

---

Published as a conference paper at ICLR 2022
220
Tacc Same
Tacc SNTS
Tase Destoctoones of3
Task: Mardiy Rules
Vocal Compotitoral
102
102
1185
Figure 6: Performance ESBN Tasks. Our proposed model outperforms the baseline across different tasks especially in the extreme OoD setting.
Quantitative Results:
We use 4-layered transformer models with parameter sharing for our ex- periments. Table highlights that Compositional Transformer outperforms standard transformers across different question types (unary, binary and ternary), even with fewer searches and/or retrievals than the number of heads in the standard transformer. This is attributed to the fact that our model exploiting the inherent compositional nature of this task by reusing the parameters for dynamic retrieval of context dependent attributes.
We refer the readers to Appendix C.2 for details regarding this task, model hyperparameters and performance ablations variety of model settings.
4.3 EQUILATERAL TRIANGLE DETECTION
This a binary classification task introduced by Ahmad & Omohundro (2009), where the aim is to decide if the set of three point clusters in the image forms an equilateral triangle. For triangle to be equilateral, the midpoints of these clusters should be equidistant from each other.
Algurithm
Searches
Retrievals
Test Accuracy
Transformer
4
4
93.8.001
95.6
Compositional Transfommer
4
96.7-04
97.0 mg
Quantitative Results. Table 3 highlights that compositional attention outperforms multi-head attention across different number of retrievals.
Table 3: Performance on Equilateral Triangle Detection. Our proposed method outperforms the baseline over different number of retrievals.
We refer the readers to Appendix C.3 for more details.
4.4 MULTI-TASK IMAGR CLASSIFICATION
We pose the problem of image classification across four different datasets - -CIFAR1O, FashionMNIST. SVHN and Equilateral Triangle Detection as a multi-task learning setup. To perform well on this, models need to figure out which information is shared across datasets and which is private a dataset.
We train a 2-layered universal transformer with four learnable embeddings, one for each task. Our proposed model replaces the multi-head attention with compositional attention the baseline. Table 2 shows that the proposed model outperforms the baseline.
Quantitative Results:
Appendix C.4 contains further details about the task and the models.
4.5 LOGICAL REASONING IN ESBN TASKS
Webb et al. (2020) introduces a suite four tasks to test the model's ability to perform logical reasoning in an OoD setting. Within each task, the input is sequence of objects and the output based on the relationships between these objects. For example, in the "Same/Differeet" task the different. Each task has
model receives two objects and has to predict whether the objects OoD criteria - eg. m95 refers to training with 5 out of 100 objects and testing with the remaining 95.
done in the original benchmark (Webb et al., 2020). Figure highlights that Compositional Attention outperforms Multi- Head Attention across all the tasks, especially higher hold-out (m) values. This shows that the proposed model is able to better handle distribution shifts than the baseline model across all tasks.
Quantitative Results: We
single layered transformer model
We refer the reader to Appendix C.5 for details regarding the tasks well the models.
8

---

Published as a conference paper at ICLR 2022
Cuioll Length
Algovithue
Seorchex
Retricnals
22
25
26
27
28
29
30
0A11 1oo
g:06
0.12
0.11gcoe
j.22
0.02,oox
g.03
0.05
0.01 zoo
a208
6.24
0.34
0.38
0.10
1.08
Compositional Tranffrmer
Looe
g00
0.24
a.27
L000
0.09 as
0.16
g000
2.06
2.30
g.26
0.04
2.08
p.002
g.0
0,14gnos
0.18
0.18
g.30
0.04.mom
0.06
0.08
Table 4: Performance on SCAN. We highlight that our proposed model outperforms the baseline across the different question types even with lower number of searches and/or retrievals.
4.6 SCAN DATASET
SCAN (Lake & Baroni, 20) a synthetic sequence to sequence task aimed at translating natural language instructions into sequence of actions for robot to follow. We follow the length extrapola- tion generalization split (Newman et al., 2020; Csordás et al., 2021) where training is done on lengths of 22 output actions and tested on larger output sequence lengths.
We train a 3-layered universal transformer as the baseline and compare with our proposed plug-in attention replacement. Table showcases that compositional attention outperforms the standard multi-head attention across multiple cutoff lengths.
Quantitative Results:
We refer the readers Appendix C.6 for further details regarding the task and models.
4.7 LANGUAGE MODELLING
Midel
Trenafermr
We perform experiments on the WikiText-100 data cor- pus (Merity et al., 2016) for the language modeling task, Here, the task is predict probabilities for next masked words, evaluated through perplexity.
Compocitionl
We use 6-layered transformer mod- els with parameter. We plot the validation perplexity against epochs in Figure 7 which highlights that our pro- posed attention mechanism not only outperforms the base- line but also converges faster. Further, see that our pro- posed model obtains test perplexity 38.8 as opposed baseline's perplexity 39.640..
Quantitative Results:
Epoch
Figure 7: Performance on Language Modeling (WikiText103). We illustrate that our proposed mechanism outper- forms the standard multi-head attention.
We refer the reader to Appendix C.7 for further details.
DISCUSSION AND CONCLUSION
Summary. In this work, revisit Multi-Hese Attention, a popular attention mechanism, and highlight its shortcomings due to the rigid association between search and retrieval mechanisms. We argue that this rigid coupling hinders re-usability of parameters and reduces the expressivity of the model. To mitigate this, we propose a novel mechanism which uses value retrieval mechanism flexibly compose searches and retrievals. Experiments on various tasks show that our proposed method outperforms standard multi-head transformers, while often using only a fraction of retrievals.
Complexity. While our proposed mechanism requires additional parameters for the computation of value scores, we highlight that this increase is often minuscule compared to the total number of parameters. Crucially, we note that this light increase in parameters per search mechanism is easily offset by reducing the number of retrievals needed. For all our experiments, our proposed models offer similar capacity as the baselines unless stated otherwise. This highlights that the improved performance is due to flexible composition of scarch and retrieval and not number of parameters. We discuss computational complexity in detail Appendix B.5.
Limitations and Conclusion. Motivated by the need for efficient factorization of knowledge and dynamic reusability of learned pieces of computations, we propose Compositional Attention, a first step towards flexible composition of search and retrieval. We also highlight some of the limitations of the proposed methodology in Appendix B.8 and hope that our work would promote research into more flexible models capable of better systematic generalization.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENTS
The authors would like to thank Felipe Codevilla, Nicolas Gontier, Disha Shrivastava, Damjan Kalajdzievski, Olexa Bilaniuk, Ioannis Mitliagkas and Gauthier Gidel for their time and effort on helping improve this work. SM would like to acknowledge the support of UNIQUE and IVADO towards his research. IR, YB and GL acknowedgge the support from Canada CIFAR AI Chair Program, as well as Samsung Electronics Co., Ldt. IR acknowledges support from the Canada Excellence Research Chairs Program. GL. acknowledges NSERC Discovery Grant TRGPIN-2018-04822
ETHICS STATEMENT
The authors do not foresee negative socictal or ethical impacts beyond those expected from general improvements in MI. Furthermore, the authors note that with added flexibility comes the increased possibility of learning negative implicit biases in real-world applications. Much like other cutting edge ML methods, careful considerations of potential negative impacts society should be considered before deploying models for applied uses.
REPRODUCIBILITY STATEMENT
The authors provide code used to run all the required experiments and will open source their code with ample documentation allow for ease of reproducibility. Further, the authors provide the exact formulation of the proposed model in Section 3 and provide the hyperparameters as well as different existing codebases used for all the experiments in Appendix C.
10

---

Published as a conference paper at ICLR 2022
REFERENCES
S. Ahmad and S. Omohundro. Equilateral triangles: challenge for connectionist vision. 2009.
Bernard J Baars. In the theatre of consciousness, global workspace theory, rigorous scientific theory of consciousness. Journal of Consciousness Studies, 4(4):292-309, 1997.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. CoRR, abs//449.0073. 2015.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Deep equilibrium models. arXiv preprint arXiv: 2019.
Yoshua Bengio. The consciousness prior. arXiv preprint arXiv: 2017.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv preprint arXiv: 10912, 2019.
Xinyun Chen, Chen Liang, Adams Wei Yu, Dawn Song, and Denny Zhou. Compositional generaliza- tion via neural-symbolic stack machines. arXiv preprint arXiv:2008. 06662, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 2020.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of transformers. arXiv preprint arXiv:2108. 12284, 2021.
S. Dehaene, H. Lau, and S. Kouider. What is consciousness, and could machines have it? Science, 358(6362):486 2017.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. CoRR, Abs/1807.033$ 0077703819 19, 2018. URL http: /ttp/vaxvvvv..///////eee/ea.aa.a wttp:/ii.wr/////////eeee 03819.
Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: aeti:811000000 2018.
David Ding, Felix Hill, Adam Santoro, and Matt Botvinick. Object-based attention for spatio- temporal reasoning: Outperforming neuro-symbolic models with flexible distributed architectures. arXiv preprint axXiv:2012.08500 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint 1:Xi::2010.11929, 2020.
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar, Addressing some limitations of transformers with feedback memory. arXiv preprint arXiv:2002.09402. 2020.
Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. arXiv preprint arXi::2011.15091 2020.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909. 10893, 2019.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006. /6225, 2020.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nico- las Heess, Michael Mozer, and Yoshua Bengio. Neural production systems. arXiv preprint 103.0:2103.01197, 2021a
11

---

Published as a conference paper at ICLR 2022
Anirudh Goyal, Aniket Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Mozer, and Yoshua Bengio. Coordination among neural modules through shared global workspace. arXiv preprint arXiv:2103.01197, 2021b.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv preprint 1410,544 2014.
Drew A Hudson and Christopher D Manning. Compositional attention networks for machine reasoning, arXiv preprint 18Xiv:880300307 2018.
Dicuwke Hupkes. Verna Dankers, Mathijs Mul. and Elia Bruni. Compositionality decomposed: how do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
Nan Rosemary Ke, Aniket Rajiv Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Jimenez Rezende, Michael Curtis Mozer, Yoshua Bengio, and Christopher Pal. Systematic evaluation of causal discovery in visual model based reinforcement learning, 2021.
URL httthhttts///pnneee /wwenreeevvveeeeeeeeeeeeeeeeeeee iew fww.fofornurmm dd=gp 9C--
Giancarlo Kerg, Bhargav Kanuparthi, Anirudh Goyal ALIAS PARTH GOYAL, Kyle Goyette, Yoshua Bengio, and Guillaume Lajoie. Untangling tradeoffs between recurrence and self-attention in artificial neural networks. Advances in Neural Information Processing Systems, 33, 2020.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, et al. Measuring composi- tional generalization: comprehensive method on realistic data. arXiv preprint arXiv: 11Xiv:191..09773, 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint urXiv:2001.04455, 2020.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sesguence-t--seqeeeee recurrent networks. In International conference on machine learning, pp- 2873-2882. PMLR, 2018a.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of requence-to-sequeee recurrent networks. Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2873-2882. PMLR, 10-15 Jul 2018b. URL http: http///pccceee..../////////.. mlr. mlrs.pess/ v80 lakelakela... kaakeelBaa..... html.
Alex Lamb, Di He, Anirudh Goyal, Guolin Ke, Chien-Feng Liao, Mirco Ravanelli, and Yoshua Bengio. Transformers with competitive ensembles of independent mechanisms. arXiv preprint arXiv:2/03.0m 1030:2203.000336, 2021.
Yuanpeng Li, Liang Zhao, Jianyu Wang, and Joel Hestness. Compositional generalization for primitive substitutions. arXiv preprint arXiv: 190000212,, 2019.
Adam Liska, Germán Kruszewski, and Marco Baroni. Memorize or generalize? searching for compositional RNN in haystack. CoRR, abs/1802.06467, 2018. URL htLLtt:///axii.oor abs/ 66467.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. arXiv preprint arXiv:2006. / 15055, 2020.
Minh-Thang Luong, Hieu Pham, and Christopher D Manning. Effective approaches to attention-based neural machine translation. arXiv preprint arXiv:1508. 040255, 2015.
Kanika Madan, Rosemary Nan Ke, Anirudh Goyal, Bernhard Bernhard Schölkopf. and Yoshua Ben- gio. Fast and slow learning of recurrent independent mechanisms. arXiv preprint arXiv:2105.08711 2021.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv:/609.00 2016.
12

---

Published as a conference paper at ICLR 2022
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint arXiv: 10650 2019.
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning combine top-down and bottom-up signals in recurrent neural networks with attention over modules. In International Conference on Machine Learning, pp. 6972-6986. PMLR, 2020.
Benjamin Newman, John Hewitt, Percy Liang, and Christopher D Manning. The eos decision and length extrapolation. arXiv preprint arXiv:2010.07174, 2020.
Rodrigo Nogueira, Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2102. 13019, 2021.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
Adam Santoro, David Raposo, David GT Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. arXiv preprint wrXiv:1706.0122,, 2017.
Muru Selvakumar, Suriyadeepan Ramamoorthy, Vaidheeswaran Archana, and Malaikannan Sankara- subbu. Compositional attention networks for interpretability in natural language question answering. arXiv preprint arXiv: /8/0. /2698, 2018.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXiv preprint arXiv:2009.007:22 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Autention is you need. CoRR, abs/ (bss//706.0362,, 2017. URL
attp://arxis: wwwoor//aao.o http//bss/111/...//.// 037662.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.09418, 2019.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
Taylor W Webb, Ishan Sinha, and Jonathan D Cohen. Emergent symbols through binding in external memory, arXiv preprint arXiv:2012.14600) 2020.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International conference on machine learning, pp. 2057.2007. PMLR, 2015.
13

---

Published as a conference paper at ICLR 2022
APPENDIX
A RELATED WORK
The advent of transformer-like models have led to advancements on various flavours of attention based models. This revolution first started with augmenting Recurrent Neural Networks (RNNs) with a form of semi- parametric memory structure through attention (Bahdanau et al., 2015) and it soon led to people questioning the need for recurrence. This line of questioning resulted in a famous class of models that get rid of recurrence in favour of just parallel self-attention computations that are quite efficient to do modern hardware (Vaswani et al., 2017). We briefly discuss the various advances along these lines and distinguish how our proposed attention algorithm is different from them.
A.1 ATTENTION
Attention has been major component of human cognition which allows humans to selectively process relevant information from the plethora of sensory stimulus we receive. The idea of selecting relevant features from a sea of information allows us to make predictions in both robust as well as compute efficient way. Inspired from neural cognition, there have been a lot of efforts in trying to introduce a notion of attention to relevant states of the input for reliable downstream prediction (Xu et al., 2015; Luong et al., 2015; Kerg et al., 2020).
major problem in Recurrent Neural Networks based systems is the problem of vanishing and exploding gradients that happens due to improper credit assignment in the model. This is because RNNs model all the information seen up to a certain time through a parametric fixed sized vector which undergoes repeated computations over all time steps. This makes the system brittle to changes in sequence lengths or in presence of long sequence of distracting information. A way to solve this problem was move away from parametric representations of the entire past and instead rely on dynamic semi eemi-parametric "memory" to allow these models to look back whenever needed (Graves al., 2014; Bahdanau et al., 2015). These works aimed at augmenting recurrence with self-attention and demonstrated that when combined with these cognition-inspe were able to extrapolate much better to larger sequence lengths.
inductive biases, ML systems
Following this, there has been lot of recent work that then aimed to remove recurrence between time- steps and rely solely on querying information through self-attention. Recent advances on multiple domains ( (Vaswani et al., 2017; Dosovitskiy et al., 2020; Ding et al., 2020; Locatello et al., 2020) showcased that removing recurrence from the picture and relying solely on parallel computations not only leads to significant improvements in performance and generalization but is also easier and faster to train on current hardware. Since the advent of these transformer based models built fundamentally on multi-head attention, the role of attention has become increasingly important across various domains like vision, language and reinforcement learning. It has also led to a lot of research on various architectural choices in fully attention-based systems, some of which we discuss in Appendix A.2.
It is, however, important note that there has been some research that highlight the need for recurrence jointly with self-attention for solving certain logical reasoning tasks efficiently (Hudson & Manning, 2018: 8; Selvakumar et al., 2018; Webb al., 2020).
A.2 TRANSFORMER VARIANTS
The ubiquity of self oeff-attention models in the current ML community has led to tremendous research aimed incorporating different inductive biases in the attention mechanism used; namely in the multi-head attention. Most of these variants aim to alter multi-head attention in way that would remove the quadratic time complexity computational bottleneck that present in standard meati-head attention. However, there are certain works that aim more on the fundamental inductive biases that the attention encodes as opposed to computational benefits. We discuss some of these variants here.
Reducing Computational Complexity. Given set of n vectors, the standard multi-head attention aims to create an n X n attention matrix that takes quadratic complexity to compute. This bottleneck prevents usage of self-attention when n is large. In light of this, lot of recent research aims to reduce this quadratic complexity to n log n or linear complexity. This is often achieved by either introducing some restrictions in the n attention matrix through locality sensitive hashing (Kitaev et al., 2020),
14

---

Published as a conference paper at ICLR 2022
sparsity (Child al., 2019), low rank approximation (Wang et al., 2020) or through random features for approximation of softmax (Choromanski et al., 2020). We refer the readers to Tay et al. (2020) for more detailed analysis of different transformer variants.
Additional Inductive Biases.
While a lot of the above transformer variations are designed to prevent the quadratic bottleneck, most of them also add certain additional inductive biases in the model. For example, the addition of sparsity not only reduces the computational complexity but also adds the additional inductive bias of sparse information routing between different elements. There certain additional variants (Lamb et al., 2021; Goyal et al., 2021b) that add other inductive biases, eg. factorized state space and global workspace bottleneck respectively in the transformer model.
A.3 MODULARITY, COMPOSITIONALITY, REUSABILITY AND BOTTLENECK
There have been recent efforts along the lines of modularized computations in an effort to improve the model's capacity to perform systematic generalization. In particular, humans are able to com- partmentalize information and act on it in disentangled, context-driven and robust fashion. These cognitive fundamentals have led to a preliminary movement of Machine Learning systems into this space. We discuss some of the essential ingredients below.
Modularity. Modularity refers to factorization of knowledge into smaller components that can inde- pendently exist and act sensory information. It can be considered as disentangled representations that allow for interventions on these different components or factorized mechanisms where each mechanism has specific purpose and can act on part whole of the sensory information. The fundamental aim of modularity is to prevent unrestricted information flow across a whole monolitihic system and instead to learn an often end-to-end fashion factorized representations and mechanisms that act on these representations. Recent works (Goyal et al., 2019; 2020; 202 la; Mittal et al., 2020; Madan et al., 2021; Lamb et al., 2021; Ke et al., 2021) along the lines of factorizing knowledge demonstrate that it often leads to increased robustness and better OoD performance.
Compositionality and Reusability. Humans are able to perform complex tasks even novel and unknown situations. This capacity often stems from the fact that our complex actions are in reality compositions of simpler primitives and our understanding of these primitives so good that we are able to dynamically combine these primitives into novel complex actions. Recent research has started looking into tasks and systems that test and allow for compositional generalization (Lake & Baroni, 201 8a; Li et al., 2019; Keysers et al., 2019; Chen et al., 2020; Hupkes et al., 2020; Goyal & Bengio, 2020), which generalization to novel combinations of the underlying primitives/mmccaaisms.. The primary reason why number of modular systems are constructed in recurrent domains is because we want the factorized mechanisms to be reusable in a number of scenarios. Reusability of knowledge (Dehghani et al., 2018; Bai et al., 2019) allows for learning of disentangled mechanisms in modular system which then has the potential to lead to efficient compositions of the learned disentangled mechanisms. Recent success of systems that use computations that can be reused multiple times demonstrates that reusability is actually an important fundamental for obtaining compositionality.
Bottleneck. Conscious attention in humans is a key ingredient to create a bottleneck of information processing, according to the Global Workspace Theory (Baars, 1997; Dehaene al., 2017). The key use of this bottleneck is to restrict information flow across the whole network, human brain or otherwise, which allows for robustness to insignificant pieces of sensory information. The usefulness of this bottleneck has been hypothesized to be linked to the sparsity and simplicity of the dependencies manipulated with System 2 cognition (Bengio, 2017: Goyal & Bengio, 2020). Recent works along these lines (Goyal al., 2021b) illustrate that modular systems with the addition of a bottleneck efficiently factorize computations and then compose them in a dynamic and context dependent fashion often lead improved performance, faster adaptation and systematic generalization (Bengio et al., 2019; Ke et al., 2021).
B PROPOSED MODEL.
In this section, we provide additional details about the general motivation, architecture setup and argument for using parameter sharing across layers. We further provide details about computational complexity of the proposed model and some ablations that consider.
15

---

Published as a conference paper at ICLR 2022
1.0
Task: Search Ratreval
Tal: 4 Search Ratioral
Tex:: Search de futroora
0.8
0.6
Compastional
0.2
0.0
hakds
heads
Haker
Hoader
hhaaes
reoinnal
Aloads
32 haAe
Arerrta heads
Ge haads
rettrral
4 hakde
Gf hoads
sorebieral
Capocity
High Capecty
Capacity
Capaciy
Caladiy
Figure 8: Performance on Contextual Retrieval Task. We compare our proposed model against standard Multi-Head attention (lower loss is better) on various setups of the task. Our proposed model outperforms the baseline across various model capacities (low and high) and number of heads.
B.1 MOTIVATION
While the setting in Figure 1 may look idealistic in the sense that it is very likely that transformer heads do not learn this interpretable single feature functions for search and retrieval, we argue that this rigidity still exists between search and retrieval in a standard multi-head attention framework. To see this, we note that the search component Search,, is parameterized by the query and key matrices Wg, and W., respectively and the retrieval component Retrieval, is parameterized by the value matrices W.,, Both these components lead to computations that are dynamic based on the input but the pairing between search and retrieval fixed, that is, Retrieval, takes Search, as its argument (notice the same h subscript), also highlighted in Equation 5. Thus, whenever there is a need to share retrieval parameterizations across multiple searches, a standard multi-head attention would lead to learning of redundancies because there is no notion of sharing of retrievals between searches.
Contrasting this with the proposed approach, Compositional Attention, we see that now there is a notion of sharing of retrievals for different searches. That is, two different searches can still opt for the same retrieval parameterization, which alleviates the rigidity and redundancy that is explained above. Note that this discussion does not depend on the model's capacity to selectively pick features as is illustrated in Figure 1. This shows that irrespective of what these searches and retrievals learn, the discussed drawbacks of multi-head attention still exist if an optimal solution requires sharing of retrievals across searches. We highlight the motivation through the idcalistic example of multiple features solely for ease of explanation and appealing to the fundamental cognitively inspired inductive bias that we try to incorporate.
We emphasize that multi- attention and the proposed compositional attention are not two separate classes of methods. In fact, our proposed mechanism is a strict superset of multi-head attention and thus presents a more general framework that subsumes the family of multi-head attention. One can see this from Equation 13 where, given enough capacity to represent any h x h matrix, we recover multi-head attention by setting the number of searches and retrievals as h and having the "Value Scores" matrix as an h x h identity matrix (or any h x h permutation matrix in general), with h being the number of heads. Thus our mechanism not only solves the redundancies highlighted in this text but also provides more general class of attention mechanism.
B.2 DIFFERENCES FROM EXISTING WORK
We propose Compositional Attention, novel attention mechanism aimed at disentangled computa- tion of search and retrieval. Unlike in multi-head attention, this allows for a flexible and dynamic composition of searches and retrievals.
This is different from MAC and its variants (Hudson & Manning, 2018; Selvakumar et al., 2018) because the proposed algorithm is a completely parallel system without recurrence. Further, we see that in MAC, disentanglement is driven by privileged information; iee through the difference between what question and image is. This privileged information may not be present across a variety of tasks (eg. language modelling, classification, etc.). Our proposed model, however, does not require privileged information and is therefore easily applicable to lot of different domains. Further, MAC
16

---

Published as a conference paper at ICLR 2022
122
Taac: Search fazz teal (Ded)
Tesk: 4 Search fiez (2o2)
Tasl: Search 1e feez eeal (CeO)
1.0
0.8
0.6
Compastional
02
0.0
haadd
Hoace
Aloads
32 haaes
&nerrtvl heads
Ge haads
hakde
G4 hoads
reoinnl
sorebieral
High Capecty
Capacity
Figure Performance on OoD Contextual Retrieval Task. We showcase that our proposed mechanism outperforms standard Multi-Head attention (lower is better) on out of distribution (OoD)
variant the various setups across various model capacities (low and high) and number of heads.
does not have multiple parallel searches and retrievals and thus, our proposed model aims to solve considerably different problem.
While one may be tempted to think of head pruning (Michel et al., 2019; Voita al., 2019) way of removing redundancies standard multi-head attention, we stress that the core goal and motivation of our work is considerably different. Pruning of head essentially means rigid search-retrieval pairing from the learned system as its utility for solving the task is negligible. However, in this work, the redundancy we want solve is when a sub-part of head is redundant but not the whole head. That is, when either the search or retrieval part of the head is redundant, but not both. Figure I highlights when only a sub-part of the head is redundant and not the whole head, and how compositional attention resolves the problem.
eliminate
Further, Compositional Attention is different from the various transformer variants Appendix A.2 because it does not aim to solve the quadratic computational bottleneck but instead adds an inductive bias that has not been explored yet. We also note that the proposed model is amenable to the various computation tricks discovered for multi-head attention.
B.3 ARCHITECTURE DETAILS
The standard transformer model (Vaswani et al., 2017) has a number of layers, where each layer is composed of two components, the multi-head attention (Section 1) which is followed by a MLP (Multi-layer perceptron) with single hidden layer. There are residual connections at the end of the multi-head attention step as well the MLP.
In this work, we follow Dehghani et al. (2018) and consider the models that have weight sharing across layers. For ease of experiments, we do not consider adaptive stopping criteria. We consider this choice because we want reusable pieces of computations, and Universal Transformers is one step towards that goal.
Our view of transformer models is that different heads perform parallel information retrieval with not only different kinds of searches but also different kinds of retrievals. Information from these parallel retrievals is then jointly processed through a linear layer, followed by another MLP. There are residual connections after the linear layer and the MLP.
For our proposed Compositional variants, we basically replace Multi Head Attention in the models with Compositional Attention while keeping all the other details the same.
B.4 MULTIPLE LAYERS AND WEIGHT SHARING
number of works demonstrate that Transformers with weight sharing are competitive with the standard transformer models (Dehghani et al., 2018; Bai et al., 2019). We also believe that reusing computations provides more pressure on the system to learn meaningful and multi-purpose parameters (eg. it is easier to learn redundant head if it is used only once if it is repeatedly used). One might be tempted to think that increasing the number of layers or removing weight sharing might compensate for the flexibility provided by proposed system. However, argue otherwise.
17

---

Published as a conference paper at ICLR 2022
Seal Search Natrional
Salk Search is lova
lok: 4 Searc ib tatricall
Nuth Compodtional
beations
Figure 10: Convergence on Contextual Retrieval Task. We see that the proposed mechanism converges faster and works well even in low data regime (low iterations).
Lets assume we have Transformer model without parameter sharing which has layers and h heads. Then, the number of unique search-retrieval pairings that can be computed by the model is 1h (h if parameter sharing). Contrasting this with compositional attention, we see that the number of unique search-retrieval pairings are actually Isr (sr if parameter sharing) where is the number of searches and r the number of retrievals. So, we use similar number of layers, compositional attention still allows for more combinatorial possibilities to be learnt. Viewed another way, at scale, the proposed mechanism has the potential to reduce the number of layers needed for tasks calling for flexible search and retrievals.
Another important point that even we have more layers (with or without parameter sharing), multi-head attention still only learn rigid combination between search and retrieval. So, if the task requires dynamic choice from all possible pairings between search and retrieval, the model will have learn each pairing in separate head combinations, whether be in the same future layers. This is because adding more layers does not change the way searches and retrievals are combined, which is what we focus on here.
B.5 COMPUTATIONAL COMPLEXITY
Number of Parameters. We keep the parameter counts within -5% each other for the compared models and the same parameter count at 140M parameters for the language modelling experiment. We also stress that our proposed models with fewer retrievals are even more tightly matched and often lower in parameters than the baseline and still outperform them on a number of tasks.
Training Time, While Compositional Attention increases the complexity the model, we note that the training time of proposed models are generally within - 10% of the baseline and hence the added complexity does not impede the model much.
FLOPs. We estimate the FLOPs of the proposed model for Equilateral Triangle Detection task using an off the shelf library and see that they are 10% of each other and the baseline. In particular, we also see that for fewer retrievals, the FLOPs are either the same or lower than the baseline.
Parallel Computations. Transformers allow for efficient implementation using GPUs due to parallel computations for each word in the sentence (or each object in the scene). Further, they allow for parallel computation of each head for each word. Correspondingly, in our proposed model, we still do parallel computations for each word in the sentence, and compute the output of different searches in parallel. The only additional complexity is another soft-attention for choice of retrieval for each search. This is also done in parallel for each search and hence retain all the major efficiencies that Multi-Head attention enjoys on GPUs.
Amenable to Different Variations. We note that lot of the current advances in standard multi-head attention, eg. sparse attention matrix, can be incorporated in the proposed model too. We can also have sparsity the retrieval end where we restrict certain searches pick from smaller set of retrievals. We believe that these analysis are important future works but out of scope of this paper.
Complexity Combinatorial Advantages. While we sometimes have more complexity than multi- head attention, this small increase in complexity is often offset by the combinatorial advantage that
18

---

Published as a conference paper at ICLR 2022
9.068
0:93
1.50 08
0.00076
0.36
9:13
0:38
6:19
Valn
0202
Wilen
Nate
2.4c
0.03033
3o 05
0.12
0.13
0.13
0.63
Vilue
Valer
Zedd
8.5
2.7w05
Volor
0.36
0.13
0.12
92
Nevileal
Acrieall
Antmral
Ammieal
Ammeal
Rerlendl
Figure 11: Specialization plots for the Contextual Retrieval Task. We plot the attention scores for ground truth retrieval learned retrieval for different task setups left: search 2 retrieval, middle: search retrieval, and right: search 4 retrieval.
we gain. particular, for h search and retrievals, multi-head attention can only compute h possible search-retrieval pairings while the proposed model can compute h? possible pairings.
Controls, In all our experiments, we control for the number of layers, searches and parameters between the baseline and the proposed model.
B.6 DETAILS ABOUT SCALING
Suppose our systems have h heads for multihead attention and s searches, r retrievals for composi- tional attention. Lets assume the imput to the system has tokens. Then, we can see that the number of parameters in multi-head attention is proportional to 3/ while in compositional attention, it is proportional to (2s + r).
Further, focusing the highest computational cost of the attention procedure (which associated with the coefficients of 22 and ignoring the coefficients of terms lincar in N ), we see that the coefficient of quadratic complexity is proportional 2h in multi-head attention and s(1 + r) in compositional attention.
This shows that depending r, there be fewer parameters in the proposed model but the time complexity is strictly higher. This is because we allow for combinatorial more search-retrieval pairings and this cannot be obtained free of cost (no free lunch theorems).
Importantly, if a task that requires h searches (s = h) and h retrievals (r = h) and a dynamic choice of any search-retrieval pair out of all possibilities (which are h²),, then multi-head attention would require A²2 heads which leads to (3h²) parameters and (24²) computational cost to fully do this task well while the proposed model would only require 3h parameters and (Λ² + h) computational cost, which would be significantly smaller and also more computationally efficient when compared to multi-head attention. We use this exact motivation best as we could in the Contextual Retrieval Task to showcase the fundamental differences and the parameter efficiency.
B.7 ABLATIONS
Retrievals. For a number of tasks, we keep the number of searches the same as the number of heads in the baseline and then ablate over the number of retrievals. Overall, we notice that most of the models outperform baselines with fewer number of retrievals highlighting the combinatorial advantage of compositional attention. Please refer Table 1, 4, 2, 6, etc. We further ablate on the model capacity and heads while keeping the searches and retrievals fixed for the proposed model. In Table 7, we demonstrate that the head redundancies indeed hurt the performance of the model however the flexible models like compositional attention perform better.
Compositional-MLLR As additional ablation of the retrieval selection mechanism, we aim to replace the dot product attention in Equation 13 for the computation of Value Scores with a Multi- Layer Perceptron (MLP) that takes the retrieval query and key as input and outputs the retrieval attention score. Our MI.P is only a linear network but we still generally notice decent performance, highlighted in Table 8 and 9.
B.8 LIMITATIONS AND FUTURE WORK
Compositional attention, while more flexible than standard multi-head attention, is still conditional on the mechanism's ability efficiently perform value retrieval. In particular, do not explicitly
19

---

Published as a conference paper at ICLR 2022
impose bottleneck on the different value matrices pressure the system into learning diverse retrievals, which when combined with joint learning of the retrieval selection mechanism can lead to sub-optimal solutions from gradient-based learning. This is highlighted in Figure 12. Potential future work aims at solving this issue.
Another interesting direction for future development would be to dynamically restrict the set of retrievals accessibly by a search to a sparse number in between 1 (standard Transformer) and all (our compositional attention). This would allow to fine tune the tradeoff between complexity and expressivity for the task at hand.
C
EXPERIMENTS
In this section, we provide further details about the tasks outlined in Section 4. We also provide task-specific architecture details below.
C.I CONTEXTUAL RETRIEVAL TASK
We define the main details for the task in Section 4.1 and provide the additional required details and motivation below.
Motivation. Our aim was with this experiment was to de- sign an experiment that consists of multiple objects, each
Train Set
Test Set
of which have number of scalar features. Correspond- ing to each object, we have a set of ground truth searches of ground truth
(0, 0) (0, 1), (0, 2), (0,
(2, I)
(1. 0), (1, 1), (1, 2), (1,3)
(2, 3)
defined by the search features and trievals defined by retrieval features. The choice for which retrieval pick is decided by the retrieval preference of the object. This is very similar to real world scenarios where objects have a number of features, and search and the
(2, 0), (2, 2)
(3, 1)
(3, 0), (3.
(3, 3)
Table 5: Contextual Retrieval Task OoD Setup. For 2-Search-4-Reree variant there are 42 unique value binations that define tasks. For the Ool) setup, we pick fraction of tasks for train- ing set and the rest for the test set.
retrieval can be done about any feature depending context. Here, we make our task easier by defining a fixed search for all objects (i.e. no search preference), providing the retrieval preference as a one-hot signal and considering only independent scalar features.
Dataset. We consider number of search-retrieval combination settings for this task. We draw ag 1) and the features 2 ~ NN(0, 1) and consider the setups of search 4 retrievals, search retrievals and 4 search 16 retrievals.
OoD Setup. For the OoD setup, show certain subset of all possible value combinations for training, and the rest for zero shot testing. We illustrate an example of train and test retrieval combinations in Figure 5 for the 2 search 4 retrieval variant.
Implementation Details. For each variant of the task, we ablate number of heads and transformer hidden di- mensions. We opt to exclude the residual network because isn't needed for the task, purely from the data generat- ing perspective. For the search, we explicitly mask out the diagonal since the ground truth search defined arg min and along with all the values, we also feed in an extra value from each hidden dimension because it contains information about the retrieval preferences.
Vatue
0.01
0.50
0.48
Valuee 2
0.00
0.00
0.52
0.48
CoarTTrall Tran
-13
Gmurd
Vauu 3
0.58
-11
Valse
0.00
0.50
0.50
Quantitative Results. We highlight the results of our ablations in Table 6 which highlights that not only Com- positional Attention outperforms the standard Multihead Attention but also does across number of hyperparam- eter choices. In particular, we often see that the proposed attention mechanism outperforms the baseline even with lower number of parameters. Further, since we gener- ate data the fly, low training iterations correspond low data regime this task. Figure 10
Berreval
Remerrl
Rameval
Retrieval
Figure 12: Contextual Retrieval Task. Failure case of proposed mechanism in specializing retrievals in 1 search retrieval setup.
20

---

Published as a conference paper at ICLR 2022
Ground Trush
Algoridhm
Dimenssio
# Poramer
Namber of Hends
Number of Scurches
Number of Retrienuls
Losx (te distributionn
Lors Ood
30k
2
0.28
1.00
64
28k
Mahi-hhed
27k
8
0.59
0.96
Search Retrioval
117k
2
0.08
0.95
128
105k
0.14
0.95
Compositiona
64
30k
2.
0.10
0.28
128
113k
4
0.81
1.05
128
109k
8
0.80
0.93
106k
32
139k
0.54
0.89
Multi head
256
422k
8
0.57
0.81
Search
410k
32
0.86
Retrieval
1.7M
0:45
0.70
512
1.M
8
0.24
0.67
0.75
128
118k
0.48
0.47
Compositional
256
458k
0.51
0.18
512
0.18
118k
0.91
1.00
128
112k
16
0.95
1.00
110k
449k
0.77
Malti-heod
256
121k
16
0.84
0.96
Search 16 Retrieval
418L
64
0.98
1.7M
0.59
512
1.6M
16
0.87
0.83
166
0.90
128
116k
0.57
0.88
Compositional
256
442k
16
0.39
0.21
512
Table 6: Performance on the Contextual Retrieval task. Performance for different number of searches and retrievals in ground truth data. Ablations are done on the number of heads and dimensionality of the transformer dimension.
demonstrates that Compositional Attention not only converges faster but also does better in low-data regime.
Qualitative Results. We visualize the activation pattern our proposed model for the different task setups in Figure 11 which shows that the model often specializes its retrievals according to the ground-truth (up to a permutation). We further see in Figure 12 that sometimes the specialization does not occur which is what discuss in Section 5.
C.2 SORT-OF-CLEVR
Dataset. Sort-of-CLEVR consist of 10k images, where each image is accompanied with 10 Non- relational and 10 Relational questions. Non-relational questions tests the model ability focus the properties a "single" object like shape, horizontal location and vertical location. A few examples for non-relational questions are (i)what is the shape of green object? (ii) what is the horizontal location of the red object? (iii) what is the vertical location of yellow object? etc. On the other hand, the relational questions tests the model's ability to reason about the attributes of one or more objects. A few examples for relational questions are (i) What is the shape of the object that is farthest frrm the green object?, (ii) What is the shape of the object that is cloxest to the red object?". A sample of the dataset is shown in Figure 13 Each input image is of size 75 x 75 and it contains 6 square or circular objects. Each object is colored using one of these 6 different colors (red, blue, green, orange, yellow, gray) make them visibly distinct. The accompanied questions (relational and non-relational) are encoded vector of size 11 where the first six entries encode the color information, the next two encode the question type and the last three encode the question sub-type all in one-hot way.
Implementation Details. We use a 4-layered transformer with shared parameters and ablate with transformer dimensions 32, 256 and 5 512 and ffn dimension 64, 512, 1024 respectively. We consider baseline with 4 and heads and for the proposed model, we use 4 searches and ablate on retrievals. We use 32 dimensions for the retrieval query and key dimensions. We train the model
21

---

Published as a conference paper at ICLR 2022
Non-relational question
Q: Is the green object on the left or on the right?
A: right
Relational question
Q: How many objects have the shape of the orange object?
A: 3
Figure 13: Sort-of-CLEVR, Samples from the dataset. Non-relational refers to unary type questions and Relational refers binary and ternary type questions. Source: Santoro et al. (2017)
Algorithm
Dinensions
Heads
Unary Accuracy
Binory Accuracy
Temary Accuracy
Transformer
2
66.2
72.8
54.5
32
Compositional Transformer
74. 113.2
73.7
53.6.08
4
98. 98.6200
84,4
64.943.3
Transformer
8
98. g:02
84.5
65.444.7
Compositional Transformer
98. 1..1
88.2
66.9
4
98.5
84.2
61.5
Transformer
512
8
98.5
81.5
62.2
Compositional Transformer
98.9
84.5
62. 12.1
Table 7: Dimensions and Heads Ablation on Sort of CLEVR. We perform ablations with increased number of dimensions and heads. For proposed model, dimensional model and 4 searches retrievals for other dimensions.
searches 2 retrievals for 32
with 0.0001 learning rate for 100 epochs. For all our experiments, we report the mean and standard deviation over seeds.
Quantitative Results. Apart from the results presented in the main paper, we also ablate on different capacities and different number of heads in standard transformer. We highlight in Table 7 that our proposed attention mechanism works well across differently sized models consistently and outperforms baselines with more number of heads. We further illustrate Table that Compositional Attention also outperforms the standard models with fewer searches and performs well even when the retrieval dot-product attention is replaced by the MLP based attention as outlined in the ablations Appendix B.7.
0.72
Unory
0.64
Binary
0.40
Temary
0.32
Retriesal
Retrievil 2
Qualitative Results. Here perform analysis the learned models on the Sort-of-CLEVR dataset. We average the value and plot value
Figure 14: Sort-of-CLEVR Re- trieval Activation. Activation statistics against the different types of questions in the dataset.
scores in Equation 13 over the entire activation statistics of retrievals vs. the type of questions from the task database Figure 14. We note that there are distinct differences in the activation pattern between unary, binary and ternary type questions.
C.3 EQUILATERAL TRIANGLE DETECTION
This is binary classification task introduced by Ahmad & Omohundro (2009), where the aim is to decide if the set of three point clusters in the image forms an equilateral triangle. The dataset consists of 64 64 images with three randomly placed point clusters. For triangle be equilateral, the midpoints of these clusters should be equidistant from each other.
22

---

Published as a conference paper at ICLR 2022
Algorithm
Searches
Retrievals
Unary Accuracy
Binary Accuracy
Termary Accuracy
98.6g02
84.4
64.9
Transformer
98.5
84.51660
65.4
98.7.m02
86,8
66.4
2
98.8
88.2
66.9
Compositional Transformer
98.9.00
89.8
67.1
98.6.03
84.9
64.8_.41
98.7
85.8
63.8
2
98.8.03
84.7
65,3g37
Compositional Transformer MLP
98.810.15
84.6
66.0.26
4
98.8
87.8
66.2
Table 8: Compositional Transformer MLP Ablation on Sort of CLEVR. We highlight that our proposed model outperforms the baseline across the different question types even with lower number of searches and/or retrievals even with the MLP ablation where the retrieval attention score is computed by an MLP instead of dot-product acttention.
Implementation Details. We follow the same setup Dosovitskiy et al. (2020) and treat this as a classification task. Each image is split into 4x4 patches which is encoded through MLP and incorporated with position encodings. These then fed 4-layered transformer (with parameter sharing) along with the CLS token which acts the classification head.
Algoriilm
Senrches
Retrievals
Tesi Accarssy
Transformrr
4
93.8gr
1
95.6.11
Compositionl Transfommer
4
2
96.7
97.0gcx
96.8.05
Compositional Transformer- MIl
2
96.1
95.3
Table 9: Performance on Equilateral Triangle Detection. We perform ablations over the number of retrievals and type of attention mechanism used retrieval. Our proposed method outperforms the baseline over multiple different setups.
We set the hidden dimension the transformer 256 and the fín dimension as 512. We use 4 heads for the baseline and for the proposed model
4 searches and ablate over 1, and 4 retrievals. The retrieval queries and keys are 32 dimensional (used in Equation 13). All the models are trained with learning rate of 0.0001 for 200 epochs with cosine annealing, similar to Goyal et al. (2021b). For all our experiments, we report the mean and standard deviation over seeds.
Qualitative Results. We illustrate in Figure 15 that Compo- sitional Attention not only disentangles search and retrieval but can also allow for tighter pairing between them if need be, depending on the gradient signal. Thus, we see that leads to increased capacity because the pairing matrix (between search and retrieval) is dynamic, context dependent and can be more than just identity matrix.
Search
0.25
0.11
0.6
0.037
03
Sabarth
0.29
0.085
0.4
-05
Smarch
0.3
0.33
0.11
0.26
Quantitative Results. We see in Table that the proposed compositional attention outperforms the multi-head attention baseline not only across different number of retrievals but also on the MLP (Appendix B.7) ablation where the retrieval selec- tion score is obtained through simple linear network instead of the dot-product attention mechanism. This shows that even different ways of having this flexible composition of search and retrieval (apart from dot-product based) can lead to improved performance.
Restiemal
Autseve
Figure 15: Search-Retrie Pair- ing in Equilateral Triangle Detec- tion. We visualize the activation statistics of learned retrievals (X axis) against learned searches show- ing that Compositional Attention can also learn tight pairing between search and retrieval.
C.4 MULTI-TASE IMAGE CLASSIFICATION
Dataset. This task is composed of four datasets CIFAR1O, FahsionMMNIST, SVHN and Equilateral Triangle Detection. Given an image from one of these datasets, the model is tasked to make the respective classification prediction. We random crop and horizontal flips data augmentation
23

---

Published as a conference paper at ICLR 2022
(a) Same/different
(b) RMTS
(c) Distribution-of-three
(d) Identity rules
Figure 16: Logical Reasoning in ESBN Tasks. Illustration of the four tasks in the suite. (a) Same/Differett. Predict whether the two objects are identical or not; they are not in this case, (b) RMTS. Match the relation in the context image the two choices; here option 2 is the right answer since similar to context objects, it also has two identical objects, (c) Distribution of 3. Find the missing object with permutations rule; option is the right answer since the square is in the context objects and missing from the second row, and (d) Identity Rules. Find the missing object with ABA rule; option is the correct answer since follows the same rule of identical objects on the edges. Source: Webb et al. (2020)
techniques. This task tests the model's ability to learn multi-purpose representations that can be used for performing well on multiple tasks that potentially share meaningful information.
Implementation Details. We train 2-layered universal transformer with 4 heads. We set the hidden dimension of the transformer as 256 and the ffn dimension as 512. For our proposed model, we keep all the settings the same, use 4 searches, and ablate over the number of retrievals. The image is cropped into 4 x 4 patches and then fed into an encoder and augmented with positional information. For prediction on each dataset, there is learnable embedding that gets fed into the transformer along with the cropped patches. The retrieval queries and keys are 16 dimensional and we use a learning rate of 0.0001. For experiments, we report the mean and standard deviation over seeds.
Quantitative Results. We showcase in Table 2 that the proposed model outperforms the standard multi-head attention across different number of retrievals.
C.5 LOGICAL REASONING IN ESBN TASKS
Dataset. This is suite of four tasks Same/Different, RMTS, Distribution of 3, and Identity Rules (Webb et al., 2020). For each task, the model gets sequence of objects as input with certain relationship between them. This relationship depends upon the task and based on it, the model is tasked to make a prediction. Further, the dataset consists of an OoD setup where the phrase m90, for example, means that training is done on 10 unique objects and testing is done on the remaining 90 objects.
Same/Different. The model gets two objects input and is tasked with predicting whether the two objects are identical or not.
RMTS. Similar to Same//ifferent, we get three pair of objects. The first pair is the context where the first two objects can be either same or different. The second and third pair are the options and the model is tasked with choosing the pair which follows the same relationship as the first pair. That is, the first pair has identical objects, then the model should pick the pair with identical objects from p and pair-3.
Distribution 3. We first get three objects as input and then get a permutation of those objects as next inputs with the last object hidden. The model is tasked with predicting which object should fill the last location from a candidate set of four objects that are shown next. In short, the model must choose an object from multiple choices which would make sure that the first three objects and the next three objects are permutations of each other.
Identity Rules. The input is given in the form of ABA where A and are some unique objects. Similar to Distribution of 3, the model then gets two more objects CD and has choose object C from candidate set of four choices. That is, the model is tasked to learn the rule of identical peripheral objects in filling up of the final spot.
24

---

Published as a conference paper at ICLR 2022
Task: Same/Different
Task: RMTS
100
80
Accuracy
60
40
20
Model
Multi-Head
Compositional
m95
Compositional
Task: Distribution of
Task: Identity Rules
Compositional
100
Compositional-6
Compositional-8
Accuracy
40
20
m0
m50
m85
m95
m0
m50
m85
m99
Figure 17: Logical Reasoning in ESBN Tasks. We see that compositional attention outperforms multi-head attention baseline different number of retrievals, especially Same/Different and RMTS. Compositional-- refers to the proposed model with retrievals.
We refer the readers to Figure 16 for visual demonstration of the four tasks.
Implementation Details. We follow the setup of Webb al. (2020) and train a single layered transformer model as baseline and replace multi-head attention with compositional attention in our proposed model. We compare the head baseline model with 8 searches in the proposed model ablated over different retrievals. We use learning rate of 0.0001 and follow the initialization as well as normalization schemes of the original codebase. We report the mean performance with standard deviation over 10 seeds for each model.
Quantitative Results. notice substantial improvements of the proposed model over the baseline on the Same/Different and RMTS task as illustrated in Figure 17, especially in the OoD setting. We see marginal improvements on the Distribution of 3 and Identity Rules too.
C.6 SCAN TASK
Dataset. The SCAN task (Lake & Baroni, 2018 8a) is aimed at systematically testing OoD performance of various models. The task requires models to translate natural language inputs. eg. "jump opposite left and walk thrice", into a sequence of actions for a robotic agent to follow, eg. "L.TURN LTTURN JUMP WALK WALK WALK". In particular, we use the length extrapolation version of this task where models trained shorter action sequences and then evaluated on longer action sequences. Thus, this task tests the capabilities of systematic generalization through compositions of primitive actions. In particular, Csordás et al. (2021) which aim to solve certain problems of the original split.
follow the task split paradigm set in Newman et al. (2020);
Implementation Details. We follow the implementation of Csordás al. (2021)* and train 3- layered Universal Transformer model with heads. We set the model dimensions as 128, the ffn
http::///gthabhccm//////wwww//e/ symbols http://ggttaaacccmmmmmmmettttttt genprrraaiateeggeeeiaiiiiita
25

---

Published as a conference paper at ICLR 2022
dimensions as 256 and the retrieval query and key dimensions 32. All the models trained with the learning rate 0.001.
Quantitative Results. The results for the task showcased in Table which illustrates that the proposed model outperforms the baseline over multiple cutoff lengths, showing that the proposed model is able to better generalize to compositional tasks.
C.7 LANGUAGE MODELLING
We perform experiments on the 1i0 data corpus (Merity et al., 2016) for the language modeling task. The corpus consists 28,475 articles in its training split and 60 the validation and test split respectively and the task is to predict probabilities for next words, evaluated by perplexity.
We use 6-layered transformer models with parameter sharing and perform experiments the fairseq codebase (Ott et al., 2019). We plot the validation perplexity against epochs in Figure 7 which highlights that our proposed attention mechanism not only outperforms the baseline but also converges faster. Further, we see that proposed model obtains test perplexity 08. as opposed to baseline's s perplexity 39.6
Quantitative Results:
Implementation Details. We the Ott et al. (2019)² repository for both pre- pre-processing of the data as well as training. We use a transformer based model with 512 hidden dimensions and 2048 fín dimensions. Our baseline has 8 heads, and proposed model has searches and 8 retrievals, and trained with 32 dimensional retrieval queries and keys. We perform a hyperparameter sweep on learning rates 0.0005, 0.001, 0.002, 0.004 and 0,007. For both the baseline and the proposed method, we choose the model with best validation perplexity (baseline 0.002 and proposed 0.004) and then run additional seeds on the chosen learning rate.
26


Reviews:
[{"review_id": "ft6vhJKyAuj", "reviewer": "Reviewer_f5jQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength: \n\nThe author presents a comprehensive analysis of the limit of multi-head attention. They point out that the rigid mapping between query-key and value pairs can lead to the learning of redundant query-key matrices and value matrices. Instead, they propose to consider all pairs between query-key and value. They define query-key as a search and value as a retrieval and then introduce new parameters to compute the attention score to compute the weighted sum over the pairs of search to all retrievals.\n\nIt compares compositional attention with multi-head attention on a set of tasks including a proof-of-concept contextual retrieval task, Sort-of-CLEVER, logical reasoning tasks, language modeling, etc. Compositional attention has shown consistent improvement on these tasks.  \n\nWeakness: \n\nThey argue that the complexity overhead of compositional attention is lightweight. More importantly, the overhead can be mitigated by reducing the number of retrievals. The appendix shows the overhead analysis on the contextual retrieval task. It would be much more helpful and convincing to show such analysis on other real world tasks.\n\nThe paper lacks the comparison with other related works focusing on head pruning. As the author pointed out, in figure 11, the compositional attention may tend to learn to focus only on several pairs of search and retrieval pairing. This is very relevant to other head pruning works. It’s worth noting what’s the difference between compositional attention and other head pruning techniques. \n\nAnother question remaining unclear is what kinds of tasks can benefit from compositional attention? Multi-head attention has been widely adopted in a lot of scenarios, such as pre-training, fine-tuning, various tasks with sufficient data or very little data. Though the author has tried to experiment on a set of tasks, it remains unclear what's the rule-of-thumb to replace compositional attention to multi-head attention.\n", "summary_of_the_paper": "This paper proposes a Compositional Attention to replace the standard multi-head attention. It argues that traditional multi-head attention has a rigid mapping between each query-key pair and the associated value. This leads to redundant parameters and less generalizability.\n", "summary_of_the_review": "The paper focuses on improving parameter redundancy and generalizability in multi-head attention. The compositional attention is well motivated to mitigate these issues. The paper is well written and has shown improvement of compositional attention. That said, it's lacking analysis of the constraints on computation overhead of compositional attention on more real-world tasks under both sufficient data and low data scenario. It's a bit unclear of the rule-of-thumb to replace the multi-head attention with compositional attention.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QGTF9RjcoEB", "reviewer": "Reviewer_TCch", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n- The paper identifies a current limitation of multi-head attention, and proposes an interesting method to overcome it.\n\n- The paper shows results on a wide variety of experimental tasks.\n\n- The paper is well-written and thorough.\n\n- The illustrative experiment proposed in the paper is an intuitive, clear example of what types of problems Compositional Attention can help with\n\n\n### Concerns\n\nI would imagine that a Transformer without compositional search/retrieval could possibly compensate for that lack of single-operation compositionality through the stacked computation from the depth of the network.  However, most of the experiments use Universal Transformers as a baseline (or a single layer for the illustrative tasks).  Since the parameters are shared between layers, it seems like that might be a limitation compared to non-shared parameters when the focus is on compositional tasks.  Have you compared to Transformers without shared parameters between layers?  \n\nThe compositional attention adds overhead compared to a standard MHA.  Although the parameter efficiency argument is well made, there is little analysis of the computational tradeoffs.  How do the Compositional Attention models perform when matched against the baselines for runtime rather than parameter counts? (I understand that is a fraught task given different hardware optimizations/settings between models, but it is still an important consideration)  \n\nOn a related note, how might Compositional Attention scale to larger models with more searches and queries and more layers?  Does the runtime get much worse than the ~10% reported in the appendix when the models are closer to some of the bigger models reported in the literature?  I’m trying to evaluate how to think of the tradeoff between heavier attention operations run fewer times versus simpler operations with more depth.", "summary_of_the_paper": "This paper proposes a mechanism for dynamically assigning query-key attention masks (searches) to one of several value matrices (retrievals), avoiding the rigid assignment of search to retrieval inherent in standard multi-head attention.", "summary_of_the_review": "The paper presents an interesting idea, based on the promising property of compositionality, to increase the power of a single MHA computation and a resulting transformer architecture.  However, I am not quite convinced by the baselines and the analysis of the computational complexity.  I am happy to raise my score if my above concerns are addressed.\n\n\n### Update after author comments\n\nThank you for the clear, detailed response and the additional experiments.  You have addressed my concerns, and I will raise my score.  Please consider including the results of the additional experiments and the related discussion of compositionality vs depth in the paper itself, even if just in an appendix.    ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HX51q-GLArb", "reviewer": "Reviewer_kYB6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Intuitively, the suggested method makes sense: it provides additional flexibility in the attention mechanism, since now the same retrievers can be re-used by different searchers, and since a search is not confined to a single retriever. Such added flexibility and re-usability could improve the generalization abilities of the attention mechanism.\n\nThe authors do a good job of evaluating the proposed mechanism, with a diverse selection of tasks: Reasoning over images (sort of clevr), logical reasoning, SCAN (length split) and language modeling and an additional specifically-designed task that requires contextual retrieval. For some of these datasets, models are evaluated on OOD samples. Since generalization seems to be the main motivation for the proposed architecture design, adding additional OOD tests where possible (e.g. as far as I understand CLEVR was only tested in i.i.d setting) could have made the results stronger. Nevertheless, experiments convincingly show an improvement in almost all cases.\n\nMy main concern regarding the experiments is that in all cases, parameters were shared between all transformer layers. Since in most common usages of the models parameters are not shared, it seems important to test how well does this method work with non-shared parameters - would the compositional attention still work well with different search/retrieve parameters?\n", "summary_of_the_paper": "The proposed mechanism disentangles the search and retrieval from the transformers architecture, where the retriever is no longer tied to a specific retriever that is in the same head - instead, there are multiple searches available, and multiple retrievers. Conditioned on the input, each search is \"tied\" to a (soft) single retriever. Then all results from all searches are concatenated.", "summary_of_the_review": "This paper proposes a modification to the transformer's attention mechanism, which shows improvement on multiple evaluation benchmarks of various types. I find this proposal to be convincingly useful, although evaluation on more o.o.d splits and experiments with non-shared parameters could have made results stronger.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4pePdRYePnw", "reviewer": "Reviewer_RbM7", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n-  The problem studied in this paper is interesting and well-motivated.\n-  The paper points out the shortcomings of rigid search-and-retrieval coupling in standard multi-head attention.\n-  The proposed compositional attention mechanism is rational.\n\nWeaknesses:\n-  The evaluation tasks selected in this paper are somewhat weird. The author should conduct experiments on some traditional or mainstream tasks  (or settings) of language or vision data.\n- The paper should also compare with other variants of multi-head attention mechanisms.\n- The paper should explore the effect of the proposed compositional attention mechanism on the latest pre-trained language models.\n\n\n\n\n\n", "summary_of_the_paper": "The paper studies the multi-head attention mechanism in Transformer. The paper first analysis the potential drawback of the rigid mapping between search and retrieval in standard attention heads, and then proposes compositional attention that disentangles search and retrieval and composes them in a dynamic and context-dependent manner. To evaluate the proposed method, the author conduct experiments on the standard Transformer model and Vision Transformer.\n\nOverall, the paper is clearly written and the problem studied in this paper is interesting and well-motivated. The proposed compositional attention mechanism is rational. The experiment conducted in this paper is somewhat thorough.\n\n\n\n", "summary_of_the_review": "The paper provides a somewhat novel insight into the shortcomings of standard multi-head attention, and also proposed a compositional attention mechanism. However, there are still some concerns about the evaluation.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper identifies a limitation with current attention in transformers where they scoring with query-key pairs is strongly tied to retrieving the value and proposes a more flexible configuration that subsumes the previous setup but provides more flexibility. The authors shows this leads to improvements in various settings.

Overall, all reviewers seem to agree there is interesting insight and results in this paper and it merits publication. Also the discussion helped stress important points regarding weight sharing and more. One concern is that the model was not evaluated on standard NLP/vision datasets (I assume alluding to GLUE/SuperGlue/SQuAD, etc.), and authors seem to hint that pre-training this is an issue for them computationally. This leaves open whether this indeed can and should replace the standard attention mechanism across the board, but is still very worthy of publication.
Decision: Accept (Spotlight)