Systematic Generalization with Edge Transformers
Authors: Leon Bergen
Timothy J. O'Donnell
Dzmitry Bahdanau
Venue: NeurIPS 2021 Poster (2021)
Keywords: systematic generalization
transformer
semantic parsing
language understanding

Abstract: Recent research suggests that systematic generalization in natural language understanding remains a challenge for state-of-the-art neural models such as Transformers and Graph Neural Networks. To tackle this challenge, we propose Edge Transformer, a new model that combines inspiration from Transformers and rule-based symbolic AI. The first key idea in Edge Transformers is to associate vector states with every edge, that is, with every pair of input nodes---as opposed to just every node, as it is done in the Transformer model. The second major innovation is a triangular attention mechanism that updates edge representations in a way that is inspired by unification from logic programming. We evaluate Edge Transformer on compositional generalization benchmarks in relational reasoning, semantic parsing, and dependency parsing. In all three settings, the Edge Transformer outperforms Relation-aware, Universal and classical Transformer baselines.

Systematic Generalization with Edge Transformers
Leon Bergen
Timothy J. O'Donnell
University of California, San Diego 1bergenSucad.edi
McGill University
Quebec Artificial Intelligence Institute (Mila)
Canada CIFAR AI Chair
Dzmitry Bahdanau
Element AI, ServiceNow company
McGill University
Quebec Artificial Intelligence Institute (Mila)
Canada CIFAR AI Chair
Abstract
Recent research suggests that systematic generalization in natural language un- derstanding remains a challenge for state-of-the-arr neural models such as Trans- formers and Graph Neural Networks. To tackle this challenge, we propose Edge Transformer, a new model that combines inspiration from Transformers and rule- based symbolic AI. The first key idea in Edge Transformers is to associate vector states with every edge, that is, with every pair of input nodes- as opposed to just every node, it is done in the Transformer model. The second major innovation is a triangular attention mechanism that updates edge representations in a way that is inspired by unification from logic programming. We evaluate Edge Transformer compositional generalization benchmarks relational reasoning, semantic parsing, and dependency parsing'. In all three settings, the Edge Transformer outperforms Relation-aware, Universal and classical Transformer baselines.
1 Introduction
Transformers (Vaswani 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 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 al., 2005; Scarselli al., 2009; Veličković 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 system similar to a continuous relaxation discuss in the next section, Transformers be implementing a kind
be interpreted kind of inferential subset of the Prolog logic programming language.
In particular, of reasoning over properties of entities-for example, RED(-) or TABLEE))-her)) where any relations
Cooe for our experiments can he found at: hwwnh/hhooo//eegee//dggg/annnnne
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

required for such reasoning for example PART-OF(- -) or GRANDMOTHER(, are inferred on the fly with the self- mechanism.
-raare
Building on this intuition, in this work we propose Edge Transformers, a generalization of the Transformer model that uses novel triangular attention mechanism that is inspired by 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 computed using all adjacent edges in way that is directly inspired by unification 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 chal- lenges. 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 (Veličković 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 Ques- tions 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 late-of-theestae start performance 87.4% on the COGS benchmark for compositional generalization in semantic parsing.
2 Intuitions
Transformer operates over set of n criitess-uuhh as words in a sentence- which we represent as nodes in graph like those displayed on the left- side of Figure 1. Each node is associated with a d-dimensional node-state X for i E n, which can be thought of as the output associated with applying some function to the node-that, representation of node properties. In transformers, computation proceeds by sequentially associating each node with a number I 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 (1) FA x(2),
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 the meaning of each property (-). Classical transformers how entity properties can be inferred from the properties other entities.
therefore be interpreted as architecture specialized to learning
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 (I, y) and MOTHER 2) implies GRANDMOTHER 2). Such general reasoning problem might be expressed in Prolog-like notation as follows.
x (1,2) F X' (1, (1,,)) x²,2, 1), (2,2).
(2)
of course possible for classical transformers to capture such reasoning patterns. But for them to do so, each transformer state x1+I(1) must encode both the properties of the relation itself, that is, MOTHER(.-)) versus GRANDMOTHER(... ) as well as all target nodes ø with which node stands in the relation, that is MOTHER 11, æ). In other words, to encode relations classical transformer must use its state representations for two distinct purposes. This mixing of concerns places 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 inspired by the process of unification in logic programming. Critical to inference rule such as
2

---

