# Disentangling and Integrating Relational and Sensory Information in
Transformer Architectures

Awni Altabaa
Department of Statistics & Data Science
Yale University
<EMAIL_ADDRESS>
&John Lafferty
Department of Statistics & Data Science
Yale University
<EMAIL_ADDRESS>
###### Abstract

The Transformer architecture processes sequences by implementing a form of
neural message-passing that consists of iterative information retrieval
(attention), followed by local processing (position-wise MLP). Two types of
information are essential under this general computational paradigm: “sensory”
information about individual objects, and “relational” information describing
the relationships between objects. Standard attention naturally encodes the
former, but does not explicitly encode the latter. In this paper, we present
an extension of Transformers where multi-head attention is augmented with two
distinct types of attention heads, each routing information of a different
type. The first type is the standard attention mechanism of Transformers,
which captures object-level features, while the second type is a novel
attention mechanism we propose to explicitly capture relational information.
The two types of attention heads each possess different inductive biases,
giving the resulting architecture greater efficiency and versatility. The
promise of this approach is demonstrated empirically across a range of tasks.

##  1 Introduction

A broad goal of machine learning research is to develop unified architectures
capable of learning and reasoning over a wide range of tasks and data
modalities, such as text, audio, time-series, and images. The generality of
the input and output formats for sequence models such as Transformers [25]
makes them promising candidates for this goal. However, a tension exists
between the goal of having a general architecture and the ability to support
inductive biases that may be favorable for certain types of tasks. Recent
research has shown that the standard Transformer architecture lacks the
ability to efficiently learn and represent relational information, leading to
several proposals for alternative architectures to encode inductive biases for
relational learning [18, 17, 20, 26, 27, 13, 3, 2]. These relational
architectures, on the other hand, fail to provide the generality required to
handle more general learning tasks. In this paper, we present an extension of
the Transformer architecture that preserves the generality of the architecture
while integrating inductive biases for processing relational information
between objects.

The Transformer can be understood as an instantiation of a broader
computational paradigm implementing a neural message-passing algorithm that
consists of iterative information retrieval followed by local processing. To
process a sequence of objects $x_{1},\ldots,x_{n}$, this takes the general
form

$\begin{split}x_{i}&\leftarrow\mathrm{Aggregate}\big{(}x_{i},\left\\{m_{j\to
i}\right\\}_{j=1}^{n}\big{)}\\\
x_{i}&\leftarrow\mathrm{Process}(x_{i})\end{split}$ (1)

In the case of Transformers, the self-attention mechanism can be seen as
sending messages from object $j$ to object $i$ that are encodings of the
sender’s features, with the message from sender $j$ to receiver $i$ given by
$m_{j\to i}=\phi_{v}(x_{j})$. These messages are then aggregated according to
some selection criterion based on the receiver’s features, typically given by
the softmax attention scores. In this work, we argue that there are two
distinct types of information that need to be encoded in the messages $m_{j\to
i}$. The first we refer to as sensory information, which represents attributes
or features of individual objects. The second is relational information about
the relationship between the sender and receiver along various dimensions. The
Transformer architecture captures the transmission of sensory information, but
does not explicitly support the transmission of relational information.

In this paper, we propose a novel type of attention mechanism that explicitly
encodes learned relations between the sender and receiver. For this attention
mechanism, the message from the sender object to the receiver object is a set
of relations between them, which can be expressed as $m_{j\to
i}=r(x_{i},x_{j})$, with the relations $r(\cdot,\cdot)$ computed through inner
products of feature maps. By combining this with the standard attention
mechanism of Transformers, we obtain a model that explicitly processes both
sensory and relational information by stacking the two types of attention
heads with $m_{j\to i}=(\phi_{v}(x_{j}),r(x_{i},x_{j}))$. This Dual Attention
architecture disentangles the two types of information in the aggregation
phase of attention, while making both types of information available during
the information processing stage.

The contributions in this paper are summarized as follows:

  1. 1.

We introduce a new relational attention mechanism that disentangles sensory
information from relational information. While standard self-attention models
retrieval of sensory information, relational attention models retrieval of
relational information.

  2. 2.

We introduce a generalized Transformer architecture that integrates sensory
and relational information through Dual Attention—a form of multi-head
attention with two distinct types of attention heads. Standard self-attention
heads encode sensory information while relational attention heads encode
relational information.

  3. 3.

We carry out an extensive set of experiments, showing that the Dual Attention
Transformer outperforms standard Transformers in terms of data efficiency. Our
experiments range across several tasks and data modalities, including visual
relational reasoning, symbolic mathematical reasoning, language modeling, and
image recognition.

(a) $\mathrm{Attention}(x,(y_{1},\ldots,y_{n}))$

(b) $\mathrm{RelationalAttention}(x,(y_{1},\ldots,y_{n}))$

Figure 1: Standard self-attention retrieves sensory information $v_{i}$ about
the attributes of individual objects while relational attention retrieves
relational information $r(x,y_{i})$ about the relationship between the objects
in the context and the receiver. Each relation is tagged with a symbol $s_{i}$
which acts as an abstract variable identifying the sender. In both cases,
information is aggregated according to the attention scores $\alpha_{i}$,
which are computed by a softmax over inner products of queries and keys.

##  2 Disentangling Attention over Sensory and Relational Information

###  2.1 Standard Attention: Attention over Sensory Information

The attention mechanism of standard Transformers can be understood as a form
of neural message-passing that performs selective information retrieval. An
object emits a query that is compared against the keys of each object in its
context via an inner product. A “match” occurs when the inner product is
large, causing the features of the attended object to be retrieved and added
to the residual stream of the receiver. Formally, standard attention takes the
form

$\begin{split}\mathrm{Attention}(x,(y_{1},\ldots,y_{n}))&=\sum_{i=1}^{n}\alpha_{i}(x,\bm{y})\phi_{v}(y_{i})\\\
\alpha(x,\bm{y})&=\mathrm{Softmax}\big{(}\left[\left\langle\phi_{q}(x),\phi_{k}(y_{i})\right\rangle\right]_{i=1}^{n}\big{)},\end{split}$
(2)

where $\phi_{q},\phi_{k}$ are learnable query/key maps controlling the
selection criterion and $\phi_{v}$ is a learnable value mapping controlling
what information is sent. The attention scores $\alpha(x,\bm{y})\in\Delta^{n}$
are used to retrieve a convex combination of the values.

The information being retrieved here is “sensory” information—that is, the
features and attributes of individual objects. There is no explicit retrieval
of information about the relationship between the features of the sender and
the receiver. The attention scores $\alpha_{i}(x,\bm{y})$ can perhaps be
thought of as (normalized) relations between objects, but these are merely
computed as an intermediate step in an information-retrieval operation, and
are ultimately entangled with the object-level features of the sender111In
principle, it is also possible that the MLP compute a relation between the
sender and receiver in the local processing step by separating their
representations and computing a comparison. However, this is difficult to do
since the representations of the objects will be additively mixed, and there
is no inductive bias pressuring the computed function to be a relation.. This
makes learning relational representations in standard Transformers
inefficient.

###  2.2 Relational Attention: Attention over Relational Information

We propose an attention mechanism with a relational inductive bias. Under the
message-passing view of Equation 1, this attention mechanism represents an
operation where the message from one object to another encodes the relation
between the sender and the receiver. We call this relational attention.

At a high level, this operation begins in the same way as the standard
attention mechanism, with each object emitting a query and a key, which are
compared via an inner product. When the inner product is high, an object is
selected. But, now, rather than retrieving the features of the selected
object, what is retrieved is the relation between the two objects. In
addition, we must send an identifier that signals to the receiver “who the
sender is”. Mathematically, this operation is defined as follows.

$\begin{split}\mathrm{RelationalAttention}(x,(y_{1},\ldots,y_{n}))&=\sum_{i=1}^{n}\alpha_{i}(x,\bm{y})\big{(}r(x,y_{i})W_{r}+s_{i}W_{s}\big{)},\\\
\alpha(x,\bm{y})&=\mathrm{Softmax}\big{(}\left[\left\langle\phi_{q}^{\mathrm{attn}}(x),\phi_{k}^{\mathrm{attn}}(y_{i})\right\rangle\right]_{i=1}^{n}\big{)}\in\Delta^{n},\\\
r(x,y)&=\big{(}\left\langle\phi_{q,\ell}^{\mathrm{rel}}(x),\phi_{k,\ell}^{\mathrm{rel}}(y)\right\rangle\big{)}_{\ell\in[d_{r}]}\in\mathbb{R}^{d_{r}},\\\
(s_{1},\ldots,s_{n})&=\mathrm{SymbolRetriever}(\bm{y};\,S_{\mathrm{lib}})\end{split}$
(3)

Thus, relational attention between the object $x$ and the context
$\bm{y}=(y_{1},\ldots,y_{n})$ retrieves a convex combination of $x$’s
relations with each object in the context, $r(x,y_{i})$, each tagged with a
symbol $s_{i}$ that identifies the sender. We will expand on the role and
implementation of the symbols in the next subsection. Here,
$\phi_{q}^{\mathrm{attn}},\phi_{k}^{\mathrm{attn}}$ are learned feature maps
that control the selection criterion for which object(s) in the context to
attend to. Another set of query/key feature maps,
$\phi_{q,\ell}^{\mathrm{rel}},\phi_{k,\ell}^{\mathrm{rel}},\ell\in[d_{r}]$,
are learned to represent the relation between the sender and the receiver.
Each inner product
$\langle\phi_{q,\ell}^{\mathrm{rel}}(x),\phi_{k,\ell}^{\mathrm{rel}}(y)\rangle$
can be thought of as a comparison of the two objects’ features under a
particular filter—i.e., a ‘relation’. The $d_{r}$ pairs of feature maps
produce a $d_{r}$-dimensional relation vector.

In some tasks, a good inductive bias on the relations is symmetry: the
relation between $x$ and $y$ is the same as the relation between $y$ and $x$.
This can be achieved by using the same feature filter for the query and key
maps (i.e., $\phi_{q}^{\mathrm{rel}}=\phi_{k}^{\mathrm{rel}}$). This imbues
the relations with added structure, making them positive semi-definite kernels
which define a pseudometric on the object space and a corresponding geometry.
We explore this and expand on this discussion in our experiments.

###  2.3 Symbol Assignment Mechanisms

For the purposes of processing relations, the receiver needs to know: 1) the
relation between itself and the objects in its context, and 2) the identity of
the object corresponding to each relation. The symbols in relational attention
are used to tag each relation with the identity of the sender (the object the
relation is with). Without this information, the result of relational
attention would only be an aggregated representation of the relations between
the receiver and the selected object(s).

A symbol identifies or “points to” an object, but, importantly, it does not
fully encode the features of the object. The second point is what makes
relational attention (Equation 3) “disentangled” from sensory or object-level
features. The sensory features of individual objects are high-dimensional and
have a lot of variability. By contrast, relational information is low-
dimensional and more abstract. If sensory features are mixed with relational
information, the sensory information could overwhelm the relational
information, preventing abstraction and generalization. Instead, symbols act
as abstract references to objects and may be thought of as connectionist
analogs of pointers in traditional symbolic architectures.

