Compositional Attention: Disentangling Search and Retrieval

1 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 et 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 a 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 is a simple idea : enable task-driven flexible connections between elements of a 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 a 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 run 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 out-of-distribution ( OoD ) generalization ( Lake & Baroni , 2018b ; Liska et 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 1 and Section 2.3 . †Correspondence authors sarthmit @ gmail.com and g.lajoie @ umontreal.ca 1Open-sourced implementation is available at https : //github.com/sarthmit/Compositional-Attention Addressing these shortcomings , there are several recent attention-enabled systems developed to allow better decomposition and re-composition of knowledge ( Goyal et al. , 2019 ; 2021a ; b ) , some of which we discuss in Appendix A . However , most of these efforts hinge on purpose-built 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 to 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 a mechanism that is easily implemented and plug-and-play for existing Transformers ( and all the models based on them ) . We propose Compositional Attention , where the search and retrieval operations can be flexibly composed : the key-query search mechanism is no longer bound to a fixed value retrieval matrix , instead it is 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 of rigid search-and-retrieval 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 a number of experiments ( Section 3 and 4 ) . ( c ) Through a series of analyses , we demonstrate how our proposed attention mechanism decomposes relational task structure as 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 ) in 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 ∈ RN×d , we first obtain queries Q , keys K and values V using linear transformations on X with learnable projection matrices Wq ∈ Rd×dk , Wk ∈ Rd×dk and Wv ∈ Rd×dv respectively . This is given by Q = XWq K = XWk V = XWv . ( 1 ) For each query , a similarity score is computed with each key using a scaled cosine similarity ( called scaled dot-product ) to give the attention weights which are used to soft-combine the values as Attention ( Q , K , V ) = Softmax ( QKT√ dk , axis = ‘ keys ’ ) V ( 2 ) where 1√ dk is the scaling factor . Multi-Head Attention : A multi-head attention mechanism combines multiple ( say , h ) independent key-value attention mechanisms in parallel to provide the model the ability to 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 o ∈ Rhdv×d : Multi-Head = Concat ( head1 , head2 . . . headh ) W o ( 3 ) where headi = 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 fundamental components - Search and Retrieval . Search : A search is parameterized by the query and key matrices , that is , Wq and Wk respectively . These parameters define a notion of compatibility metric between pairs of element xj and xk ∈ X : Search ( Q , K ) = Softmax ( QKT√ dk , axis = ‘ keys ’ ) ( 4 ) where Q = XWq and K = XWk . The above computation gives the compatibility between an element xj with other elements xk ’ s under the compatibility metric defined by the Search parameters . Retrieval : A retrieval is parameterized by a value matrix Wv describing the kind of features in the input elements in X that are relevant and need to be accessed for the downstream task : Retrieval ( Search ( Q , K ) , V ) = Search ( Q , K ) V ( 5 ) where V = XWv . Note that each Retrieval defines the kind of attributes to access from input x′ks and can take any Search result as its input . Multi-head Attention as a rigid pairing of Searches and Retrievals : Given the above definitions , one can see how standard multi-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 ith retrieval is performed only on the ith search . Multi-head attention thus amounts to a special case of Equation 4 and 5 headi = Retrieval ( Search ( Qi , Ki ) , Vi ) ( Note : Same subscript i ) ( 6 ) Viewing multi-head attention through these fixed search-retrieval pairings foreshadows a possible generalization of searches and retrievals which we 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 a 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 a feature ( defined by query-key matrices Wq and Wk ) based on which compatibility between objects is evaluated . Further , each head ’ s retrieval allows the model to access a particular feature ( defined by value matrix Wv ) 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-key matrices and ( b ) Retrieval Redundancy which leads to learning of redundant value matrices . We highlight these two redundancies jointly using a simple example illustrated in Figure 1 , 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 to 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 1 and retrieving location exists in ( b ) learned by Head 2 , there is no way to combine them . Thus , another head is needed to obtain the search of Head 1 and retrieval of Head 2 . This leads to parameter redundancy and a missed opportunity to factorize knowledge more efficiently , since these pieces of learned knowledge individually exist in Head 1 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 B.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 h 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 × R such pairings , with S the number of search types and R the number of retrieval types . 3 COMPOSITIONAL ATTENTION - DISENTANGLING SEARCH AND RETRIEVAL . We propose a novel attention mechanism that relaxes static search-retrieval pairing in favour of a 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 . Similar to heads , we start by defining S parallel search mechanisms . That is , we have S different query-key parameterizations Wqi and Wki respectively . The output of each search is as defined in Equation 4 . In essence , for each search i , we get Qi = XWqi and Ki = XWki ( 7 ) Searchi = Softmax ( QiK T i√ dk , axis = ‘ keys ’ ) ( 8 ) Next , we define R different retrieval mechanisms , which correspond to the R different Wvj matrices . These matrices are used to obtain different attributes from the input . Formally , this is summarized as Vj = XWvj ( 9 ) where Vj highlights accessing of different attributes . Then , corresponding to each search , all possible retrievals are done . This is similar to Equation 5 and is defined as Retrievalij = Searchi Vj ( 10 ) for all i , j . This step gives us all the hypothetical retrievals for each search and out of this , one retrieval per search needs to be instantiated . This instantiation is done through a secondary attention mechanism computed using retrieval queries Qi and retrieval keys Kij , that are obtained as Qi = XW qi Kij = Retrievalij W k ( 11 ) where the parameter W qi ∈ Rd×dr is a different matrix for each search indexed by i and , together with W k , is tasked with driving the pairing between search and retrieval . We broadcast the matrix Qi from RN×dr to RN×1×dr and define Ki ∈ RN×R×dr by Ki = Concat ( Ki1 , Ki2 , . . . , KiR ) . ( 12 ) Thus , through these retrieval queries and keys , the required instantiation per search is done as CAtti = Softmax ( QiKi T √ dr , axis = ‘ retrieval ’ ) ︸ ︷︷ ︸ Value Scores Retrievalij ( 13 ) 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 of these parallel searches are combined by concatenating them and passing them through a linear network : CompositionalAttention = Concat ( CAtt1 , CAtt2 , . . . , CAttS ) W o ( 14 ) where W o ∈ RSdv×d . Note that in this mechanism , the choice of retrieval for each search is not fixed , and instead is dynamically modulated by Qi and Ki respectively . We refer the readers to Figure 2 for a visual depiction of the computation graph . Compositional Attention allows the model to have ( a ) Different number of searches and retrievals , i.e . S andR respectively , ( b ) Dynamic selection of shared retrievals for each search and ( c ) Representation capacity of S × R ( Search – Retrieval ) pairings . Thus , we highlight that Compositional Attention disentangles search and retrieval and solves the redundancies of multi-head attention ( Section 2.3 ) .

Summary: 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.  Overall, 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.