Disentangling and Integrating Relational and Sensory Information in Transformer Architectures

## \thesubsection Improved Symbolic Reasoning in Sequence-to-Sequence tasks: Mathematical Problem Solving

Next, we evaluate DAT on a set of mathematical problem-solving tasks based on the benchmark contributed by \citetsaxtonAnalyzingMathematicalReasoning2019. We use this as a proxy for “symbolic reasoning”. Mathematical problem-solving is an interesting test for neural models because it requires more than statistical pattern recognition—it requires inferring laws, axioms, and symbol manipulation rules. The benchmark consists of a suite of mathematical problem-solving datasets, with each dataset consisting of a set of question-answer pairs. The tasks range across several modules or topics including solving equations, adding polynomials, expanding polynomials, differentiating functions, predicting the next term in a sequence, etc. For example, an example of question in the ‘polynomials__expand’ task is “Expand (5*x - 3) * (2*x + 1)” with the target “10 * x ** 2 - x - 3”.
This is modeled as a sequence-to-sequence task with character-level encoding. We compare DAT against a Transformer using matching encoder-decoder architectures. We use 2-layer models with the total number of heads fixed to $8$ in both the encoder and the decoder. We compare an encoder-decoder Transformer with $\nhsa=8$ against two configurations of DAT: one with $\nhsa=\nhra=4$ for the encoder and $\nhsa=8,\nhra=0$ for the decoder (config 1) and another with $\nhsa=\nhra=4$ for both the encoder and decoder (config 2). The number of cross-attention heads is $8$ in all cases. The DAT models use position-relative symbols as their symbol assignment mechanism.

Figure: Figure \thefigure: Validation accuracy over the course of training on mathematical problem-solving tasks. DAT learns faster and reaches higher accuracy. Solid lines indicate mean over 5 trials with different random seeds, and shaded regions indicate 95% bootstrap confidence intervals.

Each model is trained for 100 epochs, and accuracy on a hold-out validation set is tracked over the course of training. For each model and task, we run 5 trials with different random seeds to compute approximate confidence intervals. We find that DAT models learn faster and reach higher accuracies compared to a standard Transformer.
\aafatalupdate math experiments – include scaling

## \thesubsection Visual Processing with Relational Inductive Biases

## \thesubsection The Benefits of Relational Inductive Biases in Vision: Image Recognition with ImageNet

In the final set of experiments, we evaluate DAT on a vision task—object classification with the ImageNet dataset \citepimagenet. This further probes DAT’ ability in different modalities as a general-purpose sequence model. This section also stress tests DAT at larger scales.

Figure: Figure \thefigure: DAT compared to a Vision Transformer on image recognition with ImageNet. DAT learns faster and achieves better performance.

Here, we again use a Vision Transformer-style architecture \citepdosovitskiyImageWorth16x162020. ImageNet’s RGB images are divided into $16\times 16$ patches, flattened, and linearly embedded into a vector. A learnable positional embedding is added to each patch embedding. We also prepend a special classification token. The sequence of patch embeddings is then fed through an Encoder and the embedding of the class token is used to generate the final classification through a fully connected layer. We compare a Vision Transformer model with $n_{h}^{sa}=16$ to an DAT model with $n_{h}^{sa}=10,n_{h}^{sa}=6$. For both, we used a model dimension $\dmodel=1024$ and $L=24$ layers. The DAT model uses position-relative symbols as the symbol assignment mechanism and symmetric relational attention.
\Creffig:vision_acc_curve depicts the training and validation accuracy over the course of training. We find that DAT learns significantly faster. Averaging over epochs, DAT has 5.0 (resp., 4.4) percentage points higher training accuracy (resp., validation accuracy) over the course of training compared to a standard Vision Transformer. At the end of training, DAT maintains a 2.9 (resp., 1.5) percentage point advantage. This suggests that relational processing is important in processing visual scenes. This matches our intuition that parsing a visual scene requires reasoning about the visual relations between different objects or parts in the scene.

## \thesubsection Improvements in Language Modeling (Tiny Stories)