Here, the “identity” of an object may mean different things in different
situations. For us, identity may be encoded by 1) position, 2) relative
position, or 3) an equivalence class over features. Correspondingly, we
consider three different symbol assignment mechanisms.

A symbol identifies or “points to” an object, but, importantly, does not fully
encode the features of the object. The second point is what makes relational
attention (Equation 3) “disentangled” from sensory or object-level features.
The sensory features of individual objects are high-dimensional and have a lot
of variability. By contrast, relational information is low-dimensional and
more abstract. If sensory features are mixed with relational information, they
would overwhelm the relational information, preventing abstraction and
generalization. Instead, symbols act as abstract references to objects,
perhaps thought of as a connectionist analog of pointers in traditional
symbolic architectures.

Here, the “identity” of an object may mean different things in different
situations. For us, identity may be encoded by 1) position, 2) relative
position, or 3) an equivalence class over features. We consider three
different corresponding symbol assignment mechanisms.

Positional Symbols. In some applications, it is sufficient to identify objects
through their position in the input sequence. We maintain a library of symbols
$S_{\mathrm{lib}}=(s_{1},\ldots,s_{\mathtt{max\\_len}})\in\mathbb{R}^{\mathtt{max\\_len}\times
d_{\mathrm{model}}}$ and assign $s_{i}$ to the $i$-th object in the sequence.
These are essentially learned positional embeddings.

Position-Relative Symbols. Sometimes, the more useful identifier is relative
position with respect to the receiver, rather than absolute position. This can
be implemented with position-relative embeddings. We learn a symbol library
$S_{\mathrm{lib}}=(s_{-\Delta},\ldots,s_{-1},s_{0},s_{1},\ldots,s_{\Delta})\in\mathbb{R}^{(2\Delta+1)\times
d_{\mathrm{model}}}$, and relational attention becomes
$x_{i}^{\prime}\leftarrow\sum_{j}\alpha_{ij}(r(x_{i},x_{j})\,W_{r}+s_{j-i}\,W_{s})$,
with “$m_{j\to i}=(r(x_{i},x_{j}),s_{j-i})$.

Symbolic Attention. In certain domains, some information about the objects’
features is necessary for identifying them for the purposes of relational
processing. Yet, to maintain a relational inductive bias, we would like to
avoid sending a full encoding of object-level features. In symbolic attention,
we learn a set of symbol vectors,
$S_{\mathrm{lib}}=(s_{1},\ldots,s_{n_{s}})\in\mathbb{R}^{n_{s}\times
d_{\mathrm{model}}}$ and a matching set of feature templates
$F_{\mathrm{lib}}=(f_{1},\ldots,f_{n_{s}})$. We retrieve a symbol for each
object by an attention operation that matches the input vectors $x_{i}$
against the feature templates $f_{j}$ and retrieves symbols $s_{j}$.

$\mathrm{SymbolicAttention}(\bm{x})=\mathrm{Softmax}\big{(}(\bm{x}\,W_{q})\,F_{\mathrm{lib}}^{\top}\big{)}S_{\mathrm{lib}}.$
(4)

Here, $S_{\mathrm{lib}},F_{\mathrm{lib}},W_{q}$ are learned parameters. This
can be thought of as implementing a learned differentiable “equivalence class
map” over feature embeddings. Crucially, the number of symbols (i.e., “feature
equivalence classes”) is finite, which enables relational attention to still
produce a relation-centric representation while tagging the relations with the
necessary identifier.

We find that different symbol assignment mechanisms are more effective in
different domains.

###  2.4 Approximation theory: What Class of Functions can Relational
Attention Compute?

To get some intuition about the type of computation that relational attention
can perform, we present the following approximation result. The following
theorem says that relational attention can approximate any function on
${\mathcal{X}}\times{\mathcal{Y}}^{n}$ which 1) selects an element in
$(y_{1},\ldots,y_{n})$, then 2) computes a relation with it. Both the
selection criterion and the relation function are arbitrary, and the selection
criterion is query-dependent. The formal statement and proof are given in
Appendix A.

######  Theorem 1 (Informal).

Let $\mathrm{Select}:{\mathcal{X}}\times{\mathcal{Y}}^{n}\to{\mathcal{Y}}$ be
an arbitrary preference selection function, which selects an element among
$(y_{1},\ldots,y_{n})$ based on a query-dependent preorder relation
$\\{\preccurlyeq_{x}\\}_{x\in{\mathcal{X}}}$. Let
$\mathrm{Rel}:{\mathcal{X}}\times{\mathcal{Y}}\to\mathbb{R}^{d_{r}}$ be an
arbitrary continuous relation function on ${\mathcal{X}}\times{\mathcal{Y}}$.
There exists a relational attention module that approximates the function
$\mathrm{Rel}(x,\mathrm{Select}(x,\bm{y}))$ to arbitrary precision.

##  3 Integrating Attention over Sensory and Relational Information

###  3.1 Dual Attention

One of the keys to the success of the Transformer architecture is the use of
so-called multi-head attention. This involves learning and computing multiple
attention operations in parallel at each layer and concatenating the output.
This enables learning multiple useful criteria for routing information between
objects. However, in standard Transformers, all these attention heads share
the same inductive bias, focusing on sensory information about individual
objects. In particular, there is no explicit support for processing relational
information between objects.

We posit that sensory and relational information are the two primary types of
information that are of relevance when processing sequences or collections of
objects. In this paper, we explore the effects of augmenting a Transformer
with a specialized attention operation with relational inductive biases. We
propose a type of multi-head attention with two distinct types of attention
heads: standard self-attention, and relational attention. Our hypothesis is
that by having both kinds of computations available to the model, it can learn
to use both and select between them depending on the current task or context.

Algorithm 1 describes the proposed module. The number of self-attention heads
$n_{h}^{sa}$ and number of relational attention heads $n_{h}^{ra}$ are
hyperparameters. The self-attention heads attend to and retrieve sensory
information while the relational attention heads attend to and retrieve
relational information. The $n_{h}=n_{h}^{sa}+n_{h}^{ra}$ heads are then
concatenated to produce the output. The result is a representation of
contextual information with integrated sensory and relational components. We
note that a symmetry inductive bias can be injected into the relations
$\bm{r}_{ij}$ by imposing that $W_{q}^{\mathrm{rel}}=W_{k}^{\mathrm{rel}}$.

Input: $\bm{x}=(x_{1},\ldots,x_{n})\in\mathbb{R}^{n\times d}$

Compute self-attention heads $\displaystyle\bm{\alpha}^{(h)}$
$\displaystyle\leftarrow\mathrm{Softmax}\big{(}{(\bm{x}\,W_{q}^{h})}{(\bm{x}\,W_{k}^{h})}^{\intercal}\big{)},\qquad$
$\displaystyle h\in[n_{h}^{sa}]$ $\displaystyle e_{i}^{(h)}$
$\displaystyle\leftarrow\sum_{j}\alpha_{ij}^{(h)}x_{j}\,W_{v}^{h},\qquad$
$\displaystyle i\in[n],h\in[n_{h}^{sa}]$ $\displaystyle e_{i}$
$\displaystyle\leftarrow\mathrm{concat}\big{(}e_{i}^{(1)},\ldots,e_{i}^{(n_{h}^{sa})}\big{)}\,W_{o}^{sa},$
$\displaystyle i\in[n]$

Assign symbols:
$\bm{s}=(s_{1},\ldots,s_{n})\leftarrow\mathrm{SymbolRetriever}(\bm{x};\,S_{\mathrm{lib}})$

Compute relational attention heads $\displaystyle\bm{\alpha}^{(h)}$
$\displaystyle\leftarrow\mathrm{Softmax}\big{(}{(\bm{x}\,W_{q,h}^{\mathrm{attn}})}{(\bm{x}\,W_{k,h}^{\mathrm{attn}})}^{\intercal}\big{)},\qquad$
$\displaystyle h\in[n_{h}^{ra}]$ $\displaystyle\bm{r}_{ij}$
$\displaystyle\leftarrow\big{(}\left\langle
x_{i}\,W_{q,\ell}^{\mathrm{rel}},x_{j}\,W_{k,\ell}^{\mathrm{rel}}\right\rangle\big{)}_{\ell\in[d_{r}]}\qquad$
$\displaystyle i,j\in[n]$ $\displaystyle a_{i}^{(h)}$
$\displaystyle\leftarrow\sum_{j}\alpha_{ij}^{(h)}\big{(}\bm{r}_{ij}\,W_{r}^{h}+s_{j}\,W_{s}^{h}\big{)},\qquad$
$\displaystyle i\in[n],\,h\in[n_{h}^{ra}]$ $\displaystyle a_{i}$
$\displaystyle\leftarrow\mathrm{concat}\big{(}a_{i}^{(1)},\ldots,a_{i}^{(n_{h}^{ra})}\big{)}W_{o}^{ra},$
$\displaystyle i\in[n]$

Output : $\big{(}\mathrm{concat}(e_{i},a_{i})\big{)}_{i=1}^{n}$

Algorithm 1 Dual Attention

Attention Masks & Causality. Any type of attention mask (e.g., causal mask for
autoregressive language modeling) can be implemented in relational attention
in the same way as for standard self-attention (i.e., mask is added to
$\alpha_{ij}^{h}$ pre-softmax).

Positional Encoding. There exists different methods in the literature on
encoding positional information in the Transformer architecture. For example,
[25] propose adding positional embeddings to the input, [21] propose adding
relative-positional embeddings at each attention operation, and [23] propose
rotary positional embeddings (RoPE) which apply a position-dependent map to
the queries and keys pre-softmax. These methods are compatible with dual
attention and are configurable options in our public implementation.

Computational complexity. The computational complexity of relational attention
scales the same as self-attention with a $O(n^{2})$ dependence on sequence
length. Like standard attention, relational attention can be computed in
parallel via efficient matrix multiplication operations.

###  3.2 The Dual Attention Transformer Architecture

A standard Transformer implements a procedure of iterative information
retrieval (attention) followed by local processing (MLP). We define our Dual
Attention Transformer in the same way, except that self-attention is replaced
by dual attention. Algorithms 2 and 3 defines an encoder and decoder block
with dual attention. Composing these blocks yields the Dual Attention
Transformer.

######  Remark 1.

In our implementation, the symbol library $S_{\mathrm{lib}}$ is shared across
layers to reduce the number of parameters. The rationale is that, since the
role of symbols is merely to act as abstract variables, they can be remapped
at each layer (e.g., a variable can be reassigned to point to a different
object).

Input : $\bm{x}\in\mathbb{R}^{n\times d}$

$\bm{x}\leftarrow\mathrm{Norm}(\bm{x}+\mathrm{DualAttn}(\bm{x}))$

$\bm{x}\leftarrow\mathrm{Norm}(\bm{x}+\mathrm{MLP}(\bm{x}))$

Output: $\bm{x}$

Algorithm 2 Dual Attention Encoder Block

Input : $\bm{x},\bm{y}\in\mathbb{R}^{n\times d}$

$\bm{x}\leftarrow\mathrm{Norm}(\bm{x}+\mathrm{DualAttn}(\bm{x}))$

