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 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.