In this section, we evaluate DAT on autoregressive language modeling. Transformer language models are typically built on what is sometimes called a “decoder-only” architecture. The model receives a sequence of tokens as input and is trained to causally predict the next token at each position.
We evaluate the language modeling capabilities of DAT, as compared to standard Transformers, using the “Tiny Stories” dataset of \citeteldanTinyStoriesHowSmall2023. The dataset consists of short stories and is intended as a benchmark for small language models. Again, for each configuration, we fix the total number of attention heads, and compare a Transformer with only standard self-attention heads to DAT models with a mix of self-attention and relational attention heads. We compare a Transformer with $\nhsa=8$ attention heads to two configurations of DAT, one with $\nhsa=6,\nhra=2$ and another with $\nhsa=\nhra=4$.

Figure: Figure \thefigure: 4 Layers

fig:tiny_stories_val_loss_curves depicts the validation loss throughout training for each model. We find that DAT models with dual head attention achieve lower loss for the same total number of attention heads. We also varied the number of layers, and observed that the trend persists as the number of layers increases. The effect is small but consistent. The two DAT configurations behave similarly, with perhaps a very slight advantage to $\nhsa=\nhra=4$ (the configuration with a balanced composition of head types).
In \Creffig:tiny_stories_val_loss_curves, the DAT models use symbolic attention as the symbol assignment mechanism and asymmetric relations in relational attention. We find that symbolic attention outperforms position-relative symbols on this language modeling task. In fact, with position-relative symbols, there is no discernable advantage over the Transformer. Symbolic attention may be well-suited to language due to its implementation of a learned differentiable equivalence class mapping, which can perhaps be thought of as a form of syntax. We also find that asymmetric relations in relational attention perform better than symmetric relations. This may be because the relevant relations in language modeling are asymmetric (e.g., asymmetric syntactic or grammatical relations such as noun-verb, subject-object, determiner-noun, etc.). We provide further discussion and present ablations in \Crefssec:appendix_lm.
We conclude this section by noting that modern large language models are applied to diverse and multi-modal tasks, where different inductive biases will be useful in different contexts. While the language models explored in this section are small, an interesting avenue for future research would be to investigate whether the observed performance benefits scale up to larger models.

## \thesubsection Improvements in Language Modeling

TODO – write this section and replace tiny stories

Figure: Figure \thefigure: 350M parameter scale ($\dmodel=1024$, $\nlayers=24$)

TODO — add fineweb citation

**Table \thetable: Language Modeling on Fineweb dataset.**
| \topruleModel | Param count | # Tokens | $\dmodel$ | $\nlayers$ | $\nhsa$ | $\nhra$ | $d_{r}$ | $n_{kv}^{h}$ | Perplexity $\downarrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| \midruleTransformer | 353M | 10B | 1024 | 24 | 16 | - | - | - | 16.94 |
| DAT | 343M | 10B | 1024 | 24 | 8 | 8 | 64 | 4 | 16.09 |
| \midruleTransformer | 757M | 10B | 1536 | 24 | 24 | - | - | - | 14.65 |
| DAT | 734M | 10B | 1536 | 24 | 12 | 12 | 64 | 6 | 14.31 |
| \midruleTransformer | 1.31B | 10B | 2048 | 24 | 32 | - | - | - | 13.63 |
| DAT | 1.27B | 10B | 2048 | 24 | 16 | 16 | 128 | 8 | 13.43 |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: How are you generating attention scores for tokens after “model” and “state” in Figure 5? Is this not a causal language model?
soundness: 2
strengths: The proposed mechanism is novel, yet is a natural extension of the attention mechanisms within a standard transformer. The idea of explicitly representing relations between entities is an important one that merits further investigation. Empirically, the DAT appears to be more data efficient than the standard transformer, which is especially important in the case of language modeling.
summary: This paper introduces a novel attention mechanism by which transformers are encouraged to explicitly represent information about relations between elements of a sequence or set of tokens. This relational information is encoded separately from sensory information, or information about the specific entities that enter into a relation. This new dual attention transformer architecture is evaluated on a diverse series of modalities and tasks, demonstrating enhanced performance in language modeling, a simple visual classification task, and a toy relational reasoning dataset.
weaknesses: While the work is interesting and the idea of building in inductive biases toward representing relational information is potentially useful, there are several framing and experimental issues that should be addressed. First, the claim that standard attention mechanisms only represent sensory information is empirically false. The authors themselves cite several works describing how attention in language models often captures syntactic information, which is inherently relational. Furthermore, recent work has explicitly found that finetuned ViTs represent sensory information in early layers, but represent abstract relational information in their later layers [1]. Finally, “register tokens” – tokens that encode little information about their local image patch, but instead represent global information about an image – have been discovered in large pretrained ViTs [2]. Their very existence challenges the notion that attention mechanisms are only transmitting “sensory information". At very least, the sensory information is not necessarily tied to the pixels and token embeddings that the activation vector appears to correspond to.