$\bm{x}\leftarrow\mathrm{Norm}(\bm{x}+\mathrm{CrossAttn}(\bm{x},\bm{y}))$

$\bm{x}\leftarrow\mathrm{Norm}(\bm{x}+\mathrm{MLP}(\bm{x}))$

Output : $\bm{x}$

Algorithm 3 Dual Attention Decoder Block

An encoder-decoder architecture with causal dual-head attention in the decoder
can be applied to sequence-to-sequence tasks, as in the original Transformer
paper [25]. An encoder-only architecture can be used for a BERT-style language
embedding model [8] or a Vision Transformer-style vision model [9]. A decoder-
only architecture with causal dual-head attention can be used for
autoregressive language modeling222“Decoder-only” is the commonly used term
for this type of architecture [15], but it can also be viewed as an encoder-
only architecture with causal attention.

We note that a standard Transformer is a special case of this architecture
where $n_{h}^{sa}=n_{h},n_{h}^{ra}=0$. Appendix B provides further discussion
on the details of the architecture and its implementation.

##  4 Empirical evaluation

We empirically evaluate the Dual Attention Transformer (abbreviated DAT)
architecture on a range of tasks covering different domains and modalities.
For each experiment, we fix the total number of heads, and compare different
configurations of DAT against a standard Transformer where all heads are self-
attention heads. The difference in performance can be interpreted as
indicating the effect of having two types of attention heads integrating
sensory and relational information. Further experimental details can be found
in Appendix C.

###  4.1 Sample Efficient Relational Reasoning: Relational Games

We begin our empirical evaluation with a benchmark contributed by [20] for
evaluating the relational reasoning capabilities of machine learning models.
The dataset, called “Relational Games”, consists of a family of binary
classification tasks, each testing a model’s ability to identify a particular
visual relationship among a series of objects. The input is an RGB image
depicting a grid of objects, and the target is a binary classification
indicating whether the particular relation holds for this input. We use this
suite of benchmarks to evaluate the sample efficiency of DAT compared to a
standard Transformer. We find that DAT is significantly more sample-efficient,
particularly at more difficult tasks.

Since the input is an image, we use a Vision Transformer-type architecture [9]
where the input image is split up into patches and then fed into the model as
a sequence. We fix the total number of attention heads to 2. We compare a
Vision Transformer with $n_{h}^{sa}=2$ to two configurations of DAT: one with
$n_{h}^{sa}=n_{h}^{ra}=1$ and one with $n_{h}^{sa}=0,n_{h}^{ra}=2$.

We evaluate learning curves by varying the size of the training set, training
each model until convergence, and evaluating on a hold-out validation set. We
repeat this 5 times with different random seeds to compute approximate
confidence intervals. This is depicted in Figure 2. We find that both
configurations of DAT are consistently more sample-efficient compared to the
standard Transformer. The effect is particularly dramatic on the ‘match
pattern’ task which is the most difficult and requires identifying a “second-
order” relation (a relation between relations).

Figure 2: Learning curves on the relational games benchmark. DAT is more
sample-efficient compared to a Transformer with the same total number of
heads. Solid lines indicate the mean over 5 trials with different random seeds
and the shaded regions indicate bootstrap 95% confidence intervals.

In this experiment, we use positional symbols as the symbol assignment
mechanism since the objects can be identified through their position on the
grid. We also impose symmetry on the relations in relational attention, which
we find to be a useful inductive bias. Intuitively, this is because the task-
relevant relations are symmetric similarity relations across different visual
attributes. We provide further discussion and present ablations in Section
C.1.

###  4.2 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 [19]. 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 $n_{h}^{sa}=8$ against two configurations of DAT: one with
$n_{h}^{sa}=n_{h}^{ra}=4$ for the encoder and $n_{h}^{sa}=8,n_{h}^{ra}=0$ for
the decoder (config 1) and another with $n_{h}^{sa}=n_{h}^{ra}=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 3: 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.

###  4.3 Improvements in Language Modeling

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 [10]. 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 $n_{h}^{sa}=8$ attention heads to two configurations of DAT,
one with $n_{h}^{sa}=6,n_{h}^{ra}=2$ and another with
$n_{h}^{sa}=n_{h}^{ra}=4$.

(a) 4 Layers

(b) 5 Layers

(c) 6 Layers

Figure 4: Validation loss curves on a language modeling task. The $x$-axis
indicates the number of tokens and the $y$-axis is the validation loss. DAT
achieves a smaller validation loss for the same total number of attention
heads.

Figure 4 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 $n_{h}^{sa}=n_{h}^{ra}=4$ (the
configuration with a balanced composition of head types).

In Figure 4, 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 Section C.3.

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.

###  4.4 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 [16]. This further probes DAT’
ability in different modalities as a general-purpose sequence model. This
section also stress tests DAT at larger scales.

Figure 5: 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 [9]. 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 $d_{\mathrm{model}}=1024$ and $L=24$ layers. The DAT model uses
position-relative symbols as the symbol assignment mechanism and symmetric
relational attention.

Figure 5 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.

##  5 Discussion

Summary. The standard attention mechanism of Transformers provides a versatile
mechanism for retrieval of sensory information in any given context, but does
not explicitly support retrieval of relational information. In this work, we
presented an extension of the Transformer architecture that disentangles and
integrates sensory and relational information through a variant of multi-head
attention with two distinct types of attention heads: standard self-attention
for sensory information and a novel relational attention mechanism for
relational information. We empirically evaluate this architecture and find
that it yields performance improvements across a range of tasks and
modalities.

Limitations. The proposed architecture introduces several hyperparameters and
possible configurations. Although we carried out ablations on the major
configuration choices (e.g., composition of head types, symmetry, symbol
assignment mechanism), a more thorough empirical examination would help
develop an improved understanding of the behavior of this architecture under
different configurations. Such a systematic study may also enable the
discovery of further modifications to improve the architecture. We also note
that our implementation of the Dual-Attention Transformer lacks the hardware-
aware optimizations of standard Transformers [5], making it slower. However,
we expect that similar optimizations are possible for dual attention and DAT.

## Code and Reproducibility

Our code is publicly available at
https://github.com/Awni00/abstract_transformer. It includes an implementation
of the Dual Attention Transformer, instructions for reproducing our
experiments, and links to detailed experimental logs.

## Acknowledgment

This research was supported by the funds provided by the National Science
Foundation and by DoD OUSD (R&E) under Cooperative Agreement PHY-2229929 (The
NSF AI Institute for Artificial and Natural Intelligence).

## References

  * [1] Awni Altabaa and John Lafferty  “Approximation of Relation Functions and Attention Mechanisms”  arXiv, 2024  DOI: 10.48550/arXiv.2402.08856
  * [2] Awni Altabaa and John Lafferty  “Learning Hierarchical Relational Representations through Relational Convolutions”  arXiv, 2024  DOI: 10.48550/arXiv.2310.03240
  * [3] Awni Altabaa, Taylor Whittington Webb, Jonathan D. Cohen and John Lafferty  “Abstractors and relational cross-attention: An inductive bias for explicit relational reasoning in Transformers”  In _The Twelfth International Conference on Learning Representations_ , 2024  URL: https://openreview.net/forum?id=XNa6r6ZjoB
  * [4] Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey E. Hinton  “Layer Normalization”, 2016  arXiv:1607.06450 [stat.ML]
  * [5] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra and Christopher Ré  “Flashattention: Fast and memory-efficient exact attention with IO-awareness”  In _Advances in Neural Information Processing Systems_ 35, 2022, pp. 16344–16359
  * [6] Yann N Dauphin, Angela Fan, Michael Auli and David Grangier  “Language modeling with gated convolutional networks”  In _International conference on machine learning_ , 2017, pp. 933–941  PMLR
  * [7] Gerard Debreu  “Representation of a preference ordering by a numerical function”  In _Decision processes_ 3 Wiley New York, 1954, pp. 159–165
  * [8] Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova  “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics_ Association for Computational Linguistics, 2019, pp. 4171–4186  DOI: 10.18653/v1/N19-1423
  * [9] Alexey Dosovitskiy et al.  “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale”  In _International Conference on Learning Representations_ , 2021  URL: https://openreview.net/forum?id=YicbFdNTTy
  * [10] Ronen Eldan and Yuanzhi Li  “TinyStories: How Small Can Language Models Be and Still Speak Coherent English?”  arXiv, 2023  DOI: 10.48550/arXiv.2305.07759
  * [11] Dan Hendrycks and Kevin Gimpel  “Gaussian Error Linear Units (GELUs)”, 2016  arXiv:1606.08415 [cs.LG]
  * [12] Hakan Inan, Khashayar Khosravi and Richard Socher  “Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling”  In _International Conference on Learning Representations_ , 2017  URL: https://openreview.net/forum?id=r1aPbsFle
  * [13] Giancarlo Kerg, Sarthak Mittal, David Rolnick, Yoshua Bengio, Blake Richards and Guillaume Lajoie  “On Neural Architecture Inductive Biases for Relational Tasks”  arXiv, 2022  DOI: 10.48550/arXiv.2206.05056
  * [14] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy and Thomas Kipf  “Object-Centric Learning with Slot Attention”  arXiv, 2020  arXiv:2006.15055 [cs, stat]
  * [15] Alec Radford, Karthik Narasimhan, Tim Salimans and Ilya Sutskever  “Improving Language Understanding by Generative Pre-Training”  OpenAI, 2018
  * [16] Olga Russakovsky et al.  “ImageNet Large Scale Visual Recognition Challenge”  In _International Journal of Computer Vision (IJCV)_ 115.3, 2015, pp. 211–252  DOI: 10.1007/s11263-015-0816-y
  * [17] Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu and Timothy Lillicrap  “Relational Recurrent Neural Networks”  In _Advances in Neural Information Processing Systems_ 31 Curran Associates, Inc., 2018
  * [18] Adam Santoro, David Raposo, David G. T. Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia and Timothy Lillicrap  “A Simple Neural Network Module for Relational Reasoning”  arXiv, 2017  arXiv:1706.01427 [cs]
  * [19] David Saxton, Edward Grefenstette, Felix Hill and Pushmeet Kohli  “Analysing Mathematical Reasoning Abilities of Neural Models”  In _International Conference on Learning Representations_ , 2019  URL: https://openreview.net/forum?id=H1gR5iR5FX
  * [20] Murray Shanahan, Kyriacos Nikiforou, Antonia Creswell, Christos Kaplanis, David Barrett and Marta Garnelo  “An Explicitly Relational Neural Network Architecture”  In _Proceedings of the 37th International Conference on Machine Learning_ , 2020
  * [21] Peter Shaw, Jakob Uszkoreit and Ashish Vaswani  “Self-Attention 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 2 (Short Papers)_ New Orleans, Louisiana: Association for Computational Linguistics, 2018, pp. 464–468  DOI: 10.18653/v1/N18-2074
  * [22] Noam Shazeer  “GLU Variants Improve Transformer”  arXiv, 2020  arXiv:2002.05202 [cs, stat]
  * [23] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen and Yunfeng Liu  “RoFormer: Enhanced Transformer with Rotary Position Embedding”  arXiv, 2023  arXiv:2104.09864 [cs]
  * [24] Hugo Touvron et al.  “Llama 2: Open Foundation and Fine-Tuned Chat Models”  arXiv, 2023  DOI: 10.48550/arXiv.2307.09288
  * [25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser and Illia Polosukhin  “Attention Is All You Need”  In _Advances in neural information processing systems_ 30, 2017
  * [26] Taylor W. Webb, Ishan Sinha and Jonathan D. Cohen  “Emergent Symbols through Binding in External Memory”  arXiv, 2021  DOI: 10.48550/arXiv.2012.14601
  * [27] Taylor W. Webb et al.  “The Relational Bottleneck as an Inductive Bias for Efficient Abstraction”  In _Trends in Cognitive Sciences_ , 2024, pp. S1364661324000809  DOI: 10.1016/j.tics.2024.04.001
  * [28] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang and Tieyan Liu  “On layer normalization in the transformer architecture”  In _International Conference on Machine Learning_ , 2020, pp. 10524–10533  PMLR
  * [29] Biao Zhang and Rico Sennrich  “Root mean square layer normalization”  In _Advances in Neural Information Processing Systems_ 32, 2019

##  Appendix A Function Class of Relational Attention: a universal
approximation result

To gain a better understanding of the types of functions that can be computed
by relational attention, we presented a simple approximation result (Theorem
1) in Section 2.4. Here, we will provide a formal statement of the result and
prove it.

Recall that relational attention is a mapping on
$\mathbb{R}^{d}\times\mathbb{R}^{n\times d}\to\mathbb{R}^{d_{\mathrm{out}}}$,
where $d$ is the dimension of the input objects and $d_{\mathrm{out}}$ is the
output dimension. For convenience, we denote the “query space” by
${\mathcal{X}}$ and the “key space” by ${\mathcal{Y}}$, though both are
$\mathbb{R}^{d}$ in this setting. Relational attention takes as input a query
$x\in{\mathcal{X}}$ and a collection of objects
$\bm{y}=(y_{1},\ldots,y_{n})\in{\mathcal{Y}}^{n}$ and computes the following

$\displaystyle\mathrm{RA}(x,\bm{y})$
$\displaystyle=\sum_{i=1}^{n}\alpha_{i}(x;\bm{y})\big{(}r(x,y_{i})\,W_{r}+s_{i}\,W_{s}\big{)},$
(5) $\displaystyle\alpha(x;\bm{y})$
$\displaystyle=\mathrm{Softmax}\Big{(}\big{[}\left\langle\phi_{q}^{\mathrm{attn}}(x),\phi_{k}^{\mathrm{attn}}(y_{i})\right\rangle\big{]}_{i=1}^{n}\Big{)}\in\Delta^{n},$
(6) $\displaystyle r(x,y_{i})$
$\displaystyle=\left(\left\langle\phi_{q,\ell}^{\mathrm{rel}}(x),\phi_{k,\ell}^{\mathrm{rel}}(y_{i})\right\rangle\right)_{\ell\in[d_{r}]}\in\mathbb{R}^{d_{r}},$
(7) $\displaystyle(s_{1},\ldots,s_{n})$
$\displaystyle=\mathrm{SymbolRetriever}\left(\bm{y};\,S_{\mathrm{lib}}\right)\in\mathbb{R}^{n\times
d_{\mathrm{out}}},$ (8)

where
$\phi_{q}^{\mathrm{attn}},\phi_{k}^{\mathrm{attn}},\phi_{q,\ell}^{\mathrm{rel}},\phi_{k,\ell}^{\mathrm{rel}}:\mathbb{R}^{d}\to\mathbb{R}^{d_{k}}$
are the feature maps defining the attention mechanism and the relation,
respectively. For this section, these are multi-layer perceptrons. Note that
in Algorithm 1 these are linear maps, but they are preceded by multi-layer
perceptron in Algorithms 2 and 3, which makes the overall function class the
same. Moreover, for this analysis we will take
$W_{r}=I,d_{\mathrm{out}}=d_{r}$ and $W_{s}=0$. We will later discuss how the
role of symbols fits within the message of the result.

The following result states that relational attention can approximate any
function of the form: 1) select an object in $(y_{1},\ldots,y_{n})$ by an
arbitrary query-dependent selection criterion, and 2) compute an arbitrary
relation $r:{\mathcal{X}}\times{\mathcal{Y}}\to\mathbb{R}^{d_{r}}$ with the
selected object. This is formalized below.

