Systematic Generalization with Edge Transformers

1 Introduction . Transformers ( Vaswani et al. , 2017 ) have become ubiquitous in natural language processing and deep learning more generally ( e.g. , Devlin et al. , 2018 ; Raffel et al. , 2020 ; Carion et al. , 2020 ) . Nevertheless , systematic ( or compositional ) generalization benchmarks remain challenging for this class of models , including large instances with extensive pre-training ( Keysers et al. , 2020 ; Tsarkov et al. , 2020 ; Gontier et al. , 2020 ; Furrer et al. , 2020 ) . Similarly , despite their increasing application to a variety of reasoning and inference problems , the ability of Graph Neural Networks ’ ( GNN ) ( Gori et al. , 2005 ; Scarselli et al. , 2009 ; Veličković et al. , 2018 ) to generalize systematically has also been recently called into question ( Sinha et al. , 2019 , 2020 ) . Addressing these challenges to systematic generalization is critical both for robust language understanding and for reasoning from other kinds of knowledge bases . In this work , inspired by symbolic AI , we seek to equip Transformers with additional representational and computational capacity to better capture the kinds of information processing that underlie systematicity . Our intuition is that classical Transformers can be interpreted as a kind of inferential system similar to a continuous relaxation of a subset of the Prolog logic programming language . In particular , as we discuss in the next section , Transformers can be seen as implementing a kind of reasoning over properties of entities—for example , RED ( · ) or TABLE ( · ) —where any relations 1Code for our experiments can be found at : github.com/bergen/EdgeTransformer 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . required for such reasoning—for example PART-OF ( · , · ) or GRANDMOTHER ( · , · ) —are inferred on the fly with the self-attention mechanism . Building on this intuition , in this work we propose Edge Transformers , a generalization of the Transformer model that uses a novel triangular attention mechanism that is inspired by a much more general family of inferential rules . To achieve this , we endow the Edge Transformer with a 3D tensor state such that every edge , that is , every pair of input nodes , contains a vector that represents relations between the two nodes . The updates of each edge are computed using all adjacent edges in a way that is directly inspired by unification in logic programming . While the use of edge features or even dynamic edge states can be found in the Transformer and GNN literatures ( Shaw et al. , 2018 ; Gilmer et al. , 2017 ; Gong and Cheng , 2019 ) , to the best of our knowledge such triangular updates are novel to our approach . We evaluate the Edge Transformers on three recently proposed compositional generalization challenges . First , on the graph version of the CLUTRR relational reasoning challenge ( Sinha et al. , 2019 ) , our model displays stronger systematic generalization than Graph Attention Networks ( Veličković et al. , 2018 ) and Relation-aware Transformers ( Shaw et al. , 2018 ) . Second , on both dependency parsing and semantic parsing versions ( Goodwin et al. , 2021a ) of the Compositional Freebase Questions benchmark ( Keysers et al. , 2020 ) our model achieves a higher parsing accuracy than that of classical Transformers and BiLSTMs . Last but not least , Edge Transformer achieves state-of-the-start performance of 87.4 % on the COGS benchmark for compositional generalization in semantic parsing . 2 Intuitions . A Transformer operates over a set of n entities—such as words in a sentence—which we represent as nodes in a graph like those displayed on the left-hand side of Figure 1 . Each node is associated with a d-dimensional node-state X ( i ) for i ∈ n , which can be thought of as the output associated with applying some function to the node—that , a representation of node properties . In transformers , computation proceeds by sequentially associating each node with a number l of node-states , with each state updated from node states at the preceding layer via the attention mechanism . Adopting a Prolog-like notation , we could write the fundamental inferential process implemented by the transformer architecture as X l+1 ( 1 ) ` A X l ( 1 ) , X l ( 2 ) , . . . , X l ( n ) . ( 1 ) In logic programming , the turnstile symbol ` means that whenever the right-hand side of the expression is true , the left-hand side must be true . In our interpretation of Transformers , the inference rules expressed by ` A are learned via the attention mechanism A , as is the meaning of each property X l ( · ) . Classical transformers can therefore be interpreted as an architecture specialized to learning how entity properties can be inferred from the properties of other entities . Despite the power of this mechanism , it still has noteworthy limitations . In particular , for many practical reasoning and NLU tasks systems must learn dependencies between relations , such as family relations : MOTHER ( x , y ) and MOTHER ( y , z ) implies GRANDMOTHER ( x , z ) . Such a general reasoning problem might be expressed in Prolog-like notation as follows . X l+1 ( 1 , 2 ) ` A X l ( 1 , 1 ) , X l ( 1 , 2 ) , . . . , X l ( 2 , 1 ) , X l ( 2 , 2 ) , . . . . ( 2 ) It is of course possible for classical transformers to capture such reasoning patterns . But for them to do so , each transformer state X l+1 ( 1 ) must encode both the properties of the relation itself , that is , MOTHER ( · , · ) versus GRANDMOTHER ( · , · ) —as well as all target nodes x with which node 1 stands in the relation , that is MOTHER ( 1 , x ) . In other words , to encode relations a classical transformer must use its state representations for two distinct purposes . This mixing of concerns places a high burden on the learning , and we hypothesize that the resulting inductive bias hinders the ability of this class of models to generalize systematically . To address this problem we propose that the fundamental object represented by the model should not be a state associated with a node but , rather , states associated with edges , like those displayed on the right of Figure 1 . Such edge representations in turn are updated based on an attention mechanism which attends to other edges . This mechanism is inspired by the process of unification in logic programming . Critical to an inference rule such as GRANDMOTHER ( x , y ) ` A MOTHER ( x , z ) MOTHER ( z , y ) is the fact that the two predicates on the right hand side of this rule share a variable z . Successfully inferring GRANDMOTHER ( x , y ) which involve finding a binding for z which satisfies both MOTHER ( z , y ) and MOTHER ( x , z ) —a process which is handled by unification in logic programming . We believe that this is a fundamental aspect of relational reasoning and to capture this we make use of a form of attention which we call triangular attention and which we further describe in the next section . 3 Edge Transformers . An edge transformer operates on a complete graph with n nodes and n2 directed edges ( we allow self-connections ) labeled with d-dimensional feature vectors . An edge transformer state therefore is a 3-dimensional tensor X ∈ Rn , n , d that consist of edges states X ( i , j ) corresponding to every edge ( i , j ) , i , j ∈ [ 1 , n ] . We will also write xij for each edge state . An edge transformer layer computes X ′ = FFN ( LN ( X +TriangularAttention ( LN ( X ) ) ) ) . ( 3 ) Here , FFN ( X ) is a fully-connected feed-forward network with one hidden layer of 4 · d units that is applied independently to each edge state xij , LN stands for the layer normalization mechanism ( Ba et al. , 2016 ) which rescales activations of each feature f ∈ [ 1 , d ] across its occurrences in edge states xij . Equation 3 follows the familiar structure of transformer ( Vaswani et al. , 2017 ) computations with two important differences : ( a ) It uses a 3D-tensor state instead of a matrix state ; and ( b ) it makes use of a novel triangular attention mechanism that we describe below . For a single edge state xij a single-head triangular attention update outputs a vector aij that is computed as follows : aij =W o ∑ l∈ [ 1 , n ] αiljvilj , ( 4 ) αilj = softmax l∈ [ 1 , n ] qilklj/ √ d , ( 5 ) qil =W qxil , ( 6 ) klj =W kxlj , ( 7 ) vilj = V 1xil V 2xlj , ( 8 ) where stands for elementwise multiplication , and W q , W k , W o , V 1 , V 2 ∈ Rd , d are matrices that are used to produce key , query , output and value vectors kil , qlj , aij and vilj respectively . Here and in the rest of the paper we omit bias terms for clarity . Informally , updates associated with an edge ( i , j ) proceed by aggregating information across all pairs of edges that share a node l , that is , all pairs ( i , l ) and ( l , j ) . The updated edge value , aij , is an attention-weighted mixture of contributions from each such pair of edges . Figure 1 visualizes some of the key differences between transformer and edge transformer computations . Note that edges ( i , j ) , ( i , l ) , and ( l , j ) form a triangle in the figure—hence the name of our attention mechanism . We define a multi-head generalization of the above mechanism in a way similar to the implementation of multi-head attention in vanilla transformers . Specifically , each head h ∈ [ 1 , m ] will perform Edge 1 Attention using the smaller corresponding matrices W q , W k , V 1 , V 2 ∈ Rd , d/m , and the joint output aij will be computed by multiplying a concatenation of the heads ’ output by W o : aij =W o [ a1ij ; . . . ; a m ij ] . ( 9 ) In practice a convenient way to implement an efficient batched version of edge attention is using the Einstein summation operation which is readily available in modern deep learning frameworks . Tied & Untied Edge Transformer Edge Transformer layers can be stacked with or without tying weights across layers . By default we tie the weights in a way similar to Universal Transformer ( Dehghani et al. , 2019 ) . We will refer to the version with weights untied as Untied Edge Transformer . Input & Output The Edge Transformer ’ s initial state X0 is defined differently for different applications . When the input is a labeled graph with ( possibly null ) edge labels , the initial state x0ij is produced by embedding the labels with a trainable embedding layer . When the input is a sequence , such as for example a sequence of words or tokens ti , self-edges x0ii are initialized with the embedding for token ti , and all edges x0ij are initialized with a relative position embedding ( Shaw et al. , 2018 ) : x0ij = { e ( ti ) + a ( 0 ) if i = j a ( i− j ) if i 6= j ( 10 ) where e is a trainable d-dimensional embedding layer , and a ( i− j ) is a relative position embedding for the position difference i− j . To produce a graph-structured output with Edge Transformer , a linear layer with cross-entropy loss can be used for every edge ( i , j ) . For classification problems with a single label for the entire input the label can be associated with one of the edges , or edge representations can be pooled . Lastly , an encoder-decoder version of Edge Transformer can be used to perform sequence-to-sequence transduction . Our approach follows that used with the original Transformer architecture ( Vaswani et al. , 2017 ) . First , an encoder Edge Transformer processes the question and produces a 3D tensor state xenc . Next , the decoder Edge Transformer generates output tokens one by one , left-to-right . To produce the output distribution for the token we feed the state of the loop edge ( i , i ) through a linear layer . To connect the decoder to the encoder we insert xenc into the decoder ’ s initial state . This operation is the Edge Transformer ’ s equivalent to how the Transformer decoder cross-attends to the Transformer encoder ’ s hidden states . To compute training log-likelihoods we adapt the Transformer decoder ’ s causal masking approach to our triangular attention mechanism .

Summary: This paper introduces a variation of the Transformer architecture where attention happens between edges as opposed to nodes. As such, the input and output to the Edge Transformer are 3D $d \times n \times n$ tensors instead of 2D $d \times n$ tensors. This new architecture, together with its novel “triangular attention” mechanism, are introduced with the objective of favoring generalization in compositional reasoning. Two sets of experiments validate the architecture’s benefits over previous architectures.