The proposed method has at least two separate important components: the representation of relational information, and the tying of key and query matrices. It appears very important to test a variant of the standard transformer subject to tied key and query matrices. For example, Figure 8 shows that removing this symmetry condition deteriorates the benefits of relational attention quite a bit, especially in low-data regimes. This should also be done for the experiments presented in section 4.3. Similarly, when using position-relative symbol assignment, a control that modifies the standard transformer with relative positional embeddings should also be included. This is especially important because prior work has demonstrated that these positional embeddings can aid transformers on synthetic generalization tasks [3].

In section 4.4, the authors suggest that relational attention might be especially useful in that it captures semantic, rather than syntactic, relations between words. However, this same phenomenon also happens in standard language models! A quick check with GPT2-small using the same stimulus reveals similar attention patterns for the token “model” (Layer 0, Head 0, verifiable using this colab: https://colab.research.google.com/github/neelnanda-io/TransformerLens/blob/main/demos/Main_Demo.ipynb) . This section should thus be revised to properly contextualize the utility of relational attention in contrast to existing models.

[1] Lepori, Michael A., et al. "Beyond the Doors of Perception: Vision Transformers Represent Relations Between Objects." arXiv preprint arXiv:2406.15955 (2024).
[2] Darcet, Timothée, et al. "Vision transformers need registers." arXiv preprint arXiv:2309.16588 (2023).
[3] Csordás, Róbert, Kazuki Irie, and Jürgen Schmidhuber. "The devil is in the detail: Simple tricks improve systematic generalization of transformers." arXiv preprint arXiv:2108.12284 (2021).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Are there experiments with the learned Symbolic Attention?
soundness: 3
strengths: - The paper is well written and easy to follow
- experiments are diverse in domains and supporting main claims
summary: This paper proposes a modification of multi-head attention to better capture the relations information between objects. The proposed alternative is named dual attention, which concatenates the results of self-attention and the results of a new module named relational attention. Here, relational attention is similar to self-attention, except that the value to aggregate is replaced by a weighted sum between an relation vector computed from each pair of objects and a symbol vector. Experiments on synthetic and real data demonstrate data efficiency and parameter efficiency.
weaknesses: - The idea does not seem very novel or original. There are many attempts to integrate relational information into the attention mechanism in the Graph Neural Network community, with the closest one I can find being "Learning Graph Representations Through Learning and Propagating Edge Features" (https://ieeexplore.ieee.org/document/10004977). Specifically, Eq. 2 directly gives the general form of the proposed relational attention. This previous work goes on with slightly different parametrization of f and g, i.e. concatenation instead of dot product etc, but has an overall very similar central idea. This paper should at least cites this line of work and compare against them as baselines.

---

confidence: 3
contribution: 4
initial_rating: 10
presentation: 3
questions: If you cut some of the text as suggested above, you might have more room for future work. Do you have thoughts about how this might apply to other architectures, such as graph neural nets?
soundness: 4
strengths: This is a strong paper, and I recommend accepting it. The authors take a well-motivated challenge (helping transformers work with relational reasoning) and define a natural extension to the transformer architecture that attack that challenge.

The basic idea of exchanging "relational" information seems solid, and the given architecture is seems like a relatively simple way to achieve that, essentially adding one more bilinear function to the mix. Crucially, the new architecture decouples strength of attention from strength of relation (this distinguishes it from an earlier proposal known as relational cross attention). One might worry about the added parameters, but of course the authors are careful to control for that variable. 

The experimental data is impressive, especially because the architecture seems to work on a broad set of tasks. It's interesting that even image recognition improves. (This might be a hint that the new architecture, although meant to represent relations, might have other benefits as well.) I appreciated the careful comparisons with baselines.

Overall, this seems like an important contribution to the literature.
summary: The paper describes a new architecture designed to make it easier for transformers to work with relational information. The authors provide  several experiments that indicate the architecture outperforms standard transformers.
weaknesses: I believe there are some opportunities for improving the exposition of this paper. To begin with, "sensory" doesn't seem like the right metaphor. I realize the cognitive science origin, but I also think it's worth being careful with brain metaphors. I wonder if it would be better to talk in terms of "unary" vs. "binary" attention heads, or "first-order" vs. "relational," perhaps.

The first paragraph (and maybe much of the second) of the introduction seem unnecessary, and it might be possible to cut them entirely.

I found the first few explanations of the architecture confusing, and didn't really understand the "type" of r or symbols until I got to the explicit formulas. I wonder if it's worth making this a little more precise earlier.

The theorem in 2.4 gets very little play, and I'm not sure how important it is. I'd recommend either relegating this entirely to the appendix, or spending a bit more time explaining why it matters here. (One issue is that plain-vanilla transformers are computationally very powerful already, so it's not clear what this theorem adds.)

The second paragraph of section 3.1 seemed redundant, and might be able to be cut.

Figure 5 is potentially interesting, but I wonder if the story could be illustrated better by picking one layer, and showing attention for all the "normal" heads vs. the "relation" heads. I also wonder if there is any "low-hanging fruit" for other visualizations. For example, for training when the relation matrices are not constrained to be symmetric, do they ever end up learning to be near-symmetric? That said, this is a long paper already, and the authors explicitly mention interpretability as future work, so this is certainly an optional change!

As just mentioned—and I can't fault the authors for this—this paper has a huge amount of material, mostly in the appendices. All of that is good and necessary, but it's easy to miss details when reviewing, which is why I've put a relatively low confidence score in my review.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: None
soundness: 2
strengths: 1. The proposed method is shown to out-perform vanilla transformers on a small set of synthetic or simple real-world tasks.
2. The paper is written clearly and easy to follow.
summary: The paper proposes to add a relational attention mechanism to the transformer architecture in order to improve relational reasoning performance.
weaknesses: 1. There is limited novelty. The relational attention sounds very similar to the graph attention network to me. Also the idea of having a perception module followed by relational reasoning networks have been explored for many years in visual reasoning domain (e.g. [1,2]). The Symbolic attention also sounds like Link mechanism in the Retriever[3]. 
2. The experiments are only performed on a small set of simpler tasks. I wonder how the proposed method will perform for more complex tasks.

[1] Amizadeh, Saeed, et al. "Neuro-symbolic visual reasoning: Disentangling." International Conference on Machine Learning. Pmlr, 2020.
[2] Wang, Duo, Mateja Jamnik, and Pietro Lio. "Abstract diagrammatic reasoning with multiplex graph networks." arXiv preprint arXiv:2006.11197 (2020).
[3] Yin, Dacheng, et al. "Retriever: Learning content-style representation as a token-level bipartite graph." arXiv preprint arXiv:2202.12307 (2022).

Meta-review: This paper proposes to add a relational attention mechanism to the transformer architecture in order to improve relational reasoning performance. Here, the added relational attention mechanism is similar to self-attention, except that the value to aggregate is replaced by a weighted sum between a relation vector computed from each pair of objects and a symbol vector. Experiments on synthetic and real data demonstrate data efficiency and parameter efficiency.

Strengths: The idea of explicitly representing relations between entities is an important one that merits further investigation. Experimental results are very promising, in particular it appears that the method is more data efficient than the standard transformer, which is especially important in the case of language modeling. 
Weaknesses: Novelty is overstated, and the paper makes little effort to relate their work to the literature (especially the prior work on graph networks and message passing networks). 

This is a difficult paper to decide on. While reviewer mxrQ is clearly excited about the work, their score of 10 seems disproportionately high. Being somewhat familiar with the graph neural networks literature I tend to agree with the criticisms of Reviewers r4Cd and qVFZ (scores 5 and 3) who point to a lack of novelty, but their reviews are very short and do not seem sufficiently thorough. The results are certainly promising and in my opinion  worth publishing. At the same time I am unable to ignore the  lack of scholarship in the paper and the insistence of the authors that their method has nothing to do with graph networks. This seems especially strange since they themselves introduce their method as a form of message passing. 

Overall I recommend rejecting this paper and would encourage the authors to make an effort to situate their work relative to the literature on message passing networks.

Decision: Reject