To formalize (1), we adopt an abstract and very general formulation of a
“selection criterion” in terms of family of preference preorders,
$\\{\preccurlyeq_{x}\\}_{x}$: for each possible query $x$, the preorder
$\preccurlyeq_{x}$ defines a preference over objects in ${\mathcal{Y}}$ to be
selected. Intuitively, “$y_{1}\preccurlyeq_{x}y_{2}$” means that $y_{2}$ is
more relevant to the query $x$ than $y_{1}$.

More precisely, for each query $x\in{\mathcal{X}}$, $\preccurlyeq_{x}$ is a
complete (for each $y_{1},y_{2}\in{\mathcal{Y}}$, either $y_{1}\preccurlyeq
y_{2}$ or $y_{2}\preccurlyeq_{x}y_{1}$), reflexive ($y\preccurlyeq_{x}y$ for
all $y\in{\mathcal{Y}}$), and transitive ($y_{1}\preccurlyeq_{x}y_{2}$ and
$y_{2}\preccurlyeq_{x}y_{3}$ implies $y_{1}\preccurlyeq_{x}y_{3}$) relation.
For each $x\in{\mathcal{X}}$, $\preccurlyeq_{x}$ induces a preordered space
$({\mathcal{Y}},\preccurlyeq_{x})$. This implicitly defines two additional
relations: $\prec_{x}$ and $\sim_{x}$. We will write $y_{1}\prec_{x}y_{2}$ if
“$y_{1}\preccurlyeq_{x}y_{2}$ and not $y_{2}\preccurlyeq_{x}y_{1}$”, and
$y_{1}\sim y_{2}$ if “$y_{1}\preccurlyeq_{x}y_{2}$ and
$y_{2}\preccurlyeq_{x}y_{1}$”.

For a collection of objects $\bm{y}=(y_{1},\ldots,y_{n})\in{\mathcal{Y}}^{n}$
and a query $x\in{\mathcal{X}}$, the preorder $\preccurlyeq_{x}$ defines a
selection function

$\mathrm{Select}(x,(y_{1},\ldots,y_{n}))\coloneq\max\left((y_{1},\ldots,y_{n}),\mathtt{key}=\preccurlyeq_{x}\right).$
(9)

That is, $\mathrm{Select}(x,\bm{y})$ returns the most relevant element with
respect to the query $x$. In particular, it returns $y_{i}$ when
$y_{i}\succ_{x}y_{j},\ \forall j\neq i$ (and may return an arbitrary element
if no unique maximal element exists in $(y_{1},\ldots,y_{n})$).

We will assume some regularity conditions on the family of preorders
$\\{\preccurlyeq_{x}\\}_{x}$ which essentially stipulate that: 1) nearby
elements in ${\mathcal{Y}}$ have a similar preference with respect to each
$x$, and 2) nearby queries in ${\mathcal{X}}$ induce similar preference
preorders.

######  Assumption 1 (Selection criterion is query-continuous and key-
continuous).

The family of preorder relations $\\{\preccurlyeq_{x}\\}_{x\in{\mathcal{X}}}$
satisfies the following:

  1. 1.

Key-continuity. For each $x\in{\mathcal{X}}$, $\preccurlyeq_{x}$ is
continuous. That is, for any sequence $(y_{i})_{i}$ such that
$y_{i}\preccurlyeq_{x}z$ and $y_{i}\to y_{\infty}$, we have
$y_{\infty}\preccurlyeq_{x}z$. Equivalently, for any $y\in{\mathcal{Y}}$,
$\\{z\in{\mathcal{Y}}:z\preccurlyeq_{x}y\\}$ and
$\\{z\in{\mathcal{Y}}:y\preccurlyeq_{x}z\\}$ are closed sets in
${\mathcal{Y}}$.

  2. 2.

Query-continuity. Under key-continuity, [7] shows that for each
$x\in{\mathcal{X}}$, there exists a continuous in utility function
$u_{x}:{\mathcal{Y}}\to\mathbb{R}$ for $\preccurlyeq_{x}$ such that
$y_{1}\preccurlyeq_{x}y_{2}\iff u_{x}(y_{1})\leq u_{x}(y_{2})$. For query-
continuity, we make the further assumption that there exists a family of
utility functions $\\{u_{x}:{\mathcal{Y}}\to\mathbb{R}\\}_{x\in{\mathcal{X}}}$
such that $u(x,y)\coloneq u_{x}(y)$ is also continuous in its first argument.

For technical reasons, for Equation 9 to make sense, we must assume that there
exists a unique element to be selected. We formulate this in terms of an
assumption on the data distribution of the space
${\mathcal{X}}\times{\mathcal{Y}}^{n}$. This is a technical assumption, and
different forms of such an assumption would be possible (e.g., instead
condition on this event).

######  Assumption 2 (Selection is unique almost always).

Let $(x,\bm{y})\sim{\mathbb{P}}_{x,\bm{y}}$. For each $\varepsilon>0$, there
exists $\eta_{\varepsilon}>0$ such that $\min_{j\neq i}\left\lvert
u_{x}(y_{i})-u_{x}(y_{j})\right\rvert>\eta_{\varepsilon}$ with probability at
least $1-\varepsilon$.

######  Theorem (Function class of relational attention).

Let ${\mathcal{X}},{\mathcal{Y}}$ be compact Euclidean spaces. Let
$\\{\preccurlyeq_{x}\\}_{x\in{\mathcal{X}}}$ be an arbitrary family of
relevance preorders on ${\mathcal{Y}}$ which are query-continuous and key-
continuous (Assumption 1). Let
$\mathrm{Select}(x,(y_{1},\ldots,y_{n}))=\max((y_{1},\ldots,y_{n}),\mathtt{key}=\preccurlyeq_{x})$
be the selection function associated with $\\{\preccurlyeq_{x}\\}_{x}$. Let
$R:{\mathcal{X}}\times{\mathcal{Y}}\to\mathbb{R}^{d_{r}}$ be an arbitrary
continuous relation function. Suppose $x,\bm{y}\sim{\mathbb{P}}_{x,\bm{y}}$
and that Assumption 2 holds (i.e., the data distribution is such that there
exists a unique most-relevant element). For any $\varepsilon>0$, there exists
multi-layer perceptrons
$\phi_{q}^{\mathrm{attn}},\phi_{k}^{\mathrm{attn}},\phi_{q}^{\mathrm{rel}},\phi_{k}^{\mathrm{rel}}$
and a choice of symbols such that,

$\left\lVert\mathrm{RA}(x,(y_{1},\ldots,y_{n}))-R(x,\mathrm{Select}(x,(y_{1},\ldots,y_{n})))\right\rVert_{\infty}<\varepsilon$

###### Proof.

