Logic and the 2-Simplicial Transformer

1 INTRODUCTION . Deep learning contains many differentiable algorithms for computing with learned representations . These representations form vector spaces , sometimes equipped with additional structure . A recent example is the Transformer ( Vaswani et al. , 2017 ) in which there is a vector space V of value vectors and an inner product space H of query and key vectors . This structure supports a kind of messagepassing , where a value vector vj ∈ V derived from entity j is propagated to update an entity i with weight qi · kj , where qi ∈ H is a query vector derived from entity i , kj ∈ H is a key vector derived from entity j , and the inner product on H is written as a dot product . The Transformer therefore represents a relational inductive bias , where a relation from entity j to entity i is perceived to the extent that qi · kj is large and positive . However , the real world has structure beyond entities and their direct relationships : for example , the three blocks in Figure 1 are arranged in such a way that if either of the supporting blocks is removed , the top block will fall . This is a simple 3-way relationship between entities i , j , k that is complex to represent as a system of 2-way relationships . It is natural to make the hypothesis that such higher-order relationships are essential to extracting the full predictive power of data , across many domains . In accordance with this hypothesis , we introduce a generalisation of the Transformer architecture , the 2-simplicial Transformer , which incorporates both 2- and 3-way interactions . Mathematically , the key observation is that higher-order interactions between entities can be understood using algebras . This is nothing but Boole ’ s insight ( Boole , 1847 ) which set in motion the development of modern logic . In our situation , an appropriate algebra is the Clifford algebra Cl ( H ) of the space H of queries and keys , which contains that space H ⊆ Cl ( H ) and in which queries and keys can be multiplied . To represent a 3-way interaction we map each entity i to a triple ( pi , l1i , l 2 i ) of vectors inH consisting of a query vector pi , a ( first ) key vector l1i and a ( second ) key vector l 2 i . Given a triple i , j , k we first form the product pil1j l 2 k in the Clifford algebra , and then extract a scalar quantity η ( pil 1 j l 2 k ) using a natural continuous function η : Cl ( H ) −→ R associated to the Z-grading of Cl ( H ) . This scalar ∗Listing order is alphabetical . Correspondence to d.murfet @ unimelb.edu.au . measures how strongly the network perceives a 3-way interaction involving i , j , k. In summary , the 2-simplicial Transformer learns how to represent entities in its environment as vectors v ∈ V , and how to transform those entities to queries and ( pairs of ) keys in H , so that the signals provided by the scalars qi · kj and η ( pil1j l2k ) are informative about higher-order structure in the environment . As a toy example of higher-order structure , we consider the reinforcement learning problem in a variant of the BoxWorld environment from ( Zambaldi et al. , 2019 ) . The original BoxWorld is played on a rectangular grid populated by keys and locked boxes of varying colours , with the goal being to open the box containing the “ Gem ” . In our variant of the BoxWorld environment , bridge BoxWorld , the agent must use two keys simultaneously to obtain the Gem ; this structure in the environment creates many 3-way relationships between entities , including for example the relationship between the locked boxes j , k providing the two keys and the Gem entity i . This structure in the environment is fundamentally logical in nature , and encodes a particular kind of conjunction ; see Appendix I . The architecture of our deep reinforcement learning agent largely follows ( Zambaldi et al. , 2019 ) and the details are given in Section 4 . The key difference between our simplicial agent and the relational agent of ( Zambaldi et al. , 2019 ) is that in place of a standard Transformer block we use a 2-simplicial Transformer block . Our experiments show that the simplicial agent confers an advantage over the relational agent as an inductive bias in our reasoning task . Motivation from neuroscience for a simplicial inductive bias for abstract reasoning is contained in Appendix J . Our use of tensor products of value vectors is inspired by the semantics of linear logic in vector spaces ( Girard , 1987 ; Melliès , 2009 ; Clift & Murfet , 2017 ; Wallbridge , 2018 ) in which an algorithm with multiple inputs computes on the tensor product of those inputs , but this is an old idea in natural language processing , used in models including the second-order RNN ( Giles et al. , 1989 ; Pollack , 1991 ; Goudreau et al. , 1994 ; Giles et al. , 1991 ) , multiplicative RNN ( Sutskever et al. , 2011 ; Irsoy & Cardie , 2015 ) , Neural Tensor Network ( Socher et al. , 2013 ) and the factored 3-way Restricted Boltzmann Machine ( Ranzato et al. , 2010 ) , see Appendix A. Tensors have been used to model predicates in a number of neural network architectures aimed at logical reasoning ( Serafini & Garcez , 2016 ; Dong et al. , 2019 ) . The main novelty in our model lies in the introduction of the 2-simplicial attention , which allows these ideas to be incorporated into the Transformer architecture . 2 2-SIMPLICIAL TRANSFORMER . In this section we first review the definition of the ordinary Transformer block and then explain the 2- simplicial Transformer block . We distinguish between the Transformer architecture which contains a word embedding layer , an encoder and a decoder ( Vaswani et al. , 2017 ) , and the Transformer block which is the sub-model of the encoder that is repeated . The fundamental idea , of propagating information between nodes using weights that depend on the dot product of vectors associated to those nodes , comes ultimately from statistical mechanics via the Hopfield network ( Appendix B ) . The ordinary and 2-simplicial Transformer blocks define operators on sequences e1 , . . . , eN of entity representations . Strictly speaking the entities are indices 1 ≤ i ≤ N but we sometimes identify the entity i with its representation ei . The space of entity representations is denoted V , while the space of query , key and value vectors is denoted H . We use only the vector space structure on V , but H = Rd is an inner product space with the usual dot product pairing ( h , h′ ) 7→ h ·h′ and in defining the 2-simplicial Transformer block we will use additional algebraic structure on H , including the “ multiplication ” tensor B : H ⊗ H −→ H of ( 10 ) ( used to propagate tensor products of value vectors ) and the Clifford algebra of H ( used to define the 2-simplicial attention ) . In the first step of the standard Transformer block we generate from each entity ei a tuple of vectors via a learned linear transformation E : V −→ H⊕3 . These vectors are referred to respectively as query , key and value vectors and we write ( qi , ki , vi ) = E ( ei ) . ( 1 ) Stated differently , qi = WQei , ki = WKei , vi = WV ei for weight matrices WQ , WK , WV . In the second step we compute a refined value vector for each entity v′i = N∑ j=1 eqi·kj∑N s=1 e qi·ks vj = N∑ j=1 softmax ( qi · k1 , . . . , qi · kN ) jvj . ( 2 ) Finally , the new entity representation e′i is computed by the application of a feedforward network gθ , layer normalisation and a skip connection e′i = LayerNorm ( gθ ( v ′ i ) + ei ) . ( 3 ) Remark 2.1 . In the introduction we referred to the idea that a Transformer model learns representations of relations . To be more precise , these representations are heads , each of which determines an independent set of transformations WQ , WK , WV which extract queries , keys and values from entities . Thus a head determines not only which entities are related ( via WQ , WK ) but also what information to transmit between them ( via WV ) . In multiple-head attention with K heads , there are K channels along which to propagate information between every pair of entities , each of dimension dim ( H ) /K . More precisely , we choose a decomposition H = H1 ⊕ · · · ⊕HK so that E : V −→ K⊕ u=1 ( H⊕3u ) and write ( qi , ( 1 ) , ki , ( 1 ) , vi , ( 1 ) , . . . , qi , ( K ) , ki , ( K ) , vi , ( K ) ) = E ( ei ) . To compute the output of the attention , we take a direct sum of the value vectors propagated along every one of these K channels , as in the formula e′i = LayerNorm ( gθ [ K⊕ u=1 N∑ j=1 softmax ( qi , ( u ) · k1 , ( u ) , . . . , qi , ( u ) · kN , ( u ) ) jvj , ( u ) ] + ei ) . ( 4 ) In combinatorial topology the canonical one-dimensional object is the 1-simplex ( or edge ) j −→ i . Since the standard Transformer model learns representations of relations , we refer to this form of attention as 1-simplicial attention . The canonical two-dimensional object is the 2-simplex ( or triangle ) which we may represent diagrammatically in terms of indices i , j , k as j % % k 99 // i ( 5 ) In the 2-simplicial Transformer block , in addition to the 1-simplicial contribution , each entity ei is updated as a function of pairs of entities ej , ek using the tensor product of value vectors uj ⊗ uk and a probability distribution derived from a scalar triple product 〈pi , l1j , l2k〉 in place of the scalar product qi · kj . This means that we associate to each entity ei a four-tuple of vectors via a learned linear transformation E : V −→ H⊕4 , denoted ( pi , l 1 i , l 2 i , ui ) = E ( ei ) . ( 6 ) We still refer to pi as the query , l1i , l 2 i as the keys and ui as the value . Stated differently , pi = WP ei , l 1 i = W L1ei , l 2 i = W L2ei and ui = WUei for weight matrices WP , WL1 , WL2 , WU . Definition 2.2 . The unsigned scalar triple product of a , b , c ∈ H is 〈a , b , c〉 = ∥∥ ( a · b ) c− ( a · c ) b+ ( b · c ) a∥∥ ( 7 ) whose square is a polynomial in the pairwise dot products 〈a , b , c〉2 = ( a · b ) 2 ( c · c ) + ( b · c ) 2 ( a · a ) + ( a · c ) 2 ( b · b ) − 2 ( a · b ) ( a · c ) ( b · c ) . ( 8 ) This scalar triple product has a simple geometric interpretation in terms of the volume of the tetrahedron with vertices 0 , a , b , c. To explain , recall that the triangle spanned by two unit vectors a , b in R2 has an area A which can be written in terms of the dot product of a and b . In three dimensions , the analogous formula involves the volume V of the tetrahedron with vertices given by unit vectors a , b , c , and the scalar triple product as shown in Figure 2 . In general , given nonzero vectors a , b , c let â , b̂ , ĉ denote unit vectors in the same directions . Then we can by Lemma C.10 ( v ) factor out the length in the scalar triple product 〈a , b , c〉 = ‖a‖‖b‖‖c‖〈â , b̂ , ĉ〉 ( 9 ) so that a general scalar triple product can be understood in terms of the vector norms and configurations of three points on the 2-sphere . One standard approach to calculating volumes of such tetrahedrons is the cross product which is only defined in three dimensions . Since the space of representations H is high dimensional the natural framework for the triple scalar product 〈a , b , c〉 is instead the Clifford algebra of H ( see Appendix C ) . For present purposes , we need to know that 〈a , b , c〉 attains its minimum value ( which is zero ) when a , b , c are pairwise orthogonal , and attains its maximum value ( which is ‖a‖‖b‖‖c‖ ) if and only if { a , b , c } is linearly dependent ( Lemma C.10 ) . Using the number 〈pi , l1j , l2k〉 as a measure of the degree to which entity i is attending to ( j , k ) , or put differently , the degree to which the network predicts the existence of a 2-simplex ( i , j , k ) , the update rule for the entities when using purely 2-simplicial attention is v′i = N∑ j , k=1 e〈pi , l 1 j , l 2 k〉∑N s , t=1 e 〈pi , l1s , l2t 〉 B ( uj ⊗ uk ) ( 10 ) whereB : H⊗H −→ H is a learned linear transformation . Although we do not impose any further constraints , the motivation here is to equip H with the structure of an algebra ; in this respect we model conjunction by multiplication , an idea going back to Boole ( Boole , 1847 ) . We compute multiple-head 2-simplicial attention in the same way as in the 1-simplicial case . To combine 1-simplicial heads ( that is , ordinary Transformer heads ) and 2-simplicial heads we use separate inner product spaces H1 , H2 for each simplicial dimension , so that there are learned linear transformations E1 : V −→ ( H1 ) ⊕3 , E2 : V −→ ( H2 ) ⊕4 and the queries , keys and values are extracted from an entity ei according to ( qi , ki , vi ) = E 1 ( ei ) , ( pi , l 1 i , l 2 i , ui ) = E 2 ( ei ) . The update rule ( for a single head in each simplicial dimension ) is then : v′i = { N∑ j=1 eqi·kj∑N s=1 e qi·ks vj } ⊕ LayerNorm { N∑ j , k=1 e〈pi , l 1 j , l 2 k〉∑N s , t=1 e 〈pi , l1s , l2t 〉 B ( uj ⊗ uk ) } , ( 11 ) e′i = LayerNorm ( gθ ( v ′ i ) + ei ) . ( 12 ) If there are K1 heads of 1-simplicial attention and K2 heads of 2-simplicial attention , then ( 11 ) is modified in the obvious way using H1 = ⊕K1 u=1H 1 u and H 2 = ⊕K2 u=1H 2 u . Remark 2.3 . Without the additional layer normalisation on the output of the 2-simplicial attention we find that training is unstable . The natural explanation is that these outputs are constructed from polynomials of higher degree than the 1-simplicial attention , and thus computational paths that go through the 2-simplicial attention will be more vulnerable to exploding or vanishing gradients . The time complexity of 1-simplicial attention as a function of the number of entities isO ( N2 ) while the time complexity of 2-simplicial attention is O ( N3 ) since we have to calculate the attention for every triple ( i , j , k ) of entities . For this reason we consider only triples ( i , j , k ) where the base of the 2-simplex ( j , k ) is taken from a set of pairs predicted by the ordinary attention , which we view as the primary locus of computation . More precisely , we introduce in addition to the N entities ( now referred to as standard entities ) a set of M virtual entities eN+1 , . . . , eN+M . These virtual entities serve as a “ scratch pad ” onto which the iterated ordinary attention can write representations , and we restrict j , k to lie in the range N < j , k ≤ N + M so that only value vectors obtained from virtual entities are propagated by the 2-simplicial attention . With virtual entities the update rule is for 1 ≤ i ≤ N v′i = { N∑ j=1 eqi·kj∑N s=1 e qi·ks vj } ⊕ LayerNorm { N+M∑ j , k=N+1 e〈pi , l 1 j , l 2 k〉∑N+M s , t=1 e 〈pi , l1l , l2m〉 B ( uj ⊗ uk ) } ( 13 ) and for N < i ≤ N +M v′i = { N+M∑ j=1 eqi·kj∑N+M s=1 e qi·ks vj } ⊕ LayerNorm ( ui ) . ( 14 ) The updated representation e′i is computed from v ′ i , ei using ( 12 ) as before . Observe that the virtual entities are not used to update the standard entities during 1-simplicial attention and the 2-simplicial attention is not used to update the virtual entities ; instead the second summand in ( 14 ) involves the vector ui = WUei , which adds recurrence to the update of the virtual entities . After the attention phase the virtual entities are discarded . The method for updating the virtual entities is similar to the role of the memory nodes in the relational recurrent architecture of ( Santoro et al. , 2018 ) , the master node in ( Gilmer et al. , 2017 , §5.2 ) and memory slots in the Neural Turing Machine ( Graves et al. , 2014 ) . The update rule has complexity O ( NM2 ) and so if we take M to be of order √ N we get the desired complexity O ( N2 ) .

Summary: The paper proposes a transformer block with higher-order interactions. More precisely, instead of computing a dot product between a query vector and a key vector, 2-simplicial attention computes scalar triple product. Instead of computing a weighted average of value vectors, 2-simplicial attention computes the weighted average of tensor products of value vectors. The resulting architecture has improved representation power which is demonstrated using experiments on bridge BoxWorld environment.