3
4
5
1442
1332
V152
1424
d1
d:3
043
0gg 32
152
D4
1/2
a3
V4 D5
11
2
Figure I: Left: Transformer self-attention computes update ug for the node 3 by mixing value vectors i; associated with each node with attention weights ag Right: triangular attention computes an update 012 for the edge (1,2) by mixing value vectors V122 computed for each triangle (1, 1, 2) with attention weights 0112- Values vg, V112, 012 and their contributions are not shown for simplicity.
GRANDMOTHER((A (zc y) FA MOTHER z) MOTHER(Z,.) y) fact that the two predicates on the right hand side of this rule share variable Successfully inferring GRANDMOTHER(R, y) which involve finding a binding for z which satisfies both MOTHER(2, y) and MOTHER(I, =)-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 form attention which we call triangular attention and which further describe in the next section.
Edge Transformers
edge transformer operates on a complete graph with nodes and n² directed edges (we allow self-connections) labeled with d-dimensional feature vectors, An edge transformer state therefore is 3-dimensional tensor X E Rn.a.d that consist of edges states (i,j) corresponding to every edge (i.j), E n). We will also write zij for each edge state. An edge transformer layer computes
X' = FFN +
Here, (FN is fully-connected forwaforwwrd network with one hidden layer of 4 d units that is applied independently to each edge state Zijz LN stands for the layer normalization mechanism (Ba al., 2016) which rescales activations of each feature f e [1, d] across its occurrences in edge states Tiy- Equation 3 follows the familiar structure transformer (Vaswani et al., 2017) computations with two important differences: (a) It uses 3D-tensor state instead of a matrix state; and (b) it makes use of a novel triangular attention mechanism that we describe below.
For single edge state Fiy a single-head triangular attention update outputs a vector diy that is computed as follows:
0ij W
Σ 1E
Outy
Uil
(4)
0gj softmax qakiy/ Vd, [e[L,x]
(5)
= W4
(6)
kjj = Wk
(7)
eay V za 3 V2 V2zij,
where @ stands for elementwise multiplication, and W W Wo, V1, V2 E Rd,d are matrices that are used to produce key, query, output and value vectors ka, qy aıj and vay 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 I, that is, all pairs (i. I) and (I,j). The updated edge value, a.j. is an attention weighted mixture of contributions from each such pair of edges. Figure I visualizes some of the key differences between transformer and edge transformer computations. Note that edges (i, (i, I), and (I.j) form triangle in the figure- hence the name of our attention mechanism.
We define multi-head generalization of the above mechanism in a way similar to the implementation of multi-ho head attention in vanilla transformers. Specifically, each head h [1,m] will perform Edge
3

---

(A)
(B)
zon
undle
AEX
PUNCE
XEROI
CUSI
AMHID
sister
brother
mother
Did
fetmale octor edtit produce
mother
grandfather
(C)
(D)
SELECT count(")
'xll as:illm actor Eim artfilmm performance character nsmm. 00tm00l
hell nefill wditor film nemm 0 mihlepp
Tell seflimgpdduu film nelilm uelilm.prodnette production company.filo nsem. mihtip pepeople.perooe gender nsmm.2usn
The
likes
that Emma preferred
walk
AGENT
AGENT
CCOMP
XCOMP
Did Juckie's female actor edis and produce Rad Plaid?
Figure 2: Illustrations of the benchmarks that reasoning scenarios for the CLUTRR benchmark, with the relation length (top) and (bottom). (B) Dependency parsing for Compositional Freebase Questions (CFQ; reproduced with permission, Goodwin et al., 2021b). (C) SPARQL query and the corresponding question from the CFQ semantic parsing dataset. (D) graph representations of COGS semantic parses (Ontanón al., 2021).
evaluate the Edge Transformer. (A) Relational
Attention using the smaller corresponding matrices W W V1 V2 E gd,d/m, and the joint output ai will be computed by multiplying a concatenation of the heads' output by Wo:
Oij W [atyi...... :.ti......0
In practice convenient way to implement an efficient batched version of edge attention is using the Einstein summation operation which 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 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 X is defined differently for different applications. When the input is a labeled graph with (possibly null) edge labels, the initial state zij produced by embedding the labels with trainable embedding layer. When the input is sequence, such as for example a sequence of words or tokens f,, self-edges x are initialized with embedding for token f;, and all edges 2., are initialized with relative position embedding (Shaw et al., 2018):
Fij = Se(t.) + (())
ii=
(10)
aa(i-))
if Ş
where is a trainable d-dimensional embedding layer, and a(i - j) is a relative position embedding for the position difference j.j.
To produce graph-structured output with Edge Transformer, linear layer with cross-entropy loss can be used for every edge (i. j). For classification problems with 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-sequeee 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

---

Table Hyperparameter settings for the Edge Transformer and for the baselines. L the number of layers, d' is the dimensionality, h is the number of heads, B is the batch size, is the learning rate, T training duration. For CFQ, "dep." stands for dependency parsing and "sem." " stands for semantic parsing.
Model
Task
L
d
P
Edge Transformer
CLUTRR 8
200
4
400
1 10
50 epochs
RAT
CLUTRR 8
320
8
200
1- 10
50 epochs
RRAT
CLUTRR 6
320
8
200
1 1-10
50 epochs
Edge Transformer
CFQ dep.
7
360
4
5 10 words
1 10
8000 steps
Transformer
CFQ dep.
7
360
10 words
10
8000 steps
Universal Transformer
CFQ dep.
8
360
4
5 10 words
1. 1-103
8000 steps
Edge Transformer
CFQ sem.
6
256
8
64
6 10
100 epochs
Universal Transformer
CFQ sem.
4
256
8
64
6-10-
100 epochs
Edge Transformer
COGS
3
64
4
100
5 10 200 epochs
state Tenc- 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, through a linear layer. To connect the decoder the encoder we insert Tene: 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 adapt the Transformer decoder's s causal masking approach to our triangular attention mechanism.
Experiments
In our experiments we compare the systematic generalization ability of Edge Transformers that of Transformers (Vaswani et al., 2017), Universal Transformers (Dehghani et al., 2019). Relation- aware Transformers (Shaw et al., 2018), Graph Attention Networks (Veličković et al., 2018) and other baselines. We focus on three synthetic benchmarks with carefully controlled train-test splits, Compositional Language Understanding and Text-based Relational Reasoning (CLUTRR), proposed by Sinha et al. (2019), Compositional Freebase Questions (CFQ) proposed by Keysers et al. (2020) and Compositional Generation Challenge based on Semantic Interpretation (COGS) by Kim and Linzen (2020). Systematic generalization refers to the ability to recombine known primitive units in novel combinations. These three benchmarks match test and train terms of primitive atoms used and rules of combination, but test novel combinations of atoms not seen in train. These benchmarks are thus appropriate for assessing whether Edge Transformers have greater ability to generalize systematically than existing models.
In all tables we report mean standard error over multiple runs. For details on the hyperparameter search procedure see Appendix A. For the chosen hyperparameter settings see Table 1.
4.1 Relational Reasoning on Graphs
In the first round of experiments, we use the CLUTRR benchmark proposed by Sinha et al. (2019). CLUTRR evaluates the ability of models infer unknown familial relations between individuals based on sets of given relations. For example, if one knows that is a son of B, C a son of B and C has an uncle called D. one can infer that D must also be the uncle of (see Figure A). The authors propose training models on scenarios that require a small number (2, 3 or 4) of supporting relations to infer the target relation and test on scenarios where more (up 10) supporting relations are required. We will refer to the number of inferential steps required to prove the target relation as the relation length For relation lengths i 2, 3, 4, we train models on the original CLUTRR training set from Sinha et al. (2019), which contains 15k examples. For relation lengths k = 2, 3, we generate larger training set containing 35k examples using Sinha al. (2019)'s original code, allowing us to measure systematic generalization performance with less variance. We use the noiseless (only the required facts are included) graph-based version of the challenge, where the input is represented as labeled graph, and the model's task is to label missing edge in this graph.
5

---

CLUTRR Results (Train k=2,3)
CLUTRR Results
(Train k=2,3,4)
1.0
1.0
Accuracy
0.6
Accuracy
0.6
0.2
0.2
Edge Transfo mer
Edge Treneformer
PAT
MMT
RRAT
1Siths al, 2019)
0.0
0.0
Relation Lengtn
Relation Length
Figure 3: Relation prediction accuracy on CLUTRR for test relations of lengths & E [2, 10). GAT results from Sinha et al. (2019) (only shown for the k = 2, 3, 4 task, as the = 2, training set was regenerated). Edge, Relation-Aware (RAT), Relation-updating Relation-aware (RRAT) Transformer results are averaged over 10 runs.
We create the Edge Transformer's initial state X by embedding edge labels. We apply linear layer to the final representation æ/j of the query edge (i.j) in order to produce logits for a cross-entropy loss. Our main baseline is Relation variant that can be conditioned on arbitrary graphs. Edges for RAT are initialized by embedding edge labels, and nodes are initialized with zero vectors. Logits are computed by applying a linear layer to concatenated representations of the nodes that occur in the queried edge. We also compare against variant of RAT which updates both nodes and edges. For this Relation-updating Relation-aware Transformer (RRAT), node updates are computed in the same way as in RAT. The update for the edge Zij computed by concatenating the representations for nodes i and j. and applying linear layer. Separate feed-forward layers are applied to nodes and edges. Logits are computed by applying linear layer to the queried edge. Finally, for the k 2, 3, 4 setting we include the Graph Attention Networks (GAT, Veličković et al. (2018)) results from the original CLUTRR paper for comparison. Figure displays the results. One can see the Edge Transformer beats all baselines by wide margins, in both k- = 2, 3 and k = 2, 3, 4 training settings.
Transformer (RAT) (RT; Shaw al., 2018), Transformer
4.2 Dependency Parsing of CFQ
In our second set of experiments we use the dependency parsing version of the CFQ challenge (Keysers et al., 2020) recently proposed by Goodwin et al. (2021b). The original CFQ challenge evaluates compositional generalization in semantic parsing-in the case of CFQ, the task of translating natural language questions into corresponding knowledge graph queries. The CFQ benchmark was constructed by generating test/train splits which minimize the difference between the test and train distributions over primitive units, like words, while maximizing the compound divergence the dissimilarity between test and train distributions over larger structures, like phrases. at test time, model has to generalize to novel syntactic constructions, such as "Was NP NP?" ("Was Alice the sister of Bob?") that were not seen during training. Keysers et al. (2020) released three different maximum compound divergence (MCD) splits. Goodwin et al. (2021b) propose dependency parsing version of the CFQ task and show that the MCD splits are challenging for a of-te-ttt dependency parser (see Figure 2, B). This dependency parsing task is convenient for evaluating Edge Transformer in an mucoder-only graph prediction setting that is simpler than the original sequence-to-sequence formulation of the CFQ task (see Section 4.3 for results in the original setting).
a result,
We use the Stanza framework for dependency parsing (Qi et al., 2020) in which we replace the default BiLSTM model with Transformer, Universal Transformer or Edge Transformer variants. The models are trained to label each pair of positions and with a dependency type label or null label. These predictions are then combined into a parse tree by Stanza's 's spanning- tereiing-tre-baeeed parsing algorithm. The model's initial state x is initialized from token and relational embeddings as

---

Table 2: Dependency parsing accuracy on Compositional Freebase Queries (CFQ). EM is Exact Match accuracy, LAS is Labeled Attachment Score. BiLSTM results are from Goodwin et al. (2021a). Transformer, Universal Transformer and Edge Transformer results are averaged over 10 runs.
Model
MCDI
MCD2
MCD3
EM
LAS
EM
LAS
EM
LAS
BiLSTM
96.6 + 1.3
99.6 0.2
71.4 2.6
93.3 + 0.9
56.8 = 2.8
90.9 = 1.0
Transformer
75.3 1.7
97.0
59.3
91.8 0.4
48.0
89.4 =
Universal Transformer
80.1 1.7
97.8 0.2
68.6 23
92.5 0.4
2.0
90.5 = 0.5
Edge Transformer
97.4 + 0.8
99.4 = 0.2
89.5 + 0.5
96.8 + 0.1
89.1 + 0.2
96.3 = 0.1
described by Equation 10. For Edge Transformer, pairwise output labels between positions i and are predicted directly from edge Zij. For the baselines, these labels are obtained by passing each pair of nodes through biaffine layer, following Dozat and Manning (2017). The training set contains approximately 100K questions.
We report the Labeled Attachment Score (LAS) that assesses partial correctness (F1) of the predicted dependency trees as well the exact match accuracy (EM) that only takes into account completely predicted trees. As can be seen in Table 2, Edge Transformer enjoys a sizable performance advantage (17-30% EM accuracy) over all baseline models on the MCD2 and MCD3 splits. On MCD1 split both the BiLSTM and Edge Transformer achieve near-pertee pear-perfect performance, while the classical Transformer lags behind.
4.3 Semantic Parsing of CFQ
We furthermore experiment with the original semantic parsing task from CFQ. In this task, natural language questions must be translated knowledge graph queries SPARQL query language (see Figure 2, C). We use the encoder-decode version of Edge Transformer to produce the sequence of query tokens from sequence of questions tokens.
The O(n²) computational complexity and memory footprint of Edge Transformer pose challenge in this round of experiments as the decoder's input size larger than for the other tasks (more than 100 tokens for some examples). To remedy this, we filter out training examples in which the SPARQL queries are more than 50 tokens long. Even after this filtering training an Edge Transformer model on CFQ semantic parsing requires 1-2 days of using 4 NVIDIA V100 GPUs. This makes larger scale hyperparameter search prohibitively expensive for this model. To tune Edge Transformer, we first perform the hyperparameter search for the computationally cheaper Universal Transformer baseline which we found to outperform the vanilla Transformer in preliminary experiments. For Edge Transformer we only try a few variations on the chosen Universal Transformer configuration (see Appendix A.2 for more details). As suggested in the original CFQ paper (Keysers al., 2020), we tune hyperparameters on the random split and keep the compositional generalization splits for testing. reduce the size the random split training data from 90K examples to 20K to increase the sensitivity to different hyperparameter choices. For each MCD split we report mean exact match accuracy for 30 Universal Transformer and Edge Transformer runs.
The results are shown in Table 3. Edge Transformer reaches the state-of- of-the-art performance among the non-pretrained encoder-decoder models. It reaches 24.67 1.27 % average MCD accuracy that approximately 3 percentage points ahead of comparable competition. Notably, Edge Transformer also beats Universal Transformer on our filtered random split (99.08 0.1 % vs 98.45 + 0.08 % accuracy).
4.4 Semantic Parsing of COGS
As our final evaluation of the ability of Edge Transformers to generalize systematically, we make use of the COGS benchmark of Kim and Linzen (2020). COGS defines a semantic parsing task from natural language sentences to linguistic denotations expressed a logical formalism derived from that of Reddy et al. (2017). The training/test split in COGS is designed to test the ability of models to generalize from observed sentences, sentences of greater complexity that exhibit novel combinations of primitive words and phrases. We train and test on the graph representation of the
7

---

Table Exact match accuracy the produced SPARQL queries for CFQ semantic parsing experi- ments. and & denote results from (Keysers et al., 2020) and (Furrer et al., 2020) respectively. We convert the 95% confidence intervals reported both papers to standard errors.
Model
MCD avg.
MCD1
MCD2
MCD3
Universal Transformer
21.26 1 1.06
42.7 1: 2.55
9,46 +: 1.2
11.62 1 0.68
Edge Transformer
24.69 + 1.27
47.73 + 3.63
13.14 + 2.12
13.2 1.38
Universal Transformer
18.9 +
1.8
8.1 1.3
11.3 0.2
Evolved Transformer
20.8 + 0.6
42.4
9.6
10.8 0.2
T5-small ?
21.4 + 1.2
42.5
11.2 1.2
10.6 - 0.3
Table 4: Top section: graph prediction accuracy for Edge Transformer and comparable graph prediction Transformer results by Ontanón et al. (2021) (marked by +). Bottom section: best COGS results by Akyürek et al. (2021) and Csordás et al. (2021) (+) obtained by predicting semantic parses as sequences of tokens.
Model
Generalization Accuracy
Edge Transformer
87.4 + 0.4
Universal Transformer- Attention Decoder
78.4
LSTM+Lex
82 + 1.
Transformer
81 + 0.01
COGS semantic parses provided by Ontanón et al. (2021). The nodes in these semantic graphs are input tokens that are labeled with word-level category information, for example, semantic role and part-of-speech. The edges connect some pairs of nodes indicate parent-child dependencies between them, similar to dependency parse (see Figure 2,D). The Edge Transformer model predicts node labels for token by applying several linear layers (one for each label type) to self-edges Va. For edge labels, each node has unique parent (for nodes with parent, we follow Ontanón et al. (2021) and treat them as being their own parent). To predict the parent of node i, we apply a linear layer to all edges Zji. and take softmax across the position axis. Hyperparameters for Edge Transformer were not tuned for this task. Instead, we used the best setting identified by Ontanón et al. (2021) 1) for the Transformer architecture (see Appendix A.3), and default settings for the optimizer.
As can be seen from Table 4, Edge Transformers exhibits a state-of-the--e performance of 87.4 + 0.4 % on COGS, outperforming the baseline graph prediction models from (Ontanón et al., 2021) by 9 percentage points on this challenging task. This performance is also higher than what is reported in the literature for the original sequence prediction version of the COGS task, including 82% accuracy of the lexicon-equipped L.TT (LSTM+Lex by Akyürek et al. (2021)) and 81% accuracy of the carefully initialized Transformer model (Csordás et al., 2021).
4.5 Ablation Experiments
To explore the performance of Edge Transformers, we experiment with several variations of the model. First, we train Edge Transformers with untied weights, using separate weights for each layer, as is usually done with classical Transformers. We refer to this variant as untied weights. Second, we simplify the Edge Transformer computations by computing the value Vaz for each triangle l, using only the edge (i, instead using both edges (i, and (l,j). This is equivalent to replacing Equation 8 with vay V V'zza- We will refer to this condition as value ablation. Third and finally, key vector k.- = W*zzi instead of kıj Whry for computing attention scores in Equation 5. We call this condition attention ablation.
We compare our base model to the above variants on the CLUTRR task, the MCD3 split of the CFQ dependency parsing task and on the COGS task. CLUTTR experiments we rerun the hyperparame- ters for each ablated model, whereas in other experiments we keep the same hyperparameters as for base Edge Transformer. Table shows results for CLUTRR and Table shows results for CFQ and COGS. One can see that value ablation consistently leads to worse performance CLUTRR and CFQ dependency parsing, confirming the importance of the unification-inspired triangular updates described in Equations 5-8 and illustrated in Figure 1. Likewise, attention ablation hurts the perfor-
8

---

Table Results of ablation experiments on CLUTRR. "Base" is Edge Transformer as described in Section 3. See Section 4 for details on ablated model variants. Results are averaged over 10 runs.
Model
Test k
6
8
9
10
Train on k=2,3
Base
82.0 + 2.0
74.3 - 2.4
66.9 + 2.6
61.8
2.9
58.1 - 2.8
Untied weights
78.9 + 1.0
71.0 1.1
63.2
- 1.1
58.5 - 1.2
55.3 = 1.2
Value ablation
+
4..6
36,9 + 3.6
33.9 3.7
30.7 3.6
Attention ablation
84.4 1.7
73.1 2.3
63.2
56.4
51.4
Train on k=2,3,4
Base
9
97.9
1.5
92.1 1.5
Untied weights
98.2 0.3
98.2 0.3
95.6 0.7
93.1 + 1.4
90.2 + 1.4
Value ablation
92.6 1.3
2.1
+
77.1 + 3.7
75.2 + 2.7
Attention ablation
95.9 2.1
95.2 2.5
92.8
89.8
2.9
78.7 + 2.5
Table 6: Results of ablation experiments on CFQ dependency parsing and COGS graph semantic parsing. "Base" stands for the Edge Transformer as described in Section 3. See Section 4 for details on ablated model variants. Results are averaged over 10 runs. EM Exact Match accuracy, LAS is Labeled Attachment Score.
Model
MCD3
COGS Gen. Accuracy
EM
LAS
Base
89.1 = 0.2
96.3 - 0.1
87.4 + 0.4
Untied weights
37.2 12.4
53.2 #4422
86.3 + 0.8
Value ablation
84.5 #1.1
95.5 +0.3
87.8 + 0.5
Attention ablation
88.1 +0.7
96.4 #0.1
86.3 + 0.7
mance on CLUTRR and COGS. Lastly, for untying weights, there is some evidence of deterioration for the CLUTRR setting where we train on k = 2, 3, and large effect for CFQ dependency parsing.
Related Work
Most variations of the Transformer architecture by Vaswani et al. (2017) aim to improve upon its quadratic memory and computation requirements. To this end, researchers have experimented with recurrency (Dai et al., 2019), sparse connectivity (Ainslie et al., 2020), and Bernel-based approximations to attention (Choromanski et al., 2021). Most relevant here the use relative positional embeddings in Relation-aware Transformers proposed by (Shaw et al., 2018) and refined by (Dai et al., 2019). While this approach introduces edges to Transformers, unlike in Edge Transformers, these edges statie; that is, they do not depend external input other than the respective edge labels.
Graph Neural Networks (GNNs; Gori et al., 2005; Scarselli et al., 2009; Kipf and Welling, 2017), and especially Graph Attention Networks (Veličković et al., 2018) are closely related Transformers, especially relation- aware variants. Unlike Transformers, GNNs operate on sparse graphs. While some GNN variants use edge features (Gilmer al., 2017: Chen al., 2019) or even compute dynamic edge states (Gong and Cheng, 2019), to the best of our knowledge no GNN architecture updates edge states based on states of other edges as we do Equations 5-8.
number of studies of systematic generalization have examined the ability of neural networks to handle inputs consisting of familiar atoms that are combined novel ways. Such research typically conducted using carefully controlled train-test splits (Lake and Baroni, 2018; Hupkes et al., 2019; Sinha et al., 2019; Bahdanau et al., 2019; Keysers al., 2020; Kim and Linzen, 2020). Various techniques have been used to improve generalization on some of the these benchmarks, including (but not limited to) data augmentation (Andreas, 2019; Akyürek et al., 2020), program synthesis (Nye et al., 2020), meta-learning (Lake, 2019), and task-specific model design (Guo et al., 2020). Most relevant to our work the Neural Theorem Proving (NTP) approach proposed by (Rocktäschel and Riedel, 2017) and developed by (Minervini et al., 2020). NTP can be described as differentiable version of Prolog backward-chaining inference with vector embeddings for literals and predicate names. Unlike NTP, Edge Transformers do not make rules explicit; we merely provide the networks

---

with an architecture that facilitates the learning of rule-like computations by the attention mechanism. Edge Transformers remain closer to general purpose neural models such as Transformers spirit and, as demonstrated by our CFQ and COGS experiments, can be successfully applied in the contexts where the input is not originally formulated knowledge graph.
Finally, the Inside-Outee Ossidd-Outside autoencoder proposed by Drozdov et al. (2019) for unsupervised con- stituency parsing is the most similar our work at the technical level. The Inside-Outside approach also features vector states for every pair of input nodes. However, in contrast to the Edge Transformer, the attention is restricted by the span structure of the input.
Discussion
We present the Edge Transformer, instantiation of the idea that neural model may benefit from an extended 3D-tensor state that more naturally accommodates relational reasoning. Our experimental results are encouraging, with clear performance margins in favor of the Edge Transformer over competitive baselines on three compositional generalization benchmarks. Our ablations experiments confirm the importance of core intuitions underlying the model design.
Edge Transformers can be seen as belonging the line of work on neuro-symbolic architectures (Mao et al., 2018; Rocktäschel and Riedel, 2017; Evans and Grefenstette, 2018, to name but few). Compared to most work the field, Edge Transformers are decidedly more neural: no logical gates or rule templates can be found in our architecture. We believe this is promising direction for further exploration, as allows model to retain the key advantages of neural architectures: their generality, interchangeability and ability to learn from data.
Edge Transformer computations and memory usage scale (n²) with respect to the input size While this makes scaling to larger inputs challenging, in our experience modern GPU hardware allows training Edge Transformers with inputs of size up to hundred entities. One way to use Edge Transformer in applications with larger inputs is to create Edge Transformer nodes only for the most salient or important components of an input. For example, in the language modelling context these could be just the heads of noun phrases and/or just the named entities. Future work could build and pretrain on raw text hybrid models that combine both Transformer- and Edge Transformer-bast processing, the latter performed only for a small subset of nodes.
10

---

Acknowledgments and Disclosure of Funding
We gratefully acknowledge the support of NVIDIA Corporation with the donation of two Titan V GPUs used for this research.
References
Ainslie, J., Ontanon, S., Alberti, C., Cvicek, V., Fisher, Z. Pham, P., Ravula, A., Sanghai, S., Wang, Q.. and Yang, L. (2020). ETC: Encoding Long and Structured Inputs in Transformers. arXiv:2004.08883 [cs, stat]. arXiv: 2004.08483.
Akyürek, F., Akyürek, A. F., and Andreas, J. (2020). Learning to Recombine and Resample Data For Compositional Generalization.
Akyürek, E., Akyürek, A. F., and Andreas, J. (2021). Learning to Recombine and Resample Data for Compositional Generalization. arXiv:2010.03706 [cs]. arXiv: 2010.03706.
Andreas, (2019). Good-Enouge Compositional Data Augmentation. arXiv:1904.09545 [es]. arXiv: 1904.09545.
Ba, J. L., Kiros, J. R., and Hinton, G. E. (2016). Layer Normalization. arXiv: [cs, stat]. arXiv: 1607.06450)
Bahdanau, D., Murty, S., Noukhovitch, M., Nguyen, T. H., Vries, H. d., and Courville, A. (2019).
Systematic Generalization: What Is Required and Can It Be Learned? In International Conference on Learning Representations, ICLR 2019.
Carion, N., Massa, F., Synnaeve, G., Usunier. N., Kirillov, A., and Zagoruyko, S. (2020), End-to-End Object Detection with Transformers. arXiv:2005. 12872 [cs]. arXiv: 2005.12872.
Chen, P., Liu, W., Hsieh, C.-Y., Chen, G., and Zhang, S. (2019). Utilizing Edge Features in Graph Neural Networks via Variational Information Maximization. arXiv:1906.055. 054488 [cs, stat]. arXiv: 1906.05488.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, I.., Belanger, D., Colwell, I.., and Weller, A. (2021). Rethinking Attention with Performers. arXiv:2009. 14794 [cs, stat]. arXiv: 2009. 14794.
Csordás, R., Irie, K., and Schmidhuber, J. (2021). The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers. arXiv:2108. 12284 [cs]. arXiv: 2108. 12284.
Dai, Z., Yang, Z., Yang. Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R. (2019). Transformer- XL: Attentive Language Models Beyond Fixed-Length Context. arXiv:1901.028 [cs, stat]. arXiv: 1901.02860.
Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, L. (2019). Universal Transformers. arXiv:1807.0388 03819 [cs, stat]. arXiv: 1807.03819.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). BERT: Pre-training of Deep Bidirec- tional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, NAACL 2019.
Dozat, T. and Manning, C. D. (2017). Deep Biaffine Attention for Neural Dependency Parsing. arXiv:1611.01334 [cs]. arXiv: 1.11.01734.
Drozdov, A., Verga, P., Chen, Y.-P., Tyyer, M., and McCallum, A. (2019). Unsupervised Labeled Parsing with Deep Inside-Qutside Recursive Autoencoders. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 1507-1512, Hong Kong, China. Association for Computational Linguistics.
Evans, R. and Grefenstette, E. (2018). Learning Explanatory Rules from Noisy Data. arXiv:1711.04574 [cs, math]. arXiv: 1711.04574.

---

Furrer, D., van Zee, M., Scales, N., and Schärli, N. (2020). Compositional Generalization in Semantic Parsing: Pre-training vs. Specialized Architectures. arXiv:2007.08977 [cs]. arXiv: 2007.08 2007.08970.
Gilmer, J., Schoenholz, S. S., Riley, F,, Vinyals, O., and Dahl, G. E. (2017). Neural Message
Passing for Quantum Chemistry. arXiv:/704.01212 [cs]. arXiv: 1704.01212.
Gong, L. and Cheng, Q. (2019). Exploiting Edge Features in Graph Neural Networks. arXiv:1809.02099 [cs, stat]. arXiv: 1809.02709.
Gontier, N., Sinha, K., Reddy, S., and Pal, C. (2020). Measuring Systematic Generalization in Neural Proof Generation with Transformers. arXiv:2009..0 14786 [cs, stat]. arXiv: 2009. 14786.
Goodwin, E.,, Reddy, S., O'Donnell, T. J., and Bahdanau, D. (2021a). Compositional generalization in dependency parsing. arXiv preprint arXiv:2 1rX.v:2110.06883
Goodwin, E.,, Reddy, S., O'Donnell, T. J., and Bahdanau, D. (2021b). Compositional generalization in dependency parsing. arXiv preprint arXiv:21
Gori, M., Monfardini, G., and Scarselli, F. (2005). A new model for learning in graph domains. In Proceedings. 2005 IEEE International Juint Conference on Neural Networks, 2005., volume pages 729-734- vol. 2. ISSN: 2161-4407.
Guo, Y., Lin, Z., Lou, J.-G., and Zhang, D. (2020). Hierarchical Poset Decoding for Compositional Generalization in Language. arXiv:2010.07792 [cs]. arXiv: 2010.07792.
Hupkes, D., Dankers, V., Mul, M., and Bruni, E. (2019). The compositionality of neural networks:
integrating symbolism and connectionism. arXiv: 0000000835111 [es, stat]. arXiv: 1908.08351.
Keysers, D., Schärli, N., Scales, N., Buisman, H., Furrer, D., Kashubin, S., Momchev, N., Sinopal- nikov, D., Stafiniak, I.., Tihon, T., Tsarkov, D., Wang, X., van Zee, M., and Bousquet, O. (2020). Measuring Compositional Generalization: A Comprehensive Method on Realistic Data. Inter- national Conference on Learning Representations. arXiv: 1912.09713.
Kim, N. and Linzen, T. (2020). COGS: A Compositional Generalization Challenge Based on Semantic Interpretation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 9087-9105, Online. Association for Computational Linguistics.
Kipf, T. N. and Welling, M. (2017). Suni-Supervised Classification with Graph Convolutional Networks. arXiv:1609.02907 [cs, stat]. arXiv: 1609.02907
Lake, B. M. (2019). Compositional generalization through meta sequence-to-sequenc learning. arXiv: 1906.05381.
Lake, B. M. and Baroni, M. (2018). Generalization without systematicity: On the compositional skills of soquence-to-sequenee recurrent networks. In Proceedings of the 36th International Conference on Machine Learning. arXiv: 1711.00350.
Mao, J., Gan, C., Kohli, P., Tenenbaum, J. B., and Wu, J. (2018). The Neuro- Neuro-Symbolic Concept Learner: Interpreting Scenes, Words, and Sentences From Natural Supervision. In International Conference on Learning Representations, ICLR 2019.
Minervini, P., Riedel, S., Stenetorp, P., Grefenstette, E., and Rocktäschel, (2020). Learning
Reasoning Strategies in End-to-End Differentiable Proving. arXiv:2007.06477 [cs]. arXiv: 2007.06477.
Nye, M. I., Solar-Lezama, A., Tenenbaum, J. B., and Lake, B. M. (2020). Learning Compositional Rules via Neural Program Synthesis. arXiv:2003.05562 [cs]. arXiv: 2003.05562.
Ontanón, S., Ainslie, J.. Cvicek, V., and Fisher, Z. (2021). Making transformers solve compositional tasks. arXiv preprint arXiv:2108.04. 043.8888..
Qi, P., Zhang, Y., Zhang, Y., Bolton, J., and Manning, C. D. (2020). Stanza: A Python Natural Language Processing Toolkit for Many Human Languages. arXiv:2003.07008 [cs]. arXiv: 2003.07082.
12

---

Raffel, C., Shazeer. N.. Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y.. Li, W., and Liu, P. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Tes Text Transformer. Journal of Machine Learning Research, 21(14)):1-67.
Reddy, S., Täckström, O., Petrov, S., Steedman, M., and Lapata, M. (2017), Universal semantic parsing. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Pro- cessing (EMNLP), Copenhagen, Denmark.
Rocktäschel, T. and Riedel, S. (2017). End-to-end Differentiable Proving. arXiv preprint 17X.::775511040.
Scarselli, F., Gori, M.. Tsoi, A. C., Hagenbuchner, M., and Monfardini, G. (2009). The Graph Neural Network Model. IEEE Transactions on Neural Networks, 20(1):61-80. Conference Name: IEEE Transactions on Neural Networks.
Shaw, P., Uszkoreit, J., and Vaswani, A. (2018).
with Relative Position Representations. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Short Papers), pages 464-468, New Orleans, Louisiana. Association for Computational Linguisties.
Sinha, K., Sodhani, S., Dong. J., Pineau, J., and Hamilton, W. L. (2019). CLUTRR: A Diagnostic Benchmark for Inductive Reasoning from Text. arXiv: /908. 06177 [cs, stat]. arXiv: 1908.06177.
Sinha, K., Sodhani, S., Pineau, J., and Hamilton, W. L. (2020). GraphLog: A Benchmark for Measuring Logical Generalization in Graph Neural Networks.
Tsarkov, D., Tihon, T., Scales, N., Momchev, N., Sinopalnikov, D., and Schärli, N. (2020). *-CFQ: Analyzing the Scalability of Machine Learning on Compositional Task. 202202002200226 [es]. arXiv: 2012.08266.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention All You Need. In Advances in Neural Information Process- ing Systems 30: Annual Conference Neural Information Processing Systems 2017. arXiv: 1706.03762.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. (2018). Graph Attention Networks. arXiv:1710.1 10903 [cs, stat], arXiv: 1710.10903.
13