Condition on the event
${\mathcal{E}}:=\\{(x,\bm{y})\in{\mathcal{X}}\times{\mathcal{Y}}^{n}\,\colon\,\min_{j\neq
i}\left\lvert u_{x}(y_{i})-u_{x}(y_{j})\right\rvert>\eta_{\varepsilon}\\}$.
Let
$i^{*}=\operatorname*{arg\,max}((y_{1},\ldots,y_{n}),\mathtt{key}=\preccurlyeq_{x})=\operatorname*{arg\,max}(u_{x}(y_{1}),\ldots,u_{x}(y_{n}))$.
By [1, Theorem 5.1], for any $\varepsilon_{1}>0$, there exists MLPs
$\phi_{q}^{\mathrm{attn}},\phi_{k}^{\mathrm{attn}}$ such that
$\alpha_{i^{*}}(x,\bm{y})>1-\varepsilon_{1}$ for any
$(x,\bm{y})\in{\mathcal{E}}$. That is, the attention score is nearly $1$ for
the $\preccurlyeq_{x}$-selected element uniformly over inputs in
${\mathcal{E}}$.

Similarly, by [1, Theorem 3.1], for any $\varepsilon_{2}>0$, there exists MLPs
$(\phi_{q,\ell}^{\mathrm{rel}},\phi_{k,\ell}^{\mathrm{rel}})_{\ell\in[d_{r}]}$
such that
$r(x,y)\coloneq(\langle\phi_{q,\ell}^{\mathrm{rel}}(x),\phi_{k,\ell}^{\mathrm{rel}}(y)\rangle)_{\ell\in[d_{r}]}$
approximates the target relation $R$ uniformly within an error of
$\varepsilon_{2}$,

$\left\lVert
R(x,y)-r(x,y)\right\rVert_{\infty}<\varepsilon_{2},\quad\text{Lebesgue almost
every }(x,y)\in{\mathcal{X}}\times{\mathcal{Y}}.$

Thus, we have

$\displaystyle\left\lVert\mathrm{RA}(x,(y_{1},\ldots,y_{n}))-R(x,\mathrm{Select}(x,(y_{1},\ldots,y_{n})))\right\rVert_{\infty}$
$\displaystyle=\left\lVert\sum_{i=1}^{n}\alpha_{i}(x;\bm{y})\,r(x,y_{i})-R(x,y_{i^{*}})\right\rVert_{\infty}$
$\displaystyle\leq\sum_{i=1}^{n}\left\lVert\alpha_{i}(x;\bm{y})\,r(x,y_{i})-R(x,y_{i^{*}})\right\rVert_{\infty}$
$\displaystyle\leq\alpha_{i^{*}}(x,\bm{y})\left\lVert
r(x,y_{i^{*}})-R(x,y_{i^{*}})\right\rVert_{\infty}+\sum_{j\neq
i^{*}}\alpha_{i}(x;\bm{y})\left\lVert
r(x,y_{i})-R(x,y_{i^{*}})\right\rVert_{\infty}$
$\displaystyle\leq(1-\varepsilon_{1})\varepsilon_{2}+\varepsilon_{1}\max_{x,y,y^{*}}\left\lVert
r(x,y)-R(x,y^{*})\right\rVert_{\infty}.$

Note that $\max_{x,y,y^{*}}\left\lVert r(x,y)-R(x,y^{*})\right\rVert_{\infty}$
is finite since ${\mathcal{X}},{\mathcal{Y}}$ are compact and $r,R$ are
continuous. Letting $\varepsilon_{1},\varepsilon_{2}$ be small enough
completes the proof. ∎

To summarize the analysis in this section, we showed that relational attention
can approximate any computation composed of selecting an object from a
collection followed by computing a relation with that object. We can
approximate any well-behaved selection criterion by formulating it in terms of
an abstract preference preorder, and approximating the corresponding utility
function (given by a Debreu representation theorem) by inner products of query
and key feature maps. We can then approximate the target relation function
similarly by inner products of a different set of query and key feature maps.

In the analysis above, we set aside the role of the symbols. Note that the
function class this approximation result proves involves retrieving a relation
from a selected object, but does not explicitly encode the identity of the
selected object. Informally, the receiver knows that it has a particular
relation with one of the objects in its context, and knows that this relation
is with an object that was selected according to a particular selection
criterion, but does not know the identity of the object beyond that. This is
the purpose of adding symbols to relational attention—the retrieved relation
is tagged with a symbol identifying the sender.

##  Appendix B Architecture & Implementation Details

In this section, we briefly discuss some details of implementation that may be
of interest to the reader. Our code is publicly available through the project
git repository and includes detailed instructions for reproducing our
experimental results. We also make available detailed experimental logs for
each experimental run which include: the git commit ID giving the version of
the code that was used to run the experiment, the exact command and script
used to run the script, the hardware used for that run, and metrics tracked
over the course of training. Our code uses the PyTorch framework.

###  B.1 Relational Attention and Dual-Head Attention

The relational attention operation is defined as part of dual-head attention
in Algorithm 1. We briefly mention some details of the implementation. Let
$n_{h}=n_{h}^{sa}+n_{h}^{ra}$ be the total number of parameters The learnable
parameters here are
$W_{q}^{h},W_{k}^{h}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{\mathrm{key}}},\,W_{v}^{h}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{h}},\,h\in[n_{h}^{sa}],W_{o}^{sa}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{\mathrm{model}}}$, for the self-attention heads, and
$W_{q,h}^{\mathrm{attn}},W_{k,h}^{\mathrm{attn}}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{\mathrm{key}}},W_{s}^{h}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{h}},W_{r}^{h}\in\mathbb{R}^{d_{r}\times
d_{h}},\,h\in[n_{h}^{ra}],W_{q,\ell}^{\mathrm{rel}},W_{k,\ell}^{\mathrm{rel}}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{\mathrm{proj}}},\ell\in[d_{r}],W_{o}^{ra}\in\mathbb{R}^{d_{\mathrm{model}}\times
d_{\mathrm{model}}}$ for the relational attention heads. We let
$d_{\mathrm{key}},d_{h}=d_{\mathrm{model}}/n_{h}$ to maintain the same
dimension for the input and output objects. By default, and for all our
experiments, we let $d_{r}=n_{h}^{ra}$ and
$d_{\mathrm{proj}}=d_{\mathrm{key}}$. To model symmetric relations, we let
$W_{q,\ell}^{\mathrm{rel}}=W_{k,\ell}^{\mathrm{rel}}$.

Note that the same $d_{r}$-dimensional relation is used for all $n_{h}^{ra}$
attention heads, with a learned linear map $W_{r}^{h}$ for each head
extracting the relevant aspects of the relation for that attention head. This
allows for useful computations to be shared across all heads. Note also that
the head dimension $d_{h}=d_{\mathrm{model}}/n_{h}$ is defined in terms of the
total number of attention heads and is the same per-head for both self-
attention and relational attention. This means that the number of components
in the $d_{\mathrm{model}}$-dimensional output corresponding to each attention
head type is proportional to the number of heads of that type. For example, if
$n_{h}^{sa}=6,n_{h}^{ra}=2$, then 75% of the $d_{\mathrm{model}}$-dimensional
output is composed of the output of self-attention heads and 25% is composed
of the output of relational attention heads. This enables tuning the relative
importance of each head type for the task.

Finally, we note that relational attention can be implemented with an einsum
operation. For example, in PyTorch, it looks something like the following.

⬇

# attn_scores: (bs, n_heads, seqlen, seqlen)

# relations: (bs, seqlen, seqlen, n_heads, head_dim)

attended_symbols = torch.einsum(’bhij,bjhd->bihd’, attn_scores, sv)

# shape: (bs, seqlen, n_heads, head_dim)

attended_relations = torch.einsum(’bhij,bijhd->bihd’, attn_scores,
proj_relations)

# shape: (bs, seqlen, n_heads, head_dim)

output = attended_symbols + attended_relations

# shape: (bs, seqlen, n_heads, head_dim)

Here,
$\texttt{sv[:,:,h,:]}=\bm{s}\,W_{s}^{h},\,\texttt{attn\\_scores[:,h,:,:]}=\bm{\alpha}^{h}$,
and $\texttt{relations[:,:,h,:]}=\bm{r}\,W_{r}^{h}$, which can all be computed
with simple matrix multiplication operations. When using position-relative
symbols, the first line is replaced with attended_symbols =
torch.einsum(’bhij,ijhd->bihd’, attn_scores, sv), assuming that
$\texttt{sv[i,j,h,:]}=s_{j-i}W_{s}^{h}$

Weight-tying self-attention and relational-attention heads. To increase
parameter efficiency and perhaps improve performance, we suggest weight-tying
the attention query/key projections across self-attention and relational
attention heads. The intuition is that the same selection criteria may be
useful for retrieving either sensory or relational information from a
particular object. By sharing these projections across heads, we can reduce
the number of trainable parameters, and introduce an additional supervisory
signal which may improve performance. We leave testing this idea for future
work.

Composing relational attention. We remark that composing relational attention
modules can be interpreted as representing hierarchical or higher-order
relations. That is, relations between relations. An example of this is the
relation tested in the match pattern task in the relational games benchmark.
After one iteration of relational attention, an object’s representation is
updated with the relations it has with its context. A second iteration of
relational attention now computes a representation of the relation between an
object’s relations and the relations of the objects in its context.

###  B.2 Encoder and Decoder Blocks

We briefly mention a few configurations in our implementation that appear in
our experiments. We aimed to make our implementation configurable to allow for
various tweaks and optimizations that have been found in the literature for
training Transformer models.

Symbol assignment. A shared symbol assignment module is used for all layers in
the model. We explore three types of symbol assignment mechanisms: positional
symbols, position-relative symbols, and symbolic attention. Different symbol
assignment mechanisms are more well-suited to different tasks. We discuss
ablation experiments we carried out on the effect of the symbol assignment
mechanism in Appendix C.

MLP block. The MLP block uses a 2-layer feedforward network with a
configurable activation function. The intermediate layer size is
$d_{\mathrm{ff}}=4\cdot d_{\mathrm{model}}$ by default. We also use the SwiGLU
“activation function” [22] in some of our experiments. SwiGLU is not merely an
activation function, but is rather a neural network layer defined as the
component-wise product of two linear transformations of the input. It is a
type of gated linear unit [6] with the sigmoid activation replaced with a
Swish activation [11], $\mathrm{SwiGLU}(x)=\mathrm{Swish}(xW+b)\otimes(xV+c)$.
This is used in the Llama series of models and was found to be a useful
modification [24].

Normalization. Either LayerNorm [4] or RMSNorm [29] can be used. Normalization
can be performed post-attention, like in the original Transformer paper [25],
or pre-attention as in [28].

Positional encoding. Our experiments use either learned positional embeddings
or RoPE [23].

##  Appendix C Experimental Details & Further Discussion

###  C.1 Relational Games (Section 4.1)

#### Experimental details

