Enhancing the Transformer with explicit relational encoding for math problem solving
Authors: Imanol Schlag
Paul Smolensky
Roland Fernandez
Nebojsa Jojic
Jürgen Schmidhuber
Jianfeng Gao
Venue:  (2020)
Keywords: Tensor Product Representation
Transformer
Mathematics Dataset
Attention

Abstract: We incorporate Tensor-Product Representations within the Transformer in order to better support the explicit representation of relation structure.
Our Tensor-Product Transformer (TP-Transformer) sets a new state of the art on the recently-introduced Mathematics Dataset containing 56 categories of free-form math word-problems.
The essential component of the model is a novel attention mechanism, called TP-Attention, which explicitly encodes the relations between each Transformer cell and the other cells from which values have been retrieved by attention. TP-Attention goes beyond linear combination of retrieved values, strengthening representation-building and resolving ambiguities introduced by multiple layers of regular attention.
The TP-Transformer's attention maps give better insights into how it is capable of solving the Mathematics Dataset's challenging problems.
Pretrained models and code will be made available after publication.

Under review as a conference paper at ICLR 2020
ENHANCING THE TRANSFORMER WITH EXPLICIT RELATIONAL ENCODING FOR MATH PROBLEM SOLVING
Anonymous authors Paper under double-blind review
ABSTRACT
We incorporate Tensor-Product Representations within the Transformer in order better support the explicit representation of relation structure. Our Tensor- Product Transformer TTP-Transfrrmer) sets a new state of the art on the recently- introduced Mathematics Dataset containing 56 categories of free-form math word- problems. The essential component of the model is novel attention mechanism, called TP-Attention, which explicitly encodes the relations between each Trans- former cell and the other cells from which values have been retrieved by attention. TP-Attention goes beyond linear combination retrieved values, strengthening and resolving ambiguities introduced by multiple layers standard attention. The TP-Transformer's 's attention maps give better insights into how it is capable of solving the Mathematics Dataset's s challenging problems. Pretrained models and code will be made available after publication.
representationn-billee
INTRODUCTION
In this paper we propose a variation of the Transformer (Vaswani et al. 2017) that is designed to allow it to better incorporate structure into its representations. We test the proposal task where structured representations are expected to be particularly helpful: math word-problem solving, where, among other things, correctly parsing expressions and compositionally evaluating them is crucial. Given as input free-form math question in the form of a character sequence like Let (g) be the second derivazive be r (7) Factor -z*8 + 6 9* + 0+ 6* .. the model must produce an answer matching the specified target character- -erracte--seuuncc - (s 3) * (3*s 2) exactly. Our proposed model is trained end-to-end and infers the correct answer for novel examples without any task-specific structural biases.
2+ 21* + 10+g. Let
We begin by viewing the Transformer as a kind of Graph Neural Network (e.g., Gori et al. 2005). For concreteness, consider the encoder component of Transformer with heads. When the h head of cell t of layer issues query and result concentrates its self-attention distribution another cell e' in layer L, we can view these two cells as joined by an edge in an information-fle graph: the information content t' in effect passes via this edge affect the state of f. The strength of this attention can be viewed as a weight on this edge, and the index h of the head can be viewed as a label. Thus, each layer of the Transformer can be viewed as a complete, directed, weighted, labeled graph. Prior NLP work has interpreted certain edges of these graphs in terms of linguistic relations (Sec. 77. and we wish enrich the relation structure of these graphs better support the explicit representation of relations within the Transformer.
Here we propose to replace each of the discrete edge labels 1...., H, with a relation vector: we cre- ate bona fide representational space for the relations being learned by the Transformer. This makes possible for the hidden representation each cell to approximate the vector embedding of sym- bolic structure built from the relations generated by that cell. This embedding is Tensor-Product Representation (TPR: Smolensky 1990) in an
TPR system (Schlag & Schhidhuberr 2018 Schmidhuber 1992) that learns "internal spotlights of attention" (Schmidhu- ber] 1993). TPRs provide a general method for embedding symbol structures vector spaces. TPRs support compositional processing by directly encoding constituent structure: the representation of structure is the sum of the representation its constituents. The representation of each constituent built compositionally from two vectors: vector that embeds the content of the constituent, the

---

