We incorporate Tensor-Product Representations within the Transformer in order to better support the explicit representation of relation structure. Our TensorProduct Transformer (TP-Transformer) sets a new state of the art on the recentlyintroduced Mathematics Dataset containing 56 categories of free-form math wordproblems. 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 standard 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.
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.
Pros:
1. This work shows better performance than baseline transformer.
2. The clustering of the subject vectors gives some insights into model's behavior .
Cons:
1. 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.
2. 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.
3. 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.
4. 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.
While 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.
----updates after reading rebuttal----
Thanks 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.