Dataset details. The relational games benchmark datasets consists of $36\times
36\times 3$ RGB images depicting a $3\times 3$ grid of objects which satisfy a
particular visual relationship. The set of objects consists of simple
geometric shapes. For example, in the occurs task, one object is present in
the top row and three in the bottom row, and the task is to determine whether
the object in the top row occurs (i.e., is among) the objects in the bottom
row. The most difficult task in the benchmark is the match pattern task, where
the grid contains a triplet of objects in the top row and another triplet of
objects in the bottom row. Each triplet satisfies some relationship (e.g.,
ABC, ABA, ABB, or AAB), and the task is to determine whether the relation in
the first triplet is the same as the relation in the second triplet. The
difficulty in solving this task is that it requires parsing a second-order
relation (a relation between relations). We remark that composing relational
attention modules naturally captures this kind of hierarchical relations: one
relational attention operation produces objects with a relational
representation and another would compute relations between those relations.

Model architectures. We use a Vision-Transformer-type architecture where the
input image is split up into patches, flattened, and passed through the
sequence model with added learned positional embeddings. We use average
pooling at the end and pass through an MLP to produce the final prediction. We
use a patch size of $12\times 12$ which separates objects according to the
grid structure. We note that in more general visual relational reasoning tasks
where there isn’t this type of grid structure, it would be appropriate to
combine our approach with an object-discovery module such as Slot Attention
[14].

We use the following hyperparameters: 2 layers, $d_{\mathrm{model}}=128$,
$d_{\mathrm{ff}}=256$, SwiGLU “activation”, dropout rate = 0.1, and pre-
LayerNormalization. For the DAT models, we use positional symbols as the
symbol assignment mechanism. The total number of heads is 2. For the
Transformer model, there are only self-attention heads: $n_{h}^{sa}=2$. For
DAT, we evaluated two configurations for the composition of head types, one
with $n_{h}^{sa}=n_{h}^{ra}=1$ and one with only relational attention heads
$n_{h}^{sa}=0,n_{h}^{ra}=2$. We also evaluated variants with and without the
constraint that the relations in relational attention are symmetric (i.e.,
$W_{q}^{\mathrm{rel}}=W_{k}^{\mathrm{rel}}$).

Training details. For each task and model, we evaluated learning curves by
varying the training set size and training the model until convergence, then
evaluating on a hold-out test set. For four out of five of the tasks, we
evaluate learning curves within the range of $250$ to $2,500$ samples, in
increments of $250$. For the more difficult match pattern, the range is from
$5,000$ to $25,000$ in increments of $5,000$. The ranges were chosen based on
the difficulty of the different tasks in order to identify the right
“resolution”. When evaluating learning curves, each training set is sampled
randomly from the full dataset. For each task, model, and training set size,
we repeat the experiment 5 times with different random seeds to compute
approximate confidence intervals (accounting for randomness in sampling the
dataset and random initialization). We use an Adam optimizer with a learning
rate of $0.001$, $\beta_{1}=0.9,\beta_{2}=0.99$, and a batch size of $512$. We
train for 50 epochs.

#### Further Discussion, Exploration, & Ablations

We performed an ablation over the symmetry inductive bias in the relations
computed in relational attention. Our implementation exposes an argument which
controls whether the relation $r(x,y)=(\langle
W_{q,\ell}^{\mathrm{rel}},W_{k,\ell}^{\mathrm{rel}}\rangle)_{\ell\in[d_{r}]}\in\mathbb{R}^{d_{r}}$
modeled in relational attention is constrained to be symmetric by setting
$W_{q,\ell}^{\mathrm{rel}}=W_{k,\ell}^{\mathrm{rel}}$. Indeed, we find
symmetry to be a useful inductive bias in this task. Figure 6 depicts learning
curves for the two configurations of DAT comparing symmetric RA against
asymmetric RA. We find that symmetry results in faster learning curves for
both configurations.

Figure 6: An ablation of the effect of symmetry in relational attention in the
relational games experiments.

###  C.2 Mathematical Problem-Solving (Section 4.2)

#### Experimental details

Dataset details. [19] propose a benchmark to assess neural models’ ability to
perform mathematical reasoning. The dataset consists of a suite of tasks in
free-form textual input/output format. The tasks cover several topics in
mathematics, including arithmetic, algebra, and calculus. For each task, the
authors programmatically generate $2\times 10^{6}$ training examples and
$10^{4}$ validation examples. Questions have a maximum length of 160
characters and answers have a maximum length of 30 characters.

Model architectures. We use an encoder-decoder architecture for this
experiment, treating it as a sequence-to-sequence task. We use character-level
encoding with a common alphabet of size 85 containing small and upper case
letters, digits 0-9, and symbols (e.g., *, /, +, -). The models have 2-layer
encoders and decoders, with $d_{\mathrm{model}}=128,d_{\mathrm{ff}}=256$, ReLU
activation, dropout rate = 0.1, and post-normalization. Sinusoidal positional
embeddings are used as the positional encoding method. For all models, the
total number of attention heads (across self-attention and relational
attention) is $8$. For the Transformer model, there are only self-attention
heads: $n_{h}^{sa}=8$ for both the encoder and decoder. For DAT, we evaluated
two configurations for the composition of head types, one with
$n_{h}^{sa}=n_{h}^{ra}=4$ in the encoder and $n_{h}^{sa}=8,n_{h}^{ra}=0$ in
the decoder (i.e., standard Transformer Decoder), and one with
$n_{h}^{sa}=4=n_{h}^{ra}=4$ in the encoder and $n_{h}^{sa}=4=n_{h}^{ra}=4$ in
the decoder. The number of cross-attention heads in the decoder is $8$ in all
cases. No symmetry constraint is made on relational attention. Position-
relative symbols are used as the symbol assignment mechanism, and the symbol
library is shared across all layers in both the encoder and decoder.

Training Details. Each model is trained on each task for 50 epochs. We use the
Adam optimizer with $\beta_{1}=0.9,\beta_{2}=0.995$, a learning rate of
$6\times 10^{-4}$, and a batch size of 128. We evaluate and track the per-
character accuracy over the course of training. We repeat this process 5 times
for each combination of model and task with different random seeds to compute
approximate confidence intervals.

#### Further Discussion, Exploration, & Ablations

This benchmark contains two validation splits: an interpolation split and an
extrapolation split. The interpolation split is generated by the same
procedure as the training split. The extrapolation split is generated with
different parameters to make the task more difficult. We refer to the original
reference for more details [19]. Figure 3 in the main text depicts the
interpolation validation accuracy. Figure 7 depicts the extrapolation
validation accuracy over the course of training. For completeness, we also
show the training accuracy in Figure 8.

Figure 7: Extrapolation validation accuracy over the course of training for
mathematical problem-solving tasks. Figure 8: Training accuracy over the
course of training for mathematical problem-solving tasks.

###  C.3 Language Modeling (Section 4.3)

#### Experimental details

Dataset details. The Tiny Stories dataset by [10] is a language modeling
benchmark designed to evaluate small language models. The dataset consists of
short stories and is synthetically generated by GPT-3.5 and GPT-4.

Model architectures. We use a Decoder-only architecture, with causal attention
for autoregressive language modeling. We fix the total number of attention
heads to be $8$ for all models. For the Transformer model, there are only
self-attention heads: $n_{h}^{sa}=2$. For DAT, we evaluated two configurations
for the composition of head types, one with $n_{h}^{sa}=6,n_{h}^{ra}=2$ (more
self-attention heads) and one with $n_{h}^{sa}=n_{h}^{ra}=4$ (a balanced
composition of head types). We experiment with
$d_{\mathrm{model}}\in\\{64,128\\}$ and number of layers $L\in\\{4,5,6\\}$.
For all models, we use $d_{\mathrm{ff}}=4d_{\mathrm{model}}$, SwiGLU
activation in the MLP, RoPE positional encoding, no bias, dropout rate = 0.1,
and pre-LayerNormalization. We use weight-tying between the embedding layer
and the final prediction layer [12]. For DAT models, we perform ablations on
the effect of the choice of symbol assignment mechanism and symmetry in
relational attention. We train models with position-relative symbols and
symbolic attention, as well as models with symmetric and asymmetric relations
in relational attention. We use a context size of 512 for all models. Text is
tokenized using the Llama SentencePiece tokenizer [24], which contains
$32,000$ tokens.

Figure 4 in the main text depicts the $d_{\mathrm{model}}=64$ models with
varying number of layers $L\in\\{4,5,6\\}$. Figure 9 shows the loss curves for
larger models with $d_{\mathrm{model}}=128$. The DAT models use symbolic
attention as the symbol assignment mechanism and asymmetric relations in
relational attention. Symbolic attention uses
$n_{s}=\texttt{context\\_length}=512$ symbols with 4 heads.

Training Details. All models are trained with the AdamW optimizer with a
constant learning rate of $0.001$ and $\beta_{1}=0.9,\beta_{2}=0.95$. We clip
gradients to a norm of 1. Models are trained for $100,000$ iterations. We use
a batch size of $128$ with $2$ gradient accumulation steps. Thus, each step
trains on $131,072$ tokens.

Table 1: Language Modeling on the Tiny Stories Dataset $d_{\mathrm{model}}$ | $L$ | $n_{h}^{sa}$ | $n_{h}^{ra}$ | Symbol Assignment Mechanism | Symmetric RA | val/loss | val/perplexity
---|---|---|---|---|---|---|---
64 | 4 | 4 | 4 | Position-Relative Symbols | False | 1.764 | 5.840
True | 1.785 | 5.963
Symbolic Attention | False | 1.729 | 5.639
True | 1.744 | 5.722
6 | 2 | Position-Relative Symbols | False | 1.768 | 5.859
True | 1.777 | 5.914
Symbolic Attention | False | 1.740 | 5.697
True | 1.745 | 5.727
8 | 0 | NA | NA | 1.775 | 5.903
5 | 4 | 4 | Symbolic Attention | False | 1.692 | 5.431
True | 1.698 | 5.467
6 | 2 | Position-Relative Symbols | True | 1.730 | 5.640
Symbolic Attention | False | 1.692 | 5.432
True | 1.704 | 5.495
8 | 0 | NA | NA | 1.730 | 5.640
6 | 4 | 4 | Position-Relative Symbols | False | 1.685 | 5.395
True | 1.704 | 5.498
Symbolic Attention | False | 1.656 | 5.239
True | 1.668 | 5.303
6 | 2 | Position-Relative Symbols | True | 1.691 | 5.424
Symbolic Attention | False | 1.663 | 5.277
True | 1.669 | 5.308
8 | 0 | NA | NA | 1.692 | 5.431
128 | 4 | 4 | 4 | Symbolic Attention | False | 1.411 | 4.102
True | 1.417 | 4.127
6 | 2 | Symbolic Attention | False | 1.412 | 4.105
True | 1.415 | 4.118
8 | 0 | NA | NA | 1.431 | 4.183
6 | 4 | 4 | Symbolic Attention | False | 1.337 | 3.811
True | 1.346 | 3.843
6 | 2 | Symbolic Attention | False | 1.340 | 3.818
True | 1.346 | 3.843
8 | 0 | NA | NA | 1.353 | 3.870
Figure 9: Loss curves of DAT language models with $d_{\mathrm{model}}=128$.

### Further Discussion, Exploration, & Ablations

(a) $d=64,L=4$, asymmetric RA.

(b) $d=64,L=4$, symmetric RA.

Figure 10: Ablation of symbol assignment mechanism. Symbolic attention
outperforms position-relative symbols.

(a) $d=64,L=4$, symbolic attention.