Reviews:
[{"review_id": "v2d9tUp6SX", "reviewer": "Reviewer_Ecbf", "summary": "The paper presents a transformer-like model that carries vectors on graph edges instead of the nodes, and updates the vectors through a new (triangular) formulation of self-attention. The Edge Transformer demonstrates better generalization than baselines on two relational tasks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "- Originality\n\nSection 5 provides a useful summary of prior work involving edge vectors in transformers and GNNs. The proposed form of triangular attention appears to be both novel and well-motivated.\n\n\n- Quality\n\nThe Edge Transformer (ET) is evaluated against two baseline models, Relation-aware Transformer (RT) and a basic Transformer, on two supervised tasks (CLUTRR and CFQ) as summarized in Table 1. The results in Figure 3 and Table 2 indicate that ET outperforms the baselines on these tasks. \n\nSection 4.1 adequately describes the CLUTRR (relational reasoning) task, and the process of applying ET to it, by embedding the edge labels to obtain the initial edge vectors, and retrieving the predicted class from the query edge output vector. But the corresponding details are missing for the RT baseline model. RT carries vectors on both nodes and edges, but updates only the node vectors. Line 162 notes that RT “can be conditioned on an arbitrary graph”, so perhaps its edge vectors are initialized in the same way as those of ET, and presumably RT’s node vectors are initialized to zero. But then what happens? Since RT never updates its edge vectors, and the query edge vector is initialized to some non-informative value, retrieving the predicted class from the unchanged query edge vector could only produce a fixed output. But Figure 3 clearly shows that RT performs better than chance. It appears that RT was somehow adapted to this task, but the paper does not explain how. Without understanding how the baseline model was actually used, we cannot evaluate the significance of the difference in performance between ET and RT on the CLUTRR task. In addition, since RT was not designed to update its edge vectors, a more appropriate baseline model for this task would probably be some GNN that was designed to carry edge vectors and update them in each round of message passing.\n\nSection 4.2 describes the CFQ (dependency parsing) task, where the input is a sequence of word tokens, and the output is a label for each token pair. As with the CLUTRR task, no detailed description is provided for how the pairwise output labels are obtained from the baseline transformer, which carries no edge vectors. The results in Table 2 indicate that ET’s edge-based triangular attention delivers improved performance over standard transformer self-attention in this task.\n\nThe paper’s contribution would be strengthened by positive results on a third task. In particular, it seems important to know how well ET performs as a simple language model. It should be relatively straightforward to train ET on Wikitext and generate a training curve showing perplexity improvement over epochs. Alternatively, ET could be evaluated against GNNs on a task that requires reasoning over edge features.\n\n\n- Clarity\n\nAs noted above, the paper is missing essential details about how the baseline models were applied to the tasks. Apart from this shortcoming, the paper is quite clear and well-written. Section 3 in particular is excellent. In section 4.3, the term ‘ablation’ would be more readily understood than ‘lesion’. \n\n\n- Significance\n\nGiven the successes of transformers in NLP and beyond, any fundamental improvement to the transformer architecture can lead to major benefits. \n\n\n- Conclusion\n\nI view this as a potentially important work that currently lacks sufficient experimental support. I do not believe that novel advances such as triangular attention should be required to achieve SOTA results on major problems prior to publication, but I do believe that carefully described baselines are essential to any meaningful evaluation.\n \n- POST-DISCUSSION UPDATE\n\nTo their credit, the authors have remained deeply engaged throughout the review process, and have greatly improved the work. I still have reservations about the CLUTRR task baseline, but this seems less important than the positive results on the newly added semantic parsing task.\n\nIf the original experimental section had not left out so many key details, I would have more confidence that the next, significantly revised version of the experimental section would cover all that is required. But I’m raising my rating by 2 points, and petitioning the authors to seek feedback from someone unfamiliar with the work to help bring the experimental section up to the same standard as Section 3. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "ZzzMDOEmbQC", "reviewer": "Reviewer_HbnY", "summary": "Paper presents edge transformer, a novel transformer architecture that takes edge information, i.e. binary relation among tokens/entities as input. To further reduce the computational overhead, the authors observed that there is no need to attend on unrelated edges and therefore introduce triangular attention to replace vanilla attention. Experiments on two datasets (CLUTRR and CFQ) demonstrates the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "[Strengths]\n\n[+] The paper is overall clear and well-written. The research problem here falls in the interest of the NeurIPS community and may have the potential impact on larger groups than NLP.\n\n[+] The proposed improvement on transformers is well-motivated. I like the example in canonical deductive reasoning and I agree with this intuition. The method is technically sound and I cannot find any issues.\n\n[+] The experiments can be limited but the superior results over other baselines are impressive. \n\n[Weaknesses]\n\nHaving said those above, I do have some concerns given its current presentations and I feel this paper can be somehow premature for publishing:\n\n[-] It is still unclear which part of this method contributes the most to the performances. The authors are suggested to conduct a comprehensive ablation study on their model. At least there are two factors to test on:\n\n1. (major) vanilla transformer with relational input embedding (10)\n2. (minor) triangular attention vs. vanilla self-attention\n\nI'm particularly interested in the results of 1. Please try to include them in the rebuttal.\n\n[-] I'm not that familiar with NLP datasets&benchmarks that focus on systematic generalization. But it seems that the two datasets in this paper do not even have a compositional generalization split and therefore the results could be less convincing given the authors' claim on improving compositionally. I will recommend including some results on the following datasets:\n\n**SCAN** https://arxiv.org/abs/1711.00350\n\n**COGS** https://arxiv.org/abs/2010.05465\n\n**SyGNS** https://arxiv.org/abs/2106.01077\n\n[-] There are more and more papers today that try to explain how transformers tackle the difficult downstream tasks by visualizing the attention matrix (ex. https://arxiv.org/abs/2104.14294). Given the illustration in sec. 2, I find it could be necessary to visualize the resulting attention over all the edges. The authors could try to answer this question in their inspection: Do the proposed transformers attend to the expected edges among input tokens? \n\n[Suggestions]\n\nI'll be willing to increase my score if the following points can be addressed in a rebuttal:\n\n- ablation study on input embedding and the proposed attention mechanism.\n- results of at least one aforementioned dataset.\n- visualization on the attention matrix and some explanations.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors addressed the limitations properly.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5"}, "extra_text": {}}, {"review_id": "Cny7J3ZoGJ8", "reviewer": "Reviewer_ozBq", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "The architecture and the proposed triangular attention appear to be a relatively natural variation of Transformers, but they are novel. They are well motivated by the problem of improving compositional generalization. The paper is very clear and has all the details needed to replicate the architecture. Overall, it is a welcome addition to the literature.\n\nSome details are lacking in the experiments section, especially with respect to the dependency parsing task (the main reference  being an unspecified paper): Is the task semantic parsing or some variation of it? What are examples of input and output? Why is this task supposed to measure compositional generalization? Expanding on these points would help.\n\nIn general, the set of experiments is fairly limited and does not appear to provide enough evidence that the Edge Transformer could be more suitable than the ordinary Transformer for solving more complex tasks (especially due to the significant computational increase). The computational disadvantage of the Edge Transformer is acknowledged in Section 6, where the authors also provide ideas on how this issue could be addressed in tasks like language modeling. Additional experiments would be very welcome, but they do not necessarily fall within the scope of this paper.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations have been addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "3299S39BqEk", "reviewer": "Reviewer_V441", "summary": "The authors propose Edge Transformers to solve  CLUTRR relational reasoning challenge and Compositional Freebase Questions benchmark. The edge transformer is based on triangular attention which is weighted sum of the representations of all the triangular nodes. The proposed method can achieve better performance than the baselines of Transformer, LSTM, and GAT.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "Overall, the proposed method is interesting. I didn't know other works focusing on weighted sum of the representations in triangular. However, this method is very close to combining GNN and Transformer. How about first use GAT to encode the graph and then run another transformer layer on all the nodes? I would like to see the performance of this baseline. I feel the authors overplay the model structure but ignore the simplest way to combine GAN and Transformer.\n\nOther comments:\n1. For dependency parsing tasks, it would be better to test the method on more benchmark datasets. Or reimplement some benchmark baselines with graph neural networks. Vanilla LSTM and Transformer are too weak baselines.\n2. The intuitions section doesn't seem novel enough. The motivation is close to graph neural networks, which can also merge the information of the nodes with relation.s", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "A little bit overclaim the novelty of edge-transformer, but overlook related works on graph neural networks.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes 'Edge Transformer' which is a Transformer architecture that assigns an embedding to the edge between two tokens. A 'triangular attention' is used to update the edge embeddings, which basically computes, for edge a-b, a weighted combination of the product of edge embeddings a-n * n-b for each intermediary node n (N^3). The authors test this method on systematic generalization datasets such as CLUTTR (link prediction) and on a subtask derived from CFQ (dependency label prediction).

--

The paper was generally found clear, although some concerns about the experimental setup need to be addressed in the revision. A reviewer found improvements over baselines for systematic generalization to be "impressive". There has been an extensive back-and-forth between reviewers and authors on the validity of the baselines for the task considered. I thank both reviewers and authors for their fruitful exchange. The authors provided new positive results for the baselines suggested by reviewers. Initial concerns rotated around the synthetic nature of the task considered: the authors improved on this point by providing new results on the full semantic parsing task of CFQ by testing on each MCD split: the Edge Transformers outperforms carefully tuned previous (non pre-trained) transformers for this task. These results must still be taken with a grain of salt, given that pre-trained models achieve better results. However, these results strengthen the paper considerably. One suggestion I can think of in this respect is whether the authors could try to start from a pre-trained T5 model, add parameters for their edge attention, and fine-tune everything on the downstream CFQ task. This could maybe give a better idea on whether this model holds promise in the general setting, e.g. gains can be obtained without retraining the all model from scratch. If the potential impact of novel attention mechanism can extend to a large set of tasks, still major concerns rotate around the computational complexity of this approach (N^3) which makes it hardly applicable in general.

Overall, the reader might be left to wonder whether this method is applicable or is beneficial when it comes to large-scale pre-training of Transformer architectures. Nevertheless, provided that the authors integrate all the precious reviewers' feedback and the new experimental results in the final version of the paper, reviewers and myself agree that this paper can make an interesting addition to the conference.
Decision: Accept (Poster)