Under review as a conference paper at ICLR 2020
'filler' - here, the vector returned by attention and second vector that embeds the structural role fills - here, a relation conceptually labeling an edge of the attention graph. The vector that embeds filler and the vector that embeds the role it fills are bound together by the tensor product to form the tensor that embeds the constituent that they together define The relations here, and the structures they define, learned unsupervised by the Transformer service of task; post-hoc analysis is then required to interpret those roles.
In the new model, the TP-Transformer, each head of each cell generates a key-, value- and query- vector, as in the Transformer, but additionally generates a role-vector (which we refer to in some contexts a 'relation vector'). The query is interpreted as seeking the appropriate filler for that role (or equivalently, the appropriate string-location for fulfilling that relation). Each head binds that filler to its role via the tensor product (or some contraction of it), and these filler/role bindings are summed to form the TPR of a structure with H constituents (details in Sec.
An interpretation of an actual learned relation illustrates this (see Fig. 4 in Sec. 5.2). One head of our trained model can be interpreted as partially encoding the relation second-argument-of The top- layer cell dominating an input digit seeks the operator of which the digit is in the second-argument role. That cell generates a vector r, signifying this relation, and retrieves value vector by describ- ing the operator from position t' that stands in this relation. The result of this head's s attention is then the binding of filler vy to role ry; this binding is added to the bindings resulting from the cell's other attention heads.
On the Mathematics Dataset (Sec. B) the new model sets a new state of the art for the overall accuracy (Sec.F) A), and for all the setividduul-rrbbleee of interpreting the learned roles for the arithmetic-problem module show that they include good approximation to the second-argument role of the division operator and that they distinguish between numbers the numerator and denominator roles (Sec.
module accuracies (Fig.2). Initial results
More generally, it is shown that Multi-Heal Attention layers not only capture a subspace of the attended cell but capture nearly the full information content (Sec. 6.1). An argument provided that multiple layers of standard attention suffer from the binding problem, and it is shown theoretically how the proposed TP-Attention avoids such ambiguity (Sec. 6.2 The paper closes with discussion of related work (Sec.) 7 and a conclusion (Sec. g
2 THE TP-TRANSFORMER
The TP-Transformer's encoder network, like the Transformer's encoder (Vaswani et al. 2017). can be described 2-dimensional lattice of cells (t, where t 1,... T are the sequence elements of the input and = 1,.., L are the layer indices with / = as the embedding layer. All cells share the same topology and the cells of the same layer share the same weights. More specifically, each cell consists of an initial layer normalization (LN) followed by a TP-Multi-Head Attention (TPMHA) sub-layer followed by a fully-connected feed-forward (FF) sub-layer. Each sub-layer is followed by layer normalization (LN) and by a residual connection (as in the original Transformer; Eq T.. Our cell structure follows directly from the official TensorFlow source code by Vaswani et (007 but with regular Multi- Attention replaced by our TPMHA layer.
The imput into cell,/ is the output of cell,,111 and doesn't depend on the state of any other cells of the same layer, which allows a layer's 's outputs to be computed in parallel.
h 2τ,1 + TPMHA((N(EE.)) LN(21:1)))
(I)
21,41 = LN + FFFLN(b..))
We represent the symbols of the input string as one-hot vectors 3i, .... BT E Rdd where d, is the size of the vocabulary and the respective columns of the matrix E Rd. are the embedding vectors of those symbols. We also include a positional representation ps using the same sinusoidal
1 The tensor product operation (when the role- aeeeembedding vectors are linearly independent) enables the sum of constituents representing the structure a whole to be uniquely decomposable back into individual pairs of roles and their fillers, if necessary.
2

---

Under review as a conference paper at ICLR 2020
encoding schema introduced by Vaswani et al. (2017). The input of the first-layer cell;,1 is 21,0²
E d2 +
r: W ((o)
21,0 = C. Tt
where (() Rdd b()) E R T. is position- and symbol-dependent role representation, and
@ is elementwise multiplication (a contraction of the tensor product: see Sec. 3ec.2.1)
2.1 TP-MULTI-HEAD ATTENTION
The TPMHA layer of the encoder consists of H heads that can be applied parallel. Every head
h, 1 < h < H applies separate affine transformations W/ whh..)
whh.(0) w/(00 R4
b) (.()) by p,00) bj h.(4)) b p.i0 Rdde to produce key, value, query, and relation vectors from the hidden
state Zt,,, where dr dz/H:
K = (+)
1t,1 + bj (R)
h = w/il)0)
Ze,s + b, (.())
(3)
1 W/ h,0 1z,1 + bj h.(())
w Zτ,1 + by (.(1)
The filler the attention head f, 1, h is
T
Σ 1=1
h hi
(4)
i... a weighted sum of T values of the same layer and attention head (see Fig. Here dti
(0, 1) is continuous degree of match given by the softmax of the dot product between the query vector position and the key vector at position i:
exp(g!:
Wh Jdu
h, at
(5)
Ci=1 exp(de. k/ Jay
The scale factor yd, can be motivated as variance-reducing factor under the assumption that the elements of q,, and kt, are uncorrelated variables with mean 0 and variance 1, in order to initially keep the values of the softmax in region with better gradients.
Finally, bind the filler o with our relation vector r/, followed by affine transformation w E Rd, b RRd= before it is summed up with the other heads' bindings to form the TPR of structure with H constituents: this is the output of the TPMHA layer.
TPMHA (2EHH, 21.T.) = Σ Wwll (of) @ of + b/o]]
(6)
Note that, in this binding, to control dimensionality, we use a contraction of the tensor product, pointwise multiplication O: this is the diagonal of the tensor product. For discussion, the Ap- pendix.
It is worth noting that the 7th TPMHA layer returns vector that is quadratic in the inputs 21,1 to the layer: the vectors vi that are linearly combined to form o% (Eq. 4h, and r both linear in the z (Eq. 33), and they are multiplied together to form the output of TPMHA (Eq, ( This means that, unlike regular attention, TPMHA can increase, over successive layers, the polynomial degree of its representations that the feed-forward layer following attention (Sec. 2.2p introduces its own non- non-linearity even in the standard Transformer, in the TP- Tr-Transformer the attention mechanism itself goes beyond mere linear re- re-combination of vectors from the previous layer. This provides further potential for the construction of increasingly abstract representations in higher layers.
a function of the original input to the Transformer. Although it is true

---

Under review as a conference paper at ICLR 2020
h
role
filler/vluue Üt
Tt.1
query
Attention
Qt,7
.I1
-1 11
11 11-
11 -
-I
Zt,1
21:T.1
simplified illustration our TP-Attention mechanism for one head at position t in layer L. The main difference from standard Attention is the additional role representation that is element-wise multiplied with the filler/value representation.
Figure 1:
2.2 FEED-FORWARD LAYER
The feed-forward layer of cell consists of an affine transformation followed by ReLU activation and second affine transformation:
FF(x): w/( ReLU(W x +)()) b(9)
Here, w E Rd, 610) Rd,, w(()) e R4xxd, (b(0) =Rd and is the function's 's argument. As in previous work, we set d 4d..
2.3 THE DECODER NETWORK
The decoder network is a separate network with a similar structure to the encoder that takes the hidden states of the encoder and auto-regressively generates the output sequence. In contrast to the encoder network, the cells of the decoder contain two TPMHA layers and one feed-forward layer. We designed our decoder network analogously to Vaswani al. (2017) where the first attention layer attends over the masked decoder states while the second attention layer attends over the final encoder states. During training, the decoder network receives the shifted targets (ercch)-forreig) while during inference we use the previous symbol with highest probability (greedy-decoding) The final symbol probability distribution is given by
Ú = softmax ET 21,1)
(8)
where & is the hidden state the last layer of the decoder at decoding step of the output sequence and E the shared symbol embedding of the encoder and decoder.
THE MATHEMATICS DATASET
The Mathematics Dataset (Saxton al.) 2019) large collection of math problems of various types, including algebra, arithmetic, calculus, numerical comparison, measurement, numerical fac- torization, and probability. Its main goal is investigate the capability of neural networks to reason formally. Each problem is structured as a character-level sequence-to-sequenee problem. The input sequence a free-form math question or command like What is the rst derivat ve of 13*a**2 627434 + + 11914106? from which model correctly predicts the target sequence 2 627434. Another example from different module Calculate 66. 62. 14. which has 808. 524 as its target sequence.
The dataset is structured into 56 modules which cover a broad spectrum of mathematics up to uni- versity level. It is procedurally generated and comes with 2 million pre-gene pre-generated training samples per module. The authors provide an interpolation dataset for every module, as well as a few extrap- olation datasets additional measure of algebraic generalization.

---

Under review as a conference paper at ICLR 2020
We merge the different training splits train-easy, train-medium, and train- waadn-hard from all modules into one big training dataset of 120 million unique samples. From this dataset we extract a character- vocabulary of 72 symbols, including start-of-seneence, end-of-sentence, and padding symbold
EXPERIMENTAL RESULTS
We evaluate our trained model on the concatenated interpolation and extrapolation datasets of the pre-generated files, achieving new state of the art: see Table For more detailed compar- ison, Fig. 2 shows the interpolation and extrapolation performance of every module separately. The TP- Transforme matches or out-pes out-performs the Transformer in every module but one (proba- bility_sswy p-sequence). Our model never quite converged, and was stopped prematurely after 1.7 million steps. We trained our model on one server with 4 V100 Nvidia GPUs for 25 days.
Table I: Model accuracy averaged over all modules. A sample is correct if all characters of the target sequence have been predicted correctly. The column ">95%" counts how many of the 56 modules achieve over 95% accuracy. Boldface marks the best-performing model up to 700k steps.
weights
steps
train
interpolation
extrapolation
>> -95%
>95%
Simple LSTM
18M
500k
57.00%
6
41.00%
Transformer (Saxton et a
30M
500k
76.00%
13
50.00%
500k
83.06%
75.33%
52.42%
Transformer (ours)
44.2M
700k
85.01% %
77.42%
14
52.00%
2
500k
85.41%
78.30%
16
52.22%
TP-Transformer (ours)
49.2M
700k
87.25%
80.67% %
18
52.48%
3
1.7M
91.04%
84.24%
25
55.40%
4.1 IMPLEMENTATION DETAILS
We initialize the symbol embedding matrix from No 1), (()) from (1,, and all other matrices (-() using the Xavier uniform initialization as introduced by Glorot & Bengio (2110). The model parameters are set to dg = 512, dy = 2048, d, 72, H 8, I. = 6. We were not able to train the Tr-aransformer. nor the regular Transformer, using the learning rate and gradient clipping scheme described by |Saxton et al. (2019). Instead we proceed as follows: The gradients are computed using Py TyTrch'ss 's Autograd engine and their gradient norm is clipped at 0.1. The optimizer use is also Adam, but with smaller learning rate 1 10-4 betal = 0.9, beta2 0.995. We train with batch size of 1024 up to 1.7 million steps.
INTERPRETING THE LEARNED STRUCTURE
We report initial results of analyzing the learned structure of the encoder network's last layer from our 700k-step TP- TP-Trnssformer.
5.1 INTERPRETING THE LEARNED ROLES
To this end, we sample 128 problems from the interpolation dataset of the arithmetic._miixe module and collect the role vectors from a randomly chosen head. We use k-means with A: = 20 to cluster the role vectors from different samples and different time steps of the final layer of the encoder. Interestingly, we find separate clusters for digits in the numerator and denominator of fractions. When there is fraction of fractions can observe that these assignments are placed such that the second fraction reverses, arguably simplifying the division of fractions into a multiplication of fractions (see Fig.
Note that Saxton et al. 2019 report vocabulary size of 95, but this figure encompasses characters that never appear in the pre-generated training and test data.
5

---

Under review as a conference paper at ICLR 2020
Accuracy
Trannssfomer (17M)
TP Transffomer (700k)
Transtomer (700k)
numbers numbers__place saanco_aauue
numbers__round _nmber
arithmetic add or_sub
add _r_sub in _elculus__difeereniiaat calculus differennlatee
anthmettic add sub jub_mmlliple
algebra_iinear__________________ 1d comparison kth biggest
comparison sort
measurement
polynomials
comparison clooeet
anthmetic mul. div mw__multiple
algebra__innerr__22 _2
comparison kth biggest composed
polynomials coefficient_narl named
comparison pair
algebra sequence
comparison closest_composer
composed
polynomials add
arithmetic__neett nearest inlegerr_ool
_
companson _orrio____ottt_tteeeeeeeee sort_composed
calculus differentiate composed
diferrentiateccompppee
polynomials___exon
algebra
senear_dd_coppcee
gegbra__ineer__ddcomm
measurement_ measurement_connvesee conversionn
module
numbers
las____nner_ccooppe composed
algebra__polymoeeeee ptw.poooialroolls_scoommm
numbers ggdd
l_____ccc___p____ composed
comparison pair htlppparrrpppoeee composed
arithmelic
algebra sequence _ae__hh___e
polynomials evaluate composed
polynomials_comppe compose
polynomials evaluate
numbers is factor composed
probability level numbers factor
numbers place wal__vvlle_ooeee composed _aluu__coo_oedded arithmetic_ mul
numbers lcm composed
probability__www ________________________________ p_s_uunccee
numbers ooimp__comossdd
algebra__polynoreee roots
numbers
numbers div htt___eaindr__coppppee
composed
numbers_ Icm
anthmetic_ surd
numbers__t _s_ js_prime
numbers list prime factors, compose
numbers dey_reeeenderr
polynomials ______ili__lpeeee
numbers base consersionvsio
numbers list_|ppiitt Brcl__pmm_____eeeeee_e factors
Figure The accuracies our implementation of the Transformer (700k steps) and the TP- Transformer (700k and 1.7M steps) for every module of the Mathematics Dataset.

---

Under review as a conference paper at ICLR 2020
10
16
6
10
9
16
6
10
16
16
15
2
W
5
3
8
0
3
5
)
2
5
17
5
16
6
6
16
10
4
9
17
W
h
(
0
2
%
2
1
15
10
9
0
8
10
10
16
18
6
16
11
13
11
10
17
W
(
1
3
8
0
(
1
3 5
?
%
17
5
8
16
18
4
6
6
16
11
13
6
11
3
10
4
17
3
2
8
%
Figure 3: Samples correctly processed problems from the arithmetic...min module. '#' and '%" are the start- and end-of-sentence symbols. The colored squares indicate the k-means cluster of the role-vector assigned by one head in the final layer in that position. Blue and gold rectangles respectively highlight numerator and denominator roles. They were discovered manually. Note how their placement correctly swapped in rows 2, 3, and 4, where a number in the denominator of a denominator is treated as if in numerator. Role-cluster corresponds to the role oney-digit-o--aa mumerator-factor, and to
other such roles are also evident.
0.01 0.01 0.
0.11 D01
0.
0.
0.
0.
0.01
0.
0.01
0.
3.02 0.04 0.05
0.11
0:31
0.1
0.77
0.74
0.7 0.04
0,02
0.03
0.08
0.03
0.00
0.1 0.11 0.12 0.11
0,02
3.09 0.06
0.11
0.4
0.41
0.03
0.11
0.08
0.16
0.24
0.01
0.
0.
0
0.
0
0.
0.02
8.
9. 9.
5000>
0 0.00 0.11 0.01 0.01
0,61 0.04 0.23
0.11
0.00
0.22
0.09 0.09 0.08 0.28
0.06
0.1
0.02
0,02 0.09 0.01
0.05
Figure 4: TP-Transformer attention maps for three examples as described in section 5.2
5.2 INTERPRETING ATTENTION MAPS
In Fig. 4 we display three separate attention weight vectors of one head of the last TP-Transformen layer of the encoder. Gold boxes are overlaid to highlight most-relevant portions. The row above attention mask indicates the symbols that take information the symbol in the bottom row. In each case, they take from //" Seen most simply in the first example, this attention can interpreted as encoding a relation second-argun -aggmen--g holding between the querying digits and the */" operator. The second and third examples show that several numerals in the denominator can participate in this relation. The third display shows how a numerator-numentt (-297) intervening between two denominator-numerals is skipped for this relation.
6 INSIGHTS AND DEFICITS OF MULTIPLE MULTI-HEAD ATTENTION LAYERS
6.1 MULTI-HEAD ATTENTION SUBSPACES CAPTURE VIRTUALLY ALL INFORMATION
was claimed by Vaswani et al. (207)) that Multi-head attention allows the model to jointly allend to information from different representation subspaces at different positions. In this section, we show that in our trained models, an individual attention head does not access merely subset of the information in the attended cell but instead captures nearly the full information content.
Let us consider toy example where the attention layer of cell,,1 only attends to celle.- In this setting, the post-attention representation simplifies and becomes
AttentionHead( 2r,1, ZI.T.) zz,1 + W( (W
2v. (12) (19)
(9)
=1,1 + o(v(Zer)))

---

Under review as a conference paper at ICLR 2020
where and the respective affine maps (see Sec. 2.1 Note that even though W(
projection into an 8 times smaller vector space, it remains to be seen whether the hidden state loses information about Z We empirically test to what extent the trained Transformer and TP- Transformer lose information. To this end, we randomly select = 100 samples and extract the hidden state of the last layer of the encoder 21,6, as well as the value representation t/ (21.6) for every head. We then train an affine model to reconstruct 2t,6 from v% (21,6), the value vector of the single head h:
2G Wh.h bn
(10)
6 - (21.6 - 21,6)²
For both trained models, the Trannosfformer and the regular Transformer, the mean squared error e averaged across all heads is only -0.017 and -0.009 respectively. This indicates that the attention mechanism incorporates not just a subspace of the states it attends to, but affine transformations of those states that preserve nearly the full information content.
6.2 THE BINDING PROBLEM OF STACKED ATTENTION LAYERS
The binding problem refers to the problem of binding features together into objects while keeping them separated from other objects. It has been studied in the context of theoretical neuroscience (von der Malsburg
1981: 1994) but also with regards to connectionist machine learning models (Hinton al.) 1984). The purpose of binding mechanism is enable the fully distributed rep- resentation of symbolic structure (like hierarchy of features) which has recently resurfaced as an
important direction for neural network research (Lake & Baromi, 2017 Bahdanau al., 2018 van
Steenkiste et al. 2019 Palangi et al. 2017 Tang et 2018).
In this section, we describe how the standard attention mechanism is suited to capture complex nested representations, and we provide an intuitive understanding of the benefit of our TP-Attention. We understand the attention layer of cell as the means by which the subject (the cell state) queries all other cells for an object. We then show how a hierarchical representation of multiple queries becomes ambiguous in multiple standard attention layers.
Consider the string (a/b) / (e/d). A good neural representation captures the hierarchical struc- ture of the string such that it will not be confused with the similar-looking but structurally different string (a/d) (c/b). Our TP-Attention makes use of a binding mechanism in order to explicitly support complex structural relations by binding together the object representations receiving high attention with subject-specific role representation. Let continue with more technical example. Consider a simplified Transformer network where every cell consists only of a single-head attention layer with residual connection: no bias terms in the maps and en introduced in the previous section (Eq, 9). In this setting, as- sume that cella,t only attends to celly, ,, and celle, only attends cella,t where a, b, c, are distinct positions of the input sequence. In this case
feed-forward layer layer normalization, and let assume
Za,t+1 = Za,t + o (1(((()))))
(11)
2e,411 Ze,t o((1i(za)))
Suppose now that, for hierarchical grouping, the next layer celle,111 attends to both cella,141 and cello,141 (equally, each with attention weight 4). This results in the representation
Ze,12 ze,1+1 + 0 1041 (1 + zc,141))/2
(12)
= Ze,1+1 + 0 (141(1141 + 2 + (v(((())))) + (v(())))))
Note that the final representation is ambiguous in the sense that it is unclear by looking only at Eq. 12 whether cella.: has picked cells, cella, Either scenario would have led to the same outcome, which means that the network would not be able to distinguish between these two different structures (as in confusing (a/b) / (c/d) with (a/d) / (c/b)). In order to resolve this ambiguity, the standard Transformer must recruit other attention heads or find suitable non-lineat maps in between attention layers, but it remains uncertain how the network might achieve clean separation.
8

---

Under review as a conference paper at ICLR 2020
Our -P-tttention mechanism, on the other hand, specifically removes this ambiguity. Now Eqs. and become:
2a,4+1 = Za,t + or (vr((z).) @ ra.1)
2,,41 Zc, + (v()) o ve,))
(13)
Ee,42 = Ze,41 + 01+1 (1 (Za,1 + Ze,1 + (r(1)(z))) o ra.) + (z(((ZZ,)) o or.1))//2
Note that the final representation is not ambiguous anymore. Binding the filler symbols (z) (our objects) with subject-specifie role representation r described in Eq. 6 breaks the structural symmetry we had with regular attention. It is now simple for the network to specifically distinguish the two different structures.
RELATED WORK
Several recent studies have shown that the Transformer- based language model BERT Devlin al. 2018) captures linguistic relations such as those expressed in dependency- parse trees. This was shown for BERT's hidden activation states (Hewitt & Manning 2019 Tenney al. 2019 and, most directly related the present work, for the graph implicit in BERT's attention weights (Coenen et al.) 2019 Lin et al. 2019p. Future work applying the T--Transformer to language tasks (like those on which BERT is trained) will enable to study the connection between the explicit relations {ri,) the TP-Transformer learns and the implicit relations that have been extracted from BERT.
8 CONCLUSION
We have introduced the TP-Transformer, which enables the powerful Transformer architecture to learn to explicitly encode structural relations using Tensor-Product Representations. the novel and challenging Mathematics Dataset, TP-Transformer beats the previously published state of the art by 8.24%. Our initial analysis of this model's final layer suggests that the TP-Transformer naturally learns to cluster symbol representations based their structural position and relation to other symbols.

---

Under review as a conference paper at ICLR 2020
REFERENCES
Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron Courville. Systematic generalization: What is required and can be learned? arXiv preprint arXiv:/811.12800 2018.
Andy Coenen, Emily Reif, Ann Yuan, Been Kim, Adam Pearce, Fernanda Viégas, and Martin Wat- tenberg. Visualizing and measuring the geometry of BERT. arXiv preprint 1906v:996.022 2019.
Jacob Devlin, Ming- Weig-Weiiee Chang, Kenton Lee, and Kristina Toutanova. BERT:
Pre-training of deep
bidirectional transformers for language understanding. arXiv preprint arXiv: 1810.04805, 2018.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256, 2010.
Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new model for learning in graph domains. In Proceedings of the IEEE International Joint Conference on Neural Networks, volume 2, pp. 729-734, IEEE, 2005.
John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representa- tions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pa- pers), pp. 4129-4138, 2019.
Geoffrey E Hinton, James L. McClelland, David E Rumelhart, et al. Distributed representations. Carnegie-Mellon University Pittsburgh, PA, 1984.
Brenden M Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks. arXiv preprint arXiv:1711.003300 2017.
Yongjie Lin, Yi Chern Tan, and Robert Frank. Open sesame: Getting inside BERT's linguistic
knowledge. arXiv preprint arXiv:1906.01658 2019.
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. In Proceedings of NAACL-HIT 2019: Demonstrations, 2019.
Hamid Palangi, Paul Smolensky, Xiaodong He, and Li Deng. Deep learning of grammatically- interpretable representations through question-ansee answering. arXiv preprint arXiv:1705.0844 2017.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical rea- soning abilities of neural models. In International Conference on Learning Representutions, 2019. URL httttt::::/ppeeeeeeeeeee itwpeerrweeeeeeeeeeeooooeeiiiiaa forummidd-11
net forum?
Imanol Schlag and Jürgen Schmidhuber. Learning to reason with third order tensor products. In
Advances in Neural Information Processing Systems (NeurIPS), pp. 9981-9993, 2018.
J. Schmidhuber. On decreasing the ratio between learning complexity and number of time-varying variables in fully recurrent nets. In Proceedings of the International Conference on Artificial Neural Networks, Amsterdam, pp. 463. Springer, 1993.
Jürgen Schmidhuber. Learning to control west- memories: An alternative to dynamic recurrent networks. Neural Computation, 1(1):131-139 1992.
P. Smolensky. Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artif. Intell,, 46(1-2): 159-216911116 November 1990. ISSN 0004-3702, doi: 10. 101660004430000000990000004 URL. http: http:p///ooo.....oor/00011111//0 90C07-M
(90)
Shuai Tang, Paul Smolensky, and Virginia R de Sa. Learning distributed representations of symbolic structure using binding and unbinding operations. arXiv preprint /8/0. 11880./2..0 /2456, 2018.
10

---

Under review as a conference paper at ICLR 2020
Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. arXiv preprint arXiv:1905.00 05950, 2019.
Sjoerd van Steenkiste, Klaus Greff, and Jürgen Schmidhuber. A perspective on objects and system- atic generalization in model-based rl. arXiv preprint arXiv:1906. (00,:0099000 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998 1998-6008, 2017.
Christoph von der Malsburg. The correlation theory of brain function (internal report 81-2). Goet- fingen: Department of Neurobiology, Max Planck Intitute for Biophysical Chemistry, 1981.
Christoph der Malsburg. The correlation theory of brain function. In E. Domany, L. van Hemmen, and K. Schulten (eds.), Models of neural networks II, pp. 95-119. Springer, Berlin, 1994.

---

Under review as a conference paper at ICLR 2020
A APPENDIX: RELATIONS BETWEEN HADAMARD- AND DUNSOR-RRODUCT-IIDIIN
A.1 GENERAL CONSIDERATIONS
In the version of the Tr-araofformrr studied in this paper, binding of relations their values is not done by the tensor product, v @ r, as in full TPRs. Rather, a contraction of the full TPR is used: the diagonal, which is the elementwise or Hadamard product v @ r To what extent does Hadamard-produce binding share relevant properties with tensor-product binding?
A crucial property of the tensor product for its use in vector representations of structure is that structure like a/b is not confusable with a, unlike the frequently-used bag-of-words encoding: in the BOW encoding of a/b, the pair of arguments to the operator are encoded simply as a a- where and b are respectively the vector encodings of and b. Obviously, this cannot be distinguished from the BOW encoding of the argument pair in b/a. //a, b + a. (Hence the name, symbol "bag". as opposed to symbol "structure".)
In tensor-produee representation of the argument pair in a/b, we have a*r,, + baT where r,, and ra are respectively distinct vector embeddings of the numerator (or first-argument) and denominator (or second-argument) roles, and * the tensor product. This is distinct from a * ra + b Tux the embedding of the argument-pair in b/a. aa. (In Sec. 6.2 of the paper. aspect of this general property, in the context of attention models, is discussed. In Sec. [ visualization of the roles and the per- role-attention show that this particular distinction, between the numerator and denominator roles, is learned and used by the trained TP-Transformer model.)
This crucial property of the tensor product, that T + b + a b Tn, is shared by the Hadamard product: if we now take x to represent the Hadamard product, the inequality remains true. To achieve this important property, the full tensor product is not required: the Hadamard product is the diagonal of the tensor product, which retains much of the product structure of the tensor product. In any application, it is an empirical question how much of the full tensor product is required to successfully encode distinctions between bindings of symbols to roles; in the TP-Transformer, it turns out that the diagonal of the tensor product is sufficient to get improvement in performance over having no cymbol-roleproodcct structure at all. Unfortunately, the compute requirements of training on the Mathematics Dataset currently makes using the full tensor product infeasible, unless the vector representations of symbols and roles are reduced to dimensions that proved to be too small for the task. When future compute makes it possible, we expect that expanding from the diagonal to the full tensor product will provide further improvement in performance and interpretabbility.
We next move beyond these general considerations and consider a setting in which Hadamard- product attention can yield an optimal approximation to tensor- tensor-product attention.
A.2 HADAMARD-PRODUCT ATTENTION AS AN OPTIMAL APPROXIMATION TO PRODUCT ATTENTION
In Eq. 6 for TPMHA, we have a sum over all / heads of an afline-transformed product of a value vector 2/ and a role vector The (Throughout this discussion, we leave the subscripts t, implicit, as well as the over-bar on 2n in Eq. [6) In hypothetical, full-TPR formulation, this product would be the tensor product tn. @ The although in our actual proposed TP-Transformer, the Hadamard (elementwise) product di @ Th (the diagonal of e7, @ rT)) is used. The appropriateness of the compression from tensor product to Hadamard product can be seen as follows.
In the hypothetical full-TPR version of TPMHA, attention would return the sum of H tensor prod- ucts. This tensor A would have rank at most H, potentially enabling a substantial degree of com- pression across all tensors the model will compute over the data of interest. Given the translation- invariance built into the Transformer via position-invvin invariant parameters, the same compression must be applied in all positions within given layer I, although the compression may vary across heads. For the compression of A we will need more than H components, as this decomposition needs to be optimal over all all tensors in that layer for all data points.
This is vector, and should not be confused with the inner product r which is a scalar: the inner product is the sum of all the elements of the Hadamard product.
12

---

Under review as a conference paper at ICLR 2020
In detail, for each head h, the compression of the tensor A, U/h rh (or matrix A, uyr is to dimension dg, which will ultimately be mapped to dimension dg (to enable addition with via the residual connection of Eq. I by the affine transformation of 6.. The optimal dą -d-dimensioonl compression for head h at layer would preserve the dj dominant dimensions of variance of the iitention-generated states for that head and layer, across all positions and inputs: kind of singular- value decomposition retaining those dimensions with the principal singular values. Denote these principal directions by {m" & n= = I...., ] and let M2 and N₂, respectively be the d4 du matrices with the orthonormal vectors {m²) } and {nk } as columns. (Note that orthonormality implies that M M1.- I and N Nr I, with I the dş dg identity matrix.)
The compression of A,, An. will lie within the space spanned by these d2 tensor products m & i.e., A, , m. @ nei in matrix form, An = M, A,,,, where A, is the dg x de diagonal matrix with elements ah Thus the dg dimensions {a } of the compressed matrix A,, that approximates A, are given by:
A Mn AN M A. MA vnry N₁ (M₁ Un)
a An
[M unle rnlo= (ön inle
where ÜR MR th, Ph N rh- Now from Eq. UR in M (Wh,()) + bl.(t)). Thus by changing the parameters Wh,(t), bhh(v) to Wh.((.) M, Wh,(o), }h.(o) = M, ba.(),) and analogously for the role parameters Wh,(r), Bh,(r), we con- vert original hypothetical TPR attention tensor A₁ to its optimal de-dimensional approximation, in which the tensor product of the original vectors th Th is replaced by the Hadamard product of the transrrytrrnssformedd vectors üh. Th. Therefore, the proposed model, which deploys the Hadamard product, learning simply needs to converge to the parameters Wh.(v), (h.(v) rather than the parameters Wh.(v), Bh.(o).
Wholo) + b).(e), so
A.3 NEURAL MACHINE TRANSLATION
To demonstrate the generalizability of contribution, lation experiment. We follow the experimental setup of the fairseq toolkit (Ot et al, 2019). The dataset modeled in previous work (Vaswani al. 2017). The dataset consists of the WMT' English to German data but adds additional news-commentary--e 12 data from WMT' 17. We use byte-pair encoding and employed label-smoothing with value of 0.1. We train the regular trans- former and our TP-Transformer on 4 P100 GPUs for 500,000 steps using the hyperparameters of the regular transformer. Throughout training, we observe that our Transfooome achieves lower negative log- likeliiod than the regular transformer.
preliminary neural machine trans-
Table 2: The results of our neural machine translation experiment. The loss is the smoothed nega- tive log-likelln likelihhodd of the word-pieces (lower better). The TP-Transformer achieves lower loss throughout training and comparable final BLEU score (higher is better).
loss 4.137
BLEU
Transformer (ours)
27.05
TP- Transffffrrer (ours)
4.111
27.10
13

---

Under review as a conference paper at ICLR 2020
Transformer Tr-rransformer
4.7
4.5
best-loss
4.4
4.2
4.1
100000
200000
300000
400000
500000
600000
steps
The best-so-far smoothed negative log- likelihoodd of word-pieces held-out data through- out training.
Figure 5:
14


Reviews:
[{"review_id": "rkx9a33nKS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper illustrates the TP-Transformer architecture on the challenging mathematics dataset. The TP-Transformer combines the transformer architecture with tensor-product representations. The experiments show a dramatic improvement of accuracies compared with SOTA models. Moreover, the paper also explains the reason why the TP-Transformer can learn the structural position and relation to other symbols with a detailed math proof.\n \nOverall, this paper is nice as it makes a milestone for math problem solving from unique perspectives. To be specific, the paper makes the following contributions:\n\n1. Demonstrate a novel architecture TP-Transformer in details;\n2. Achieve a better accuracies in the challenging mathematics dataset than the SOTA transformer models;\n3. Illustrate in fundamental math that why TP-Transformer can learn the structural position and relation, and solve the binding problems of stacked attention layers.\n \nHere are a few minor questions that may further improve the paper:\n\n1. The conclusion states that TP-Transformer beats the previously published SOTA by 8.24%. However, it does not match to the experiment results (see section 4).\n\n2. In figure 5, there are 4 tasks in the bottom with accuracies lower than 0.5. It would be nice to provide more insights on this.\n \n3. It would be interesting to see whether it transferable to the other downstream tasks (such as natural language understanding) besides the experiments on the challenging mathematics dataset.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "Hyx0sLGk5S", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Motivated by the fact that the attention mechanism in transformers is symmetric which might not be able to disambiguate different orders, this work proposes to use a subject vector (in addition to query, key states) for each attention head, and multiply it elementwise with the context vector for each head before merging the heads. Experiments on a mathematics dataset shows superior performance compared to the normal transformer. Qualitatively, the proposed model exhibits attentions that are more interpretable, and clustering by the subject vector gives some insights into how the model solved this problem.\n\nPros:\n1. This work shows better performance than baseline transformer.\n2. The clustering of the subject vectors gives some insights into model's behavior .\n\nCons:\n1. In terms of experiments, the proposed approach adds a few million parameters to normal transformer (table 1), but in terms of interpolation it only improves 3% (extrapolation improves 0.5%) at 700k steps. The comparison would be fairer if the normal transformer can be given more parameters.\n2. In terms of experiments, this approach is only evaluated on the mathematics dataset, but the argument for relational encoding is pretty general. It would be nice if experiments on other tasks are shown in addition to the math dataset.\n3. In terms of motivation, the claim that there're ambiguities introduced by multiple layers of  regular attention needs to be supported by evidence. I think (which authors also pointed out) the feedforward network and non-linearties can disambiguate as well.\n4. In terms of interpretablity, there's claim that the learned attention maps more interpretable than transformer. Can there be more quantitative measures? It appears to me that both are hard to interpret.\n\nWhile this work shows superior performance on the mathematics dataset, I have a few concerns about the generalizability of this proposed architectural change to other problems, as well as the fairness of comparison to baseline. Therefore, I am inclined to reject this paper.\n\n----updates after reading rebuttal----\nThanks for adding the new NMT experiment in Appendix A3. My concern is that the proposed TP-Transformer is not very effective on NMT. Therefore, I'm keeping my score.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "HJeOaw0CKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper, the authors incorporated tensor-product representations within the Transformer.  By creating an attention mechanism called TP-Attention, they explicitly encode the relations between each Transformer cell and the other cells, whose values are retrieved by attention. By introducing tensor products, the proposed algorithm can empirically perform well for noncommutative operations with multiple arguments, such as division. The authors trained models with the proposed algorithm on the Mathematics Dataset and compared the performances with two baselines (simple LSTM and the original Transformer). At last, several model snapshots are provided to help interpret several key elements of the model: the learned roles, the attention maps, the TP-transformer columns and so on. \n\nOverall, the paper is well-written. The experimental results generally support the high-level intuition behind the introduction of tensor-product representation. I would recommend accepting this paper. \n\nSome quick questions:\n\n1. It was claimed in the Conclusion section that the performance of the proposed algorithm beats the previously published state of the art by 8.24%. I guess the number comes from the 2nd and the last row of interpolation accuracy in Table 1. However, these two results are obviously trained for different numbers of iterations: The baseline algorithm was trained for 500k steps, while the proposed algorithm is trained for 1.7M steps. Is it a fair comparison? If the proposed algorithm is also trained for 500k steps, the improvement is around 2.3%. \n\n2. Why is the extrapolation accuracy results for TP-Transformer missing in Table 1? \n\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a change in the attention mechanism of Transformers yielding the so-called "Tensor-Product Transformer" (TP-Transformer). The main idea is to capture filler-role relationships by incorporating a Hadamard product of each value vector representation (after attention) with a relation vector, for every attention head at every layer. The resulting model achieves SOTA on the Mathematics Dataset. Attention maps are shown in the analysis to give insights into how TP-Transformer is capable of solving the Mathematics Dataset's challenging problems. 

While the modified attention mechanism is interesting and the analysis is insightful (and improved with the addition of an experiment in NMT after the rebuttal), the reviewers expressed some concerns in the discussion stage:

1. The comparison to baseline is not fair (not to mention the 8.24% claim in conclusion). The proposed approach adds 5 million parameters to a normal transformer (table 1, 5M is a lot!), but in terms of interpolation, it only improves 3% (extrapolation improves 0.5%) at 700k steps. The rebuttal claimed that it is fair as long as the hidden size is comparable, but I don't think that's a fair argument. I suspect that increasing the feedforward hidden size (d_ff) of a normal transformer to match parameters (and add #training steps to match #train steps) might change the conclusion.
2. The new experiment on WMT further convinces me that the theoretical motivation does not hold in practice. Even with the added few million more parameters, it only improved BLEU by 0.05 (we usually consider >0.5 as significant or non-random). This might be because the feedforward and non-linearity can disambiguate as well. 

I also found the name TP-Transformer a bit misleading, since what is proposed and tested here is the Hadamard product (i.e. only the diagonal part of the tensor product). 

I recommend resubmitting an improved version of this paper with  stronger empirical evidence of outperformance of regular Transformers with comparable number of parameters.
Decision: Reject