(b) $d=64,L=4$, position-relative symbol assignment

Figure 11: Ablation of Symmetry in relational attention. The DAT models with
asymmetric relations perform slightly better than those with a symmetry
constraint.

Table 1 shows the end-of-training validation loss and perplexity for all
models we evaluated. For each $d_{\mathrm{model}},L$ the best-performing model
is bolded. We find that the DAT model with asymmetric relational attention and
symbolic attention consistently performs the best, beating out the Transformer
by a small margin.

Figure 10 depicts an ablation over symbol type. For
$d_{\mathrm{model}}=64,L=4$, it shows training curves for DAT models with two
types of symbol retrieval mechanism: symbolic attention and position-relative
symbol. We find that the symbolic attention variant learns faster and reaches
a smaller loss. This holds both when relational attention uses symmetric or
asymmetric relations. We posit the explanation that the differentiable
equivalence class mapping implemented by symbolic attention may capture a form
of syntax.

Figure 11 shows an ablation over the symmetry of relational attention.
Relational attention can be constrained to represent symmetric relations by
imposing that $W_{q}^{\mathrm{rel}}=W_{k}^{\mathrm{rel}}$. This is sometimes a
useful inductive bias, as observed in the relational games experiments.
However, in the language modeling experiments, we found that it was not a
useful inductive bias and that asymmetric relational attention outperformed
symmetric relational attention. One possible explanation is that the types of
relations that are relevant in parsing language are in general asymmetric. For
example, syntactic or grammatical relations such as noun-verb, subject-object,
determiner-noun, etc. are asymmetric.

###  C.4 Image Recognition (Section 4.4)

#### Experimental details

Dataset details. For this experiment, we use the ImageNet Large Scale Visual
Recognition Challenge 2012 (ILSVRC2012) dataset [16]. The dataset is roughly
140GB in size and consists of RGB images hand-labeled with the presence or
absence of 1,000 object categories. We train on images resized to $224\times
224$. In the training set, we perform a random crop and resize to $224\times
224$, then normalize channel-wise with means $(0.485,0.456,0.406)$ and
standard deviations $(0.229,0.224,0.225)$. For the validation set, we resize
to $256\times 256$ then center crop to $224\times 224$, and normalize in the
same way.

Model architectures. We use a Vision Transformer-style architecture [9].
ImageNet’s $224\times 224\times 3$ 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
$d_{\mathrm{model}}=1024$, $L=24$ layers, MLP hidden dimension
$d_{\mathrm{ff}}=4096$, SwiGLU activation, no bias, dropout rate = 0.1, and
pre-LayerNormalization. The DAT model uses position-relative symbols as the
symbol assignment mechanism and symmetric relational attention.

Training details. Both models are trained with the AdamW optimizer, with a
constant learning rate of $5\times 10^{-4}$ and
$\beta_{1}=0.9,\beta_{2}=0.99$. We used a batch size of 32 and 32 gradient
accumulation steps.

##  Appendix D Comparison to [3]: Abstractors and Relational Cross-Attention

A closely related work is [3], which proposes a Transformer-based module
called the “Abstractor” with relational inductive biases. The core operation
in the Abstractor is a variant of attention dubbed “relational cross-
attention” (RCA). In this section, we will discuss the relation between the
Dual Attention Transformer and the Abstractor.

###  D.1 Comparison between RA (this work) and RCA [3]

[3] propose a variant of attention called relational cross-attention which
shares some characteristics with our proposal of what we’re calling
“relational attention” in this work. In this discussion, we will use the
acronyms RCA and RA, respectively to distinguish between the two.

RCA processes a sequence of objects $\bm{x}=(x_{1},\ldots,x_{n})$ and produces
a sequence of objects $\bm{x}^{\prime}=(x_{1}^{\prime},\ldots,x_{n}^{\prime})$
via the following operation

$\displaystyle\bm{x}^{\prime}$
$\displaystyle\leftarrow\sigma_{\mathrm{rel}}\left(\phi_{q}(\bm{x}){\phi_{k}(\bm{x})}^{\intercal}\right)\bm{s},$
$\displaystyle\bm{s}$ $\displaystyle=\mathrm{SymbolRetriever}(\bm{x})$

where $\phi_{q},\phi_{k}$ are query and key transformations, and the symbols
$\bm{s}$ take the same role as in this work. $\sigma_{\mathrm{rel}}$ is
referred to as a “relation activation”. It may be either softmax or an
element-wise activation (e.g., tanh, sigmoid, or linear). For the purposes of
this discussion, let us consider $\sigma_{\mathrm{rel}}=\mathrm{Softmax}$,
which was used in the majority of the experiments in [3].

To facilitate the discussion, let us write RA and RCA side-by-side using a
common notation.

RA (this work) RCA [3] $\displaystyle(x_{1}^{\prime},\ldots,x_{n}^{\prime})$
$\displaystyle\leftarrow\mathrm{RA}(\bm{x};S_{\mathrm{lib}}),$
$\displaystyle(x_{1}^{\prime},\ldots,x_{n}^{\prime})$
$\displaystyle\leftarrow\mathrm{RCA}(\bm{x};S_{\mathrm{lib}})$ $\displaystyle
x_{i}^{\prime}$
$\displaystyle=\sum_{j=1}^{n}\alpha_{ij}\,\big{(}W_{r}\,r(x_{i},x_{j})+W_{s}\,s_{j}\big{)},$
$\displaystyle x_{i}^{\prime}$
$\displaystyle=\sum_{j=1}^{n}\alpha_{ij}\,s_{j},$ $\displaystyle\bm{\alpha}$
$\displaystyle=\mathrm{Softmax}\big{(}\phi_{q}(\bm{x}){\phi_{k}(\bm{x})}^{\intercal}\big{)},$
$\displaystyle\bm{\alpha}$
$\displaystyle=\mathrm{Softmax}\big{(}\phi_{q}(\bm{x}){\phi_{k}(\bm{x})}^{\intercal}\big{)},$
$\displaystyle r(x,y)$
$\displaystyle=\big{(}\left\langle\phi_{q,\ell}^{\mathrm{rel}}(x),\phi_{k,\ell}^{\mathrm{rel}}(y)\right\rangle\big{)}_{\ell\in[d_{r}]},$
$\displaystyle(s_{1},\ldots,s_{n})$
$\displaystyle=\mathrm{SymbolRetriever}(\bm{x};\,S_{\mathrm{lib}})$
$\displaystyle(s_{1},\ldots,s_{n})$
$\displaystyle=\mathrm{SymbolRetriever}(\bm{x};\,S_{\mathrm{lib}})$

RCA can be understood as self-attention, but the values are replaced with
symbols (i.e.,
$\mathrm{Attention}(Q\leftarrow\bm{x},\,K\leftarrow\bm{x},\,V\leftarrow\bm{s})$).
By viewing the attention scores $\alpha_{ij}$ as relations, this has the
effect of producing a relation-centric representation. The rationale is that
in standard self-attention, the attention scores form a type of relation, but
these relations are only used in an intermediate processing step in an
information-retrieval operation. The relations encoded in the attention scores
are entangled with the object-level features, which have much greater
variability. This thinking also motivates the design of RA in the present
work.

RCA can be understood as computing a pairwise relation
$\langle\phi_{q}^{\mathrm{attn}}(x_{i}),\phi_{k}^{\mathrm{attn}}(x_{j})\rangle$
between $x_{i}$ and each $x_{j}$ in the context, and retrieving the symbol
$s_{j}$ associated with the object $x_{j}$ with which the relation is
strongest. That is, RCA treats the relations and the attention scores as the
same thing. By contrast, the attention operation and computation of relations
are disentangled in RA. The attention component is modeled by one set of
query/key maps $\phi_{q}^{\mathrm{attn}},\phi_{k}^{\mathrm{attn}}$ and the
relation component is modeled by another set of query/key maps
$(\phi_{q,\ell}^{\mathrm{rel}},\phi_{k,\ell}^{\mathrm{rel}})_{\ell\in[d_{r}]}$.

The intuitive reason for this choice is that, for many tasks, the optimal
“selection criterion” will be different from the task-relevant relation. For
example, in a language modeling task, you may want to attend to objects on the
basis of proximity and/or syntax while being interested in a relation based on
semantics. Similarly, in a vision task, you may want to attend to objects on
the basis of proximity, while computing a relation across a certain visual
attribute.

In RA, the symbols maintain the role of identifying the sender. But instead of
being the whole message, they are attached to a relation.

###  D.2 Comparison between DAT and the Abstractor

We now briefly discuss the differences in the resulting architectures. [3]
propose an encoder-like module called the Abstractor which consists of
essentially replacing self-attention in an Encoder with relational cross-
attention. That is, it consists of iteratively performing RCA followed by an
MLP. The paper proposes several ways to incorporate this into the broader
Transformer architecture. For example, some of the experiments use a
$\texttt{Encoder}\to\texttt{Abstractor}\to\texttt{Decoder}$ architecture to
perform a sequence-to-sequence task. Here, the output of a standard
Transformer is fed into an Abstractor, and the Decoder cross-attends to the
output of the Abstractor. In another sequence-to-sequence experiment, [3] use
an architecture where the Decoder cross-attends to both the Encoder and the
Abstractor, making use of both sensory and relational information. In
particular, the standard encoder and decoder blocks are the same (focusing on
sensory information), but an additional module is inserted in between with a
relational inductive bias.

By contrast, our approach in this paper is to propose novel encoder and
decoder architectures imbued with two distinct types of attention heads, one
with an inductive bias for sensory information and the other with an inductive
bias for relational information. This has several potential advantages. The
first is versatility and generality. The Abstractor architectures that were
explored in [3] only explicitly support sequence-to-sequence or discriminative
tasks. For example, they do not support autoregressive models like modern
decoder-only language models (e.g., of the form we experiment with in Section
4.3). Moreover, even in sequence-to-sequence tasks, Abstractor architectures
only support relational processing over the input sequence, but they do not
support relational processing over the target sequence (since the decoder does
not have RCA). Another potential advantage of DAT is simplicity. The
Abstractor paper proposes several architectures and configurations for the
Encoder/Abstractor/Decoder modules, introducing several hyperparameters that
are not trivial to choose. Moreover, it is unclear how to interpret this kind
of architecture as the number of layers increases, and the original paper does
not experiment with scaling up the number of layers. The final potential
advantage is increased expressivity. In DAT, the two types of attention heads
exist side by side in each layer. This allows relational attention heads to
attend to the output of the self-attention heads at the previous layer, and
vice-versa. This yields broader representational capacity, and potentially
more interesting behavior as we scale the number of layers.

###  D.3 How would RCA perform in an DAT-style dual head-type architecture?

One question one might ask is: how would an DAT-style dual head-type
architecture perform if we used [3]’s RCA instead of the RA head-type proposed
in this work? We carried out a few ablation experiments to answer this
question.

Figure 12 compares learning curves on the relational games benchmark between
standard DAT (with RA-heads) and a version of DAT with [3]’s RCA heads. We
find that the two models perform similarly, with most differences small enough
to be within the margin of error. This figure depicts the configuration with
asymmetric RA and positional symbols.

Figure 13 depicts the validation loss curves on the Tiny Stories language
modeling task, comparing standard DAT against a version with RCA heads. Here,
we find that our relational attention heads yield better-performing models,
with the RCA-head variant of DAT performing no better than a standard
Transformer with a matching total number of heads.

Figure 12: Learning curves for DAT with RA compared with DAT with RCA on the
relational games benchmark. The performance is similar, with most differences
within the margin of error. Figure 13: Ablation of relational attention type.
The solid line depicts the form of relational attention proposed in this work.
The dotted line depicts RCA as proposed by [3]. We find that our relational
attention mechanism performs better, whereas RCA performs no better than a
Transformer.


# Quantum Pair State Transfer on Isomorphic Branches

Hiranmoy Pal  National Institute of Technology Rourkela, Odisha-769008, India.

###### Abstract

Quantum state transfer plays an important role in quantum information
processing. The evolution of certain pair states in a quantum network with
Heisenberg $XY$ Hamiltonian depends only on the local structure of the
network, and it remains unchanged even if the global structure is altered. All
graphs with high fidelity vertex state transfer may be considered as
isomorphic branches of the graph underlying a large quantum network to exhibit
high fidelity pair state transfer. Among other graphs, one may construct
infinite family of trees admitting perfect pair state transfer.

Introducton.—The continuous-time quantum walk [1] on a quantum spin network
modelled by a graph $G$ with the Heisenberg $XY$ Hamiltonian as described in
[2] is governed by the transition matrix $U(t)=\exp{\left(itA\right)},$ where
$A$ is the adjacency matrix of $G.$ The Laplacian matrix may be considered
instead of $A$ when defining the transition matrix $U(t)$ whenever $XYZ$
interaction model is adopted [3, 4]. However, for a regular graph, the study
of state transfer relative to the adjacency and Laplacian models are
equivalent in a sense that both considerations provide the same information.
The pair state $\frac{1}{\sqrt{2}}\left(|a\rangle-|b\rangle\right)$
corresponds to a pair of vertices $(a,b)$ in $G$ where the characteristic
vectors $|a\rangle$ and $|b\rangle$ are the vertex states of $a$ and $b,$
respectively. Perfect state transfer [5, 2] occurs between $|a\rangle$ and
$|b\rangle$ whenever the fidelity of transfer $|\langle a|U(t)|b\rangle|$
attains its maximum value $1.$ In case of perfect pair state transfer, the
fidelity of the pair states assumes the maximum. A study on Laplacian pair
state transfer can be found in [6]. There we observe that among paths and
cycles, only the paths on three or four vertices, and the cycle on four
vertices admit Laplacian perfect pair state transfer. Here we consider pair
state transfer relative to the adjacency matrix of a graph. In past two
decades, several network topologies having high fidelity vertex state transfer
was analysed, such as the Cartesian powers of the path on two or three
vertices [2], circulant graphs [7, 8], path $P_{n}$ on $n$ vertices [9], etc.
These graphs are considered as isomorphic branches to a larger network, and
then we investigate the existence of pair state transfer as below.

Theorem 1.—Let $u_{1},u_{2},\ldots,u_{k}$ be vertices in a graph $H,$ and
$S_{1},S_{2},\ldots S_{k}$ be sets of vertices (not necessarily disjoint) in a
graph $X_{1}.$ Suppose $X_{1}$ is isomorphic to a graph $X_{2}$ with the
isomorphism $f:X_{1}\to X_{2}.$ Consider the graph $G$ obtained from the
disjoint union of $X_{1},X_{2}$ and $H$ by joining edges from $u_{j}$ to all
vertices in $S_{j}\cup f\left(S_{j}\right).$ If $U(t)$ is the transition
matrix of $G$ then for every vertex $a$ in $X_{1}$

$\displaystyle
U(t)\left(|a\rangle-|f(a)\rangle\right)=(I-P)\sum\limits_{\lambda\in\sigma\left(X_{1}\right)}e^{it\lambda}E_{\lambda}|a\rangle,$

where $P$ is the matrix of the automorphism $\hat{f}$ of $G$ that switches
each vertex of $X_{1}$ to its $f$-isomorphic copy in $X_{2},$ and fixing all
other vertices of $G.$ Assuming $F_{\lambda}$ to be the eigenprojector of
$X_{1}$ corresponding to $\lambda$ in $\sigma\left(X_{1}\right),$ the set of
all eigenvalues $X_{1}$, the matrix $E_{\lambda}$ is given by

$E_{\lambda}=\frac{1}{2}\begin{bmatrix}F_{\lambda}&-F_{\lambda}&0\\\
-F_{\lambda}&F_{\lambda}&0\\\ 0&0&0\end{bmatrix}.$

Proof of Theorem 1.—The spectral decomposition of the transition matrix of $G$
gives

$\displaystyle
U(t)\left(|a\rangle-|f(a)\rangle\right)=\sum\limits_{\lambda\in\sigma\left(G\right)}e^{it\lambda}E_{\lambda}\left(|a\rangle-|f(a)\rangle\right),$
(1)

where $E_{\lambda}$ is the eigenprojector associated to
$\lambda\in\sigma\left(G\right).$ Consider the equitable partition on $G$
formed by the orbits of the automorphism $\hat{f}.$ The eigenvectors of $G$
arising from the equitable partition are constant on the orbits of $\hat{f}$
(see Remark 3.9.6 in [10]). Consequently, the eigenvectors that contribute to
the sum on the right of (1) is of the form $\left[x^{T},-x^{T},0\right]^{T},$
where $x$ is an eigenvector of $X_{1}.$ Since
$|a\rangle-|f(a)\rangle=(I-P)|a\rangle$ and $E_{\lambda}$ commutes with $P,$
we have the desired result.

Suppose the premise of Theorem 1 Holds. If $X_{1}$ has perfect vertex state
transfer between $|a\rangle$ and $|b\rangle$ then $G$ admits perfect pair
state transfer between $\frac{1}{\sqrt{2}}\left(|a\rangle-|f(a)\rangle\right)$
and $\frac{1}{\sqrt{2}}\left(|b\rangle-|f(b)\rangle\right).$ In Figure 1, the
graph $G$ is the path $P_{5},$ where $X_{1}=P_{2}$ and $H$ is an isolated
vertex $u.$ Since perfect vertex state transfer occurs in $X_{1}$ at
$t=\frac{\pi}{2}$ perfect pair state transfer occurs in $P_{5}$ between
$\frac{1}{\sqrt{2}}\left(|a\rangle-|c\rangle\right)$ and
$\frac{1}{\sqrt{2}}\left(|b\rangle-|d\rangle\right)$ at the same time.

$u$$a$$b$$c$$d$$X_{1}$$X_{2}$$H$$u$$a$$b$$c$$d$$X_{1}$$X_{2}$$H$ Figure 1:
The path $P_{5}$ and its perturbation.

In fact, if $u$ is identified to a vertex in any arbitrary graph $H$ (see
Figure 1) then perfect pair state transfer between
$\frac{1}{\sqrt{2}}\left(|a\rangle-|c\rangle\right)$ and
$\frac{1}{\sqrt{2}}\left(|b\rangle-|d\rangle\right)$ remains unchanged.
Considering $H$ to be a tree, we find infinitely many trees with perfect pair
state transfer. In above constructions, we observe that the state transfer of
some pair states behave like a local phenomena. Graphs with high fidelity pair
sate transfer are obtained using Theorem 1 as well. A generalized version of
the previous observations can thus be realised as below.

Theorem 2.—Let $u_{1},u_{2},\ldots,u_{k}$ be vertices in a graph $H,$ and
$S_{1},S_{2},\ldots S_{k}$ be sets of vertices (not necessarily disjoint) in a
graph $X_{1}.$ Suppose $X_{1}$ is isomorphic to $X_{2},X_{3},\ldots,X_{l}$
with the isomorphisms $f_{r}:X_{1}\to X_{r}.$ Consider the graph $G$ obtained
from the disjoint union of $X_{1},X_{2},\ldots,X_{l}$ and $H$ by joining edges
from $u_{j}$ to all vertices in
$S_{j}\cup\left[\bigcup\limits_{r=2}^{l}f_{r}\left(S_{j}\right)\right].$ At
every instance of time, the fidelity of transfer between $|a\rangle$ and
$|b\rangle$ in $X_{1}$ is equal to that of
$\frac{1}{\sqrt{2}}|a\rangle-|f_{r}(a)\rangle$ and
$\frac{1}{\sqrt{2}}|b\rangle-|f_{r}(b)\rangle.$

The path on three vertices $P_{3}$ exhibit perfect vertex state transfer
between the end vertices $a$ and $b$ at $t=\frac{\pi}{\sqrt{2}},$ and the
state of the internal vertex $c$ remains unchanged at the same time.
Therefore, one has perfect pair state transfer in $P_{3}$ between
$\frac{1}{\sqrt{2}}\left(|a\rangle-|c\rangle\right)$ and
$\frac{1}{\sqrt{2}}\left(|b\rangle-|c\rangle\right).$ Likewise graphs having
vertex state transfer also exhibit pair state transfer. The results presented
here gives a non-trivial relationship between vertex state transfer and pair
state transfer, which can be used to design various quantum networks having
pair state transfer.

###### Acknowledgements.

The research is funded by Science and Engineering Research Board (Project:
SRG/2021/000522).

## References

  * Farhi and Gutmann [1998] E. Farhi and S. Gutmann, Quantum computation and decision trees, Phys. Rev. A 58, 915 (1998).
  * Christandl  _et al._ [2004] M. Christandl, N. Datta, A. Ekert, and A. J. Landahl, Perfect state transfer in quantum spin networks, Phys. Rev. Lett. 92, 187902 (2004).
  * Bose  _et al._ [2009] S. Bose, A. Casaccino, S. Mancini, and S. Severini, Communication in XYZ all-to-all quantum networks with a missing link, Int. J. Quantum Inf. 7, 713 (2009).
  * Pal [2022] H. Pal, Laplacian state transfer on graphs with an edge perturbation between twin vertices, Discrete Math. 345, 112872 (2022).
  * Bose [2003] S. Bose, Quantum communication through an unmodulated spin chain, Phys. Rev. Lett. 91, 207901 (2003).
  * Chen and Godsil [2020] Q. Chen and C. Godsil, Pair state transfer, Quantum Inf. Process. 19, 321 (2020).
  * Bašić [2013] M. Bašić, Characterization of quantum circulant networks having perfect state transfer, Quantum Inf. Process. 12, 345 (2013).
  * Pal and Bhattacharjya [2017] H. Pal and B. Bhattacharjya, Pretty good state transfer on circulant graphs, Electron. J. Combin. 24, Paper No. 2.23, 13 (2017).
  * Godsil  _et al._ [2012] C. Godsil, S. Kirkland, S. Severini, and J. Smith, Number-theoretic nature of communication in quantum spin systems, Phys. Rev. Lett. 109, 050502 (2012).
  * Cvetković  _et al._ [2010] D. Cvetković, P. Rowlinson, and S. Simić,  _An introduction to the theory of graph spectra_ , London Mathematical Society Student Texts, Vol. 75 (Cambridge University Press, Cambridge, 2010) pp. xii+364.