Abstractors and relational cross-attention: An inductive bias for explicit relational reasoning in Transformers
Authors: Awni Altabaa
Taylor Whittington Webb
Jonathan D. Cohen
John Lafferty
Venue: ICLR 2024 poster (2024)
Keywords: relational representation learning
attention
transformers
sequence models
abstract representations
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: An extension of Transformers is proposed for explicit representation of relational information.

Abstract: An extension of Transformers is proposed that enables explicit relational reasoning through a novel module called the *Abstractor*. At the core of the Abstractor is a variant of attention called *relational cross-attention*. The approach is motivated by an architectural inductive bias for relational learning that disentangles relational information from object-level features. This enables explicit relational reasoning, supporting abstraction and generalization from limited data. The Abstractor is first evaluated on simple discriminative relational tasks and compared to existing relational architectures. Next, the Abstractor is evaluated on purely relational sequence-to-sequence tasks, where dramatic improvements are seen in sample efficiency compared to standard Transformers. Finally, Abstractors are evaluated on a collection of tasks based on mathematical problem solving, where consistent improvements in performance and sample efficiency are observed.

Published as a conference paper at ICLR 2024
ABSTRACTORS AND RELATIONAL AN INDUCTIVE BIAS FOR EXPLICIT RELATIONAL REASONING IN TRANSFORMERS
Awni Altabaa', Taylor W Jonathan D. Cohen", John Lafferty
1 Yale University, PCCLA, SPrinceton University
ABSTRACT
An extension of Transformers is proposed that enables explicit relational reasoning through a novel module called the Abstractor. At the core of the Abstractor a variant of attention called relational cross-attention. The approach is motivated by an architectural inductive bias for relational learning that disentangles relational information from object-level features. This enables explicit relational reasoning, supporting abstraction and generalization from limited data. The Abstractor is first evaluated simple discriminative relational tasks and compared existing relational architectures. Next, the Abstractor is evaluated on purely relational to-suueneee tasks, where dramatic improvements are seen in sample efficiency compared to standard Transformers. Finally, Abstractors are evaluated collection of tasks based on mathematical problem solving, where consistent improvements in performance and sample efficiency are observed.
INTRODUCTION
The ability to infer and process relations and reason in terms of analogies lies at the heart of human abilities for abstraction and creative thinking (Snow et al., 1984; Holyoak, 2012). This capability is largely separate from our ability to acquire semantic and procedural knowledge through sensory tasks, such as image and audio processing. Modern deep learning systems can often capture this latter type of intelligence through efficient function approximation. However, deep learning has seen limited success with relational and abstract reasoning, which requires identifying novel associations from limited data and generalizing to new domains.
Recognizing the importance of this capability, machine learning research has explored several novel frameworks for relational learning (Graves al., 2014; Pritzel et al., 2017; Santoro et al., 2017; Battaglia et al., 2018; Barrett et al., 2018; Shanahan et al., 2020; Whittington et al., 2020; Webb et al., 2021; Mondal et al., 2023). In this paper, we propose framework that casts relational learning in terms of Transformers. The success of Transformers lies in the use of attentional mechanisms to support richly context-sensitive processing (Vaswani et al., 2017; Wolf al., 2020; Kerg et al., 2020). However, it is clear that Transformers are missing core capabilities required for modeling human thought (Mahowald al., 2023), including an ability support analogy and abstraction. While large language models show surprising ability to complete some analogies (Webb et al., 2023), this ability emerges implicitly only after processing vast amounts of data.
The Transformer architecture has the ability to model relations between objects implicitly through its attention mechanisms. However, we argue in this paper that standard attention produces entangled representations encoding mixture of relational information and object-level features, resulting in suboptimal sample efficiency for learning relations. In this work, we propose an extension of Transformers that enables explicit relational reasoning through novel module called the Abstractor. At the core of the Abstractor is a variant of attention called relational cross-ati cross-attention. Our approach is motivated by architectural inductive bias for relational learning we call the "relational bottleneck," which separates relational information from extraneous object-level features (see Webb et al., 2024, for cognitive science perspective on this idea).
growing body of literature has focused on developing machine learning architectures for relational representation learning. An early example is the Relation Network (Santoro et al., 2017), which

---

Published as a conference paper at ICLR 2024
proposes modeling pairwise relations by applying MLP to the concatenation of object represen- tations. Shanahan et al. (2020) proposed the PrediNet architecture, which learns representations of relations in a manner inspired by predicate logic. Webb et al. (2021)'s ESBN is memory augmented LSTM network that aims to factor representations into "sensory" and 'relational". Another related architecture is CoReINet (Kerg et al., 2022), which reduces relational learning to modeling similarity matrix. More recently, Altabaa and Lafferty (2023) tackled the problem of learning representations of hierarchical relations by formalizing notion of "relational convolution".
The Transformer is a common baseline against which other approaches are compared in this literature. These works show that explicitly relational architectures outperform the Transformer on several synthetic discriminative relational tasks, sometimes by large margins (Shanahan et al., 2020; Webb al., 2021; Kerg al., 2022; Altahaa and Lafferty, 2023). We offer an explanation, arguing that while the Transformer architecture is versatile enough to learn such relational tasks given enough data, does not support relational representation explicitly and thus can suffer in terms of sample-efficiene The Abstractor extends the Transformer framework by introducing an inductive bias for learning representations of relations that are disentangled from extraneous object-level features.
Our experiments first compare the Abstractor to other relational architectures on discriminative relational tasks, finding that the Abstractor both more flexible and achieves superior sample efficiency. We then evaluate whether the Abstractor augment a Transformer to improve relational reasoning by evaluating on synthetic requence-t--sequeee relational tasks, which has so far been unexplored in the literature explicitly relational architectures. Finally, evaluate an Abstractor- supported architecture on set of mathematical problem-solving tasks to evaluate the potential of the idea on tasks more representative of real-world applications. We observe consistent, sometimes dramatic, gains in sample efficiency.
2 RELATIONAL CROSS-ATTENTIC AND THE ABSTRACTOR MODULE
At high level, the primary function of Abstractor is to compute abstract relational features of its inputs.' That is, given a sequence of input objects .1..... .n. the Abstractor learns to model a relation r(:,)) and computes function the set of pairwise relations between objects (r(x;, 25)) At the heart of the Abstractor module is inductive bias we call the relational bottleneck, that disentangles relational information from the features of individual objects.
2.1 MODELING RELATIONS AS INNER PRODUCTS
A "relation function" maps a pair of objects 31, 22 E X' to a vector representing the relation between the two objects. We model pairwise relations as inner products between appropriately encoded (or "filtered') object representations. In particular, we model the pairwise relation function r(.,.) E Rd.
in terms of d, learnable 'query' encoders p/.-.., dg and 'key' encoders d/...... $4.
r(x, y) (z) (k())) (0)))) ( (v)), (x). (1) E Rd..
Modeling relations as inner products (z)( (y)) ensures that the result represents a comparison between the two objects' features. More precisely, inner product relations induce a geometry on the object space X, with notions of distance, angles, and orthogonality. Altabaa and Lafferty (2024) analyzes the approximation properties of inner products of neural networks for relation functions.
Considering all pairwise relations yields relation tensor, R (r(x1,;;)) E Raaxnd,
2.2 RELATIONAL CROSS-ATTENTION
The core operation in Transformer is attention. For an input sequence X (xI,..., Raxd self-attention transforms the sequence via, X' + Softmax (o, )ox (X), where B 0 &y are functions applied independently to each object in the sequence. Note that R :=
In this paper, we will use the name 'Abstractor' to refer to hoth the module and to model architectures which contain the Abstractor module main component.
2

---

Published as a conference paper at ICLR 2024
(x(X))x(X)) is relation matrix in the sense defined above. Self-attention admits an interpretation as a form of 'neural message- saassageppssing" (Gilmer et al., 2017) as follows
x MessagePassingg (1)))(z)), (.)) Rey))
Σ (x))
(2)
where Mj-xi (rg) Ray) is the message from object j to object i, encoding the sender's features (2yy and the relation between the two objects Rij = (dy (x2,) (x(r)) Here, R Sofumax (R) is the
relation matrix. Hence, the processed representation obtained by self- attention an entangled mixture of relational information and object-level features.
Our goal is to learn relational representations that are abstracted away from object-level features in order to achieve more sample-efficien learning and improved generalization in relational reasoning. This is not naturally supported by the entangled representations produced by standard self-attention. We achieve this via a simple modification of attention-w replace the values &u (2;) with vectors that identify objects, but do encode their features. We call those vectors symbols. Hence, the message sent from object j to object is now mj++ = (sj, Ray), the relation between the two objects, together with the symbol identifying the sender object.
A MessagePassingg (1 (u))),m)))
Σ Raysy.
Symbols act abstract references objects, akin to pointers in traditional symbolic architectures. They do not encode information about the contents or features of the objects, but rather refer to objects. This results in improved sample efficiency and generalization by restricting the search space of computations, and allowing abstraction through shared symbols. We call the vectors {*t} 'symbols' in the same sense that we call 'r' a symbol an equation like y x²2they reference an object with an unspecified value. Suppose for now that each object z; is assigned a symbol in manner that satisfies this property. We will discuss -asbol-assignment mechanisms in the next subsection.
This modification yields variant of attention that we call relational cross-attention (RCA), given by
- drel (Xx(X) (x(x))) S.
where S (S),.... ..s....,,,,) 8n) are the symbols, rel is the relation activation function, and og. Φs correspond to the query and key transformations. When the relation activation function orel is softmax, this corresponds to Attention(Q <- X, K + X, V + S). In contrast, self-attention corresponds to Attention(Q + X, K X, V + X), mixing relational information with object-level features.
We observe in our experiments that allowing Grel to be a configurable hyperparameter can lead to performance benefits in some tasks. Softmax has the effect of normalizing the relation between a pair of objects (i based on the strength of i's relations with the other objects in the sequence. In some tasks this is useful. In other tasks, this may mask relevant information, and element- wise activations (e.g., tanh, sigmoid, or linear) may be more appropriate.
Relational cross-attention implements type of information bottleneck, that call the "relational bottleneck", wherein the resultant representation encodes only relational information about the object sequence (Figure 1) and does not encode information about the features of individual objects. This enables a branch of the model to focus purely on modeling the relations between objects, yielding greater sample efficiency in tasks that rely on relational reasoning.
Multi-dimensional relations can be modeled through multi-head relational cross- cross-attention. In our experiments, 0%,- of are linear maps Wi. W and multi-head relational cross- crossnattention given by
RelationallCrosstteenniooo X, S) = concat ".,,...,A,,, W.. ") sw:
(5)
where A = Grel
2.3 SYMBOL ASSIGNMENT MECHANISMS
Abstraction relies on the assignment of symbols to individual objects without directly encoding their features. We propose three different mechanisms for assigning symbols objects.
3

---

Published as a conference paper at ICLR 2024
Encoder states
Relational Bottleneck
ststract
F
E
Fg
F
Myy
Relation tensor (o,(4))dd2t)))
Ry = (4 (x)))
Relation
A
A2
A
An
12
X1
X2
X
Xn
S1
$2
Sq
Input sequence
Iaput sequence
Symbols
(a) E' + Sel[Attention(X)
(b) + RelttonnlCCrssAtteetiinkk. S)
Figure 1: Comparison of relational cross-attention with self-attention. Red represents object-level features, blue represents relational features, and purple represents mixed representations. Relational cross-attention computes relational information disentangled from the features of individual objects.
Positional symbols. The simplest symbol assignment mechanism is to assign symbols to objects sequentially based on the order they appear in the sequence. That is, we maintain a library of symbols Sib = (81,-... Snax, len) TRaax__enx Lew._leexdde Amodel and assign the symbol s; to the i-th object. The symbol library Slbb can either learned parameters of the model fixed positional embeddings.
Posation-relative symbols. Similar to relative positional embeddings (Shaw et al., 2018; Kazemnejad et al., 2023), we can compute relational cross-altention with position-relative symbols via A + E, Raj where the symbol library is Sib = (... .-..,8-1, 1,,0,, 80, .1....).
jo
Symbolic attention. In this case we learn a library of symbols Sub ($1..... Sn.) E Xdmoded together with associated "binding vectors" Bm, = (bi...., bn where is the number of symbols the library. Through an attention mechanism, symbols are bound to objects n; based on their relations to the symbol binding vectors. A multi-head variant of symbolic attention is naturally defined by concatenating the symbols retrieved for each head. Formally,
SymmoliccAtentionnXX X) concat (5 (S0))...,,S S(Hu))
(6)
((XW)) !) Shb.
S Softmax
In terms of representational capacity, this is similar to cross-attending to the symbol parameters Süb-- Note that symbolic attention weakens the relational bottleneck since object-level features are used to retrieve a symbol for each object. However, the symbols are shared across objects and sequences, and the dependence only with respect low-dimensional projection of the object-level features, which we may think of encoding the object's "syntactic role." Encoding such information in the symbols allows identifying objects by their role, rather than merely their position or relative position.
The choice of symbol assignment mechanism determines the way in which relational information encoded the abstract states. For example, with positional symbols, A, encodes the relations between object and each object identifying by its position (or relative position in the case of risition-relative symbols). contrast, with symbolic attention, each object j is identified by its "syntactic role." as determined by its relation to the binding vectors.
2.4 THE ABSTRACTOR MODULE
We now describe the Abstractor module. Like the Encoder in Transformer, this module that processes an input sequence of objects X = (Iz,..., 2I,) producing processed sequence of objects A = (A)...,, An) that represents features of the input sequence. In an Encoder, the output objects represent mix of object-level features and relational information. In an Abstractor, the output objects are abstract states that represent relational information, abstracted away from the features of individual objects. The core operation in Abstractor module is relational cross- attention. Mirroring an Encoder, an Abstractor module can comprise several layers, each composed of relational cross-attention followed by feedforward network. Optionally, residual connections
4

---

Published as a conference paper at ICLR 2024
and nnyer-normalization can be applied suggested by Vaswani al. (2017)². The algorithmic description is presented Algorithm 1. In Appendix B, we theoretically study the approximation properties of the Abstractor module, showing that it is universal approximator of relation functions.
Algorithm I: Abstractor module
:object sequence: X (II1...;Znn) Rn) Raxd
A SymbolicAttention(X)
// or, A Sim, for positional symbols
for / < 1 to L do
40) <- AelationnlCCoosAtttetiooo ((-)))
400 <- A + 1t- // residual connection (optional)
A0) + LayerNorm (AO)) // (optional; can also be done pre-RCA)
A0) <- FeedForward (A(1)
end
Output: ALL)
The hyperparameters of an Abstractor module include the number of layers L, the relation dimension d, (i.e., number of heads), the projection dimension dę (i.e., key dimension), the relation activation function Tret. and the model dimension dmoode-- The learnable parameters, each layer, the projection matrices W W E RduoddeX4d Xd,, i [dr], the symbol library Sjb. and the parameters of the feedforward network. In our experiments. we use 2-layer feedforward network with hidden dimension dæ and ReLU activation. The implementation in the publicly available code adds few additional hyperparameters, including whether restrict the learned relations to be symmetric (via W; = W and whether to apply self-attention after relational cross- attentot (which enables exchange of relational information between abstract states).
3 ABSTRACTOR ARCHITECTURES
Whereas a Transformer Encoder performs "general-purposs"" processing, extracting representations of both object-level and relational information, an Abstractor module is more specialized and produces more enriched relational representations. An Abstractor module can be integrated into a broader transformer-bart based architecture, for enhanced relational processing.
To facilitate the discussion of different architectures, we distinguish between types of tasks. In purely relational prediction task, there exists a sufficient statistic of the imput which is purely relational and encodes all the information that is relevant for predicting the target. The experiments of (Webb et al., 2021; Kerg et al., 2022) are examples of purely relational discriminative tasks. We consider discriminative relational tasks in Section 4.1. An example of a purely relational sequence-to-xequencc task is the object-sorting task described in Section 4.2. Many real-world tasks, however, are not purely relational. partially- prediction task, the relational information is important but is not sufficient for predicting the target. The math problem-solving experiments in Section 4.3 are Natural language understanding is also an example of a partially-relational task.
partially-
The way that an Abstractor module is integrated into a broader model architecture should be informed by the underlying prediction task. Figure depicts several Abstractor architectures each with different configurations. Architecture (a) depicts configuration in which the Abstractor processes the relational features in the input, and the decoder attends to the abstract states A. Architecture (b) depicts configuration in which the input objects are first processed by Encoder, followed by an Abstractor for relational processing, and the decoder again attends to the abstract states. These architectures would be appropriate for purely relational tasks since the decoder attends only the relational representations' in the abstract states. Architectures (c) and (d), in which information can also pass directly from the encoder to the decoder, would be more appropriate for more general tasks that are only partially relational. For example, in architecture (c), the model branches into two parallel processing streams in which an Encoder performs general-purpose purpose processing and an Abstractor
2 nyye---orraiiziiionn can also he applied before relational cross-attention, in the pre-INN Trans- former (Xiong et al., 2020).
5

---

Published as a conference paper at ICLR 2024
(8)
Abstractor
(b)
X
Encoder
rastoo
Decader
X
Abstractor
[d)
Abstractor
Decoder
(e)
Abstractor
Figure 2: Examples of Abstractor-based model architectures.
performs more specialized processing of relational information. The decoder attends both the encoder states and the abstract states. These architectures use the "multi-attention decoder" described in Appendix A. Finally, architecture (e) depicts composition of Abstractors, wherein the abstract states produced by one Abstractor module are used as input to another Abstractor. This results in computing "higher-order relational information (i.e., relations on relations).
EXPERIMENTS
4.1 DISCRIMINATIVE RELATIONAL. TASKS
modeling asymmetric relations. We generate N 64 "random objects" rep- resented by Gaussian vectors, o ~ (N, E R and associate an order relation to them 01 02 < < oN- Note that < is anfi-symmetric. Of the N 4096 possible pairs 05). 15% are held out as a validation set (for early stopping) and 35% as a test set. We evaluate learning curves by training on the remaining 50% and computing accuracy on the set size). Note that the models are evaluated on pairs they have never seen. Thus, the models will need generalize based the transitivity of the relation.
Order relations:
(10 trials for each training
We compare five models: an Abstractor, standard (symmetric) CoRelNet, an asymmetric variant of CoReINet, PrediNet, and an MLP. The MLP, which is non-relational architecture, is completely unable to learn the task. Among the relational architectures, we observe that standard CoRelNet is also completely unable to learn the task, whereas the Abstractor and asymmetric CoReINet learn the transitive relation (Figure 4a). PrediNet has limited success in learning the task. This can be explained by the observation that symmetric inner products (e.g., in standard CoReINet) don't 't have the representational capacity to model asymmetric relations, whereas the asymmetric inner products with different learned left and right encoders do.
SET: modeling mutti-dimensional relations. SET is a cognitive task in which players are presented with sequence of cards, each of which contains figures that vary along four dimensions (color, number, pattern, and shape) and they must find triplets of cards that obey deceptively simple rule: along each dimension, cards "set" must either have the same value or three unique values (Figure 3). In this experiment, the task classify triplets of card images as "set" or not.
$
0
0
S
Again, compare an Abstractor, CoReINet, PrediNet, and an MLP. The shared architecture CNN + {-} + Flatten - Dense, where {.}
Figure 3: The SET game
one of the aforementioned modules. The CNN embedder is obtained through a pre-training task. We report learning curves in Figure 4b (10
trials per training set size). We find that the Abstractor model significantly outperforms the other baselines. We attribute this to its relational inductive biases and its ability to model multi dimensional relations. In this task, there exist four different relations (one for each attribute) that needed to

---

Published as a conference paper at ICLR 2024
Learning Agymmetnic Relation
SET Classification
Comparison to symbolic baseline SET tack
1.0
SET Clas Acsmaac
100
0.5
Galaltee Aeemmai
Canaite Soltman
0.8
112
1000
2003
1000
Traiiing
Sim
(a) The < relation can be learned with asymmetric but not sym- metric inner products.
(b) The Abstractor's 's ability model multi-dimensional rela- tions enables it to solve SET.
(c) Comparison Abstractor trained on card images and MLP with hand-encoded relations.
Figure 4: Experiments on discriminative relational tasks and comparison CoRelNet.
determine whether a triple of cards forms a set. We hypothesize that the ability to model relations as multi-dimensional is also part of the reason that the Abstractor is sample efficient in learning the order relation the previous experiment- even though the underlying relation is "one-dimensional" having multi-dimensional representation enables greater robustness and multiple avenues towards good solution during optimization.
SET (continued): comparison to "neuro-symbolic" model. In Figure 4c, to evaluate the quality of the representations produced by Abstractors, we compare an Abstractor-basee model to "neuro- symbolic" model which receives as input binary representation of the four relevant relations. We train 1-head Abstractors separately for each of the four attributes learn "same/different" relations, where the task is to decide if an input pair of cards is the same or different for that attribute. We then use the Wa and We parameters learned for these relations initialize the relations in a multi-head Abstractor. The Abstractor is then trained on dataset of triples of cards, half of which form "set".
This is compared a baseline neuro-symbolic model where, instead images, the input a vector with 12 bits, explicitly encoding the relations. A two-layer MLP is then trained to decide if the triple forms a "set". The MLP using the symbolic representation represents an upper bound on the performance achievable by any neural network model. This comparison shows that the relational representations learned by an Abstractor result in sample efficiency that is not far from that obtained with pre-specifien symbolic encodings of the relevant relations.
4.2 OBJECT-SORTING: PURELY RELATIONAL
TASKS
In the following set of experiments, we consider sequence-tor -oueencett--squence tasks which are purely relational, and compare an Abstractor-supporee model to standard Transformer. We consider the task of sorting sequences of random objects. This task "purely relational" the sense that there exists relation (order) which is sufficient statistic for solving the task-the features of individual objects beyond this relation are extraneous. This is a more controlled setting that tests the hypothesis that the inductive biases of the Abstractor confer benefits in modeling relations. The experiments in the present section demonstrate that the Abstractor enables a dramatic improvement in sample-efficiency soguence-to-sequeeee relational tasks.
Random Object Sorting
Random Object Sorting Generalization
Sorting Accuracy
Sorting Accuracy
0.50
0.5
Aberartor
0.25
aetor Abluion Mode
rzaieed
Trandformer
0.00
Transforar
Presrained
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Training Set Size
Training Set Size
(a) Leariing curves on sorting sequences of ran- dom objects. The abstractor is dramatically more eamplc-efficient
(b) Learning curves with and without pre- pre-training similar sorting task. The Abstractor benefits from pre-training.
Figure 5: Experiments on object-sorting, a purely relational sequence-to-sequencc task.
7

---

Published as a conference paper at ICLR 2024
Task: pol pol.yomiia _ee__ppaad_ddddd__________
Task: 1igebra__ineea_
Question: Expand (2*x + * 1).
Question: Salve for 2: 5*z + - 9. Answer: 7/5
Answer: 2 + x - 3
Figure 6: Examples of input/target sequences from the math problem-solving dataset.
Superior sample-efficiency on relational seq2seq tasks. We generate random objects in the following way. First, we generate two sets of random attributes A = A 02, d3, az}. j.j... a, ~ (0,, e R and B = {61,... (bi,...... b12), b; (0, E R.. To each set of attributes, we associate the strict ordering relation 01 < 0g x d3 x a4 and bj x b2 < ................................ < b12, respectively. Our random objects are formed by the Cartesian product of these two attributes O = A B, yielding N = 48 objects. We associate with O the following strict ordering relation: (a,, by) < (ax; br) if a; < O& or if a; - and bj- < br Given a set of objects in O, the task is sort it according to x. The input sequences are randomly permuted sequences of 10 objects in O and the target sequences are the indices of the object sequences in sorted order (i.e.. 'argsort') ). The training data are sampled uniformly from the set of length- 10 sequences in O. We also generate non-overlapping validation and testing datasets.
We evaluate learning curves on an Abstractor, a standard Transformer, and an "Ablation" model (10 trials for each training set size). The Abstractor uses architecture (b) in Figure 2 with learned positional symbols. The Encooderto--Abstraator interface uses relational cross-attention and the interface uses standard cross-attention. The Ablation Model tests the effects of relational eross-attention in the Abstractor model-it is architecturally identical to the Abstractor model with the crucial exception that the Encoder-to-AAstracce cross-attention. The hyperparameters of the models are chosen so that the parameter counts are similar (details in Appendix C). We find that the Abstractor is dramatically more sample-efficient than both the standard Transformer and the Ablation model (Figure 5a).
Abstractor-to-Decoddd
interface instead uses standard
Ability to generalize to similar tasks. We also used the object-sorting task and the dataset generated as described above to test the Abstractor's ability to generalize from similar relational tasks through pre-training. The main task uses the same dataset described above. The pre-training task involves the same object set O but with a modified order relation. The ordering in attribute A is randomly permuted, while the ordering in attribute B is kept the same. We pre-train an Abstractor and a Transformer on the pre-training task and then, using those learned weights for initialization, evaluate learning curves on the original task. Since the Transformer requires more training samples to learn the object-sorting task, we use a pre-trainn pre-training set size of 3 000, chosen to be large enough for the Transformer to learn the pre-training task.
This experiment assesses the models' ability to generalize relations learned on one task to a new task. Figure 5b shows the learning curves for each model with and without pre-training. We observe that when the Abstractor is pre-trained, its learning curve on the object-sorting task is significantly accelerated, whereas the Transformer does not benefit from pre-training. We attribute this to the fact that the Transformer's learned representations are entangled with extraneous object-level features, which prevents generalization; by contrast, the Abstractor's disentangled relational representations can be more easily mapped to the new task.
4.3 MATH PROBLEM-SOLVING: PARTIALIY-RELATIONAA SEQUENCE- TO- TO-SEQUENCE TASKS
The object-sorting experiments in the previous section are "purely relational" in the sense that the set of pairwise - order relations is a sufficient statistic for solving the task. In a general sequence-to- sequence task, however, there may not be relation that is sufficient statistic. Nonetheless, relational reasoning may still be crucial for solving the task, and the enhanced relational reasoning capabilities architectures
of the Abstractor may enable performance improvements. The ppartiall--relationn described in Section 3 enable branch of the model to focus on relational reasoning while another branch performs general processing involving object-level attributes. In this section, we evaluate such an Abstractor model (using architecture (d) of Figure 2) on set of math problem-solving tasks based on the dataset proposed by Saxton et al. (2019).
The dataset consists of several math problem- solvenng tasks, with each task having set of question- answer pairs. The tasks include solving equations, expanding products of polynomials, differentiating functions, predicting the next term in sequence, etc. A sample of question-answer pairs is displayed in Figure 6. The overall dataset contains 2 x 10 training examples and 10 validation examples per task. Questions have maximum length of 160 characters and answers have maximum length
8

---

Published as a conference paper at ICLR 2024
differits
algebra__eqqeenceeeetttt
Ligebra_.Jincaro
Accuracy (val)
TF Accuracy (val)
Accuracy (val)
0.8
0.9
0.8
0.6
20
40
20
40
0
20
40
epoch
epoch
epoch
palynomials expand
polynomials add
TF Accuracy (val)
TF Accuracy (val)
Model
0.80
0.85
Transformer
Transformer
0.75
Abstractor (pos, symbols)
0.80
Abstractor (symbolic attn)
0.70
epoch
opoch
Figure 7: Training curves comparing an Abstractor-based architecture standard Transformer on mathematics problem-solving tasks.
of 30 characters. We use character-level encoding with a common alphabet size 95 (including upper/lower case characters, digits, punctuation, and special start/end/pad tokens).
We compare an Abstractor model using architecture (d) in Figure 2 to a standard Transformer. We evaluate Abstractor with positional symbols and Abstractor with symbolic attention (see Section 2.3). Since the Abstractor-based models with architecture (d) have an Abstractor module addition Encoder and Decoder, we compare against two versions of the Transformer in order to control for parameter count. In the first, the Encoder and Decoder have identical hyperparameters the Abstractor model. In the second, the feedforward network such that the overall parameter count is approximately the same as for the Abstractor model. We refer to the first model as "Transformer" and the second as "Transformer4" in the figures. The precise architectural details and hyperparameters are described in Appendix C.
increase the model dimension and hidden layer size of
We evaluate the three models on five subtasks: differentiating functions (calculus); predicting the next term in sequence (algebra); solving linear equation (algebra); expanding polynomials; and adding polynomials. For each, we train on training split and track teacher-foreing accuracy (excluding null characters) the validation split. For each combination of model and task, repeat the experiment five times and report error bars as twice the standard error of the mean.
Figure 7 shows the validation tcacher-forcing accuracy during the course of training, which use as a proxy for sample efficiency. We observe an improvement in accuracy compared to both "Transformer' and "Transformert" across all tasks. The larger Transformer tends perform better than the smaller Transformer, but the Abstractor-based model consistently outperforms both, with the symbolic attention mechanism showing the greatest improvement. This indicates that the performance improvement stems from the architectural modification and inductive bias. We conjecture that "partially-relation Abstractor architecture (e.g.. architecture (d)) implements two branches of information processing. The Encoder performs more general-purpose purppose processing of the input sequence, while the Abstractor performs more specialized relational processing. The Decoder then has access to both representations, enabling it to perform the task more effectively.
5 CONCLUSION
In this work, we propose variant of attention that produces representations of relational information disentangled from object-level attributes. This leads to the development of the Abstractor module, which fits naturally into the powerful framework of Transformers. Through a series of experiments, we demonstrate the potential of this new framework to achieve gains in sample efficiency in both purely relational tasks and more general sequence modeling tasks. This work opens up several avenues for future research, including better understanding of the strengths and weaknesses of different architectural variants, work towards a more streamlined scalable architecture, and exploring the framework's use in increasingly complex real-world problems.

---

Published as a conference paper at ICLR 2024
CODE AND REPRODUCIBILITY
Code, detailed experimental logs, and instructions for reproducing our experimental results are available at: https: /ttl:://ggtttttttte wttp://...co../.to////ttttttaeea
ACKNOWLEDGMENTS
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). JL was supported in part by NSF grant CCF-1839308 and NSF grant DMS-2015397. JDC was supported by a Vannevar Bush Faculty Fellowship.
REFERENCES
Altabaa, Awni and John Lafferty (2023). "Relational Convolutional Networks: A Framework for Learning Representations of Hierarchical Relations". arXiv: 2310. 03240 [cs]. Altabaa, Awni and John Lafferty
(2024). "Approximation of relation functions and attention mecha- nisms". In: arXiv: 2402. 08856 [cs LG].
Barrett, David, Felix Hill, Adam Santoro, Ari Morcos, and Timothy Lillicrap (2018). "Measuring abstract reasoning in neural networks". In: International conference on machine learning. PMLR. Battaglia, Peter W. et al. (2018). "Relational inductive biases, deep learning, and graph networks". arXiv: 1806. 01 21 LG). Cybenko, George (1989). "Approximation by superpositions of sigmoidal function". In: Mathemat- ics of Control, Signals, and Systems.
Gilmer, Justin, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl (2017).
"Neural message passing for quantum chemistry". In: International conference on machine learning. PMLR.
Graves, Alex, Greg Wayne, and Ivo Danihelka (2014). "Neural Turing Machines". arXiv: 1410.
5401 [cs NE].
Holyoak, K. J. (2012). "Analogy and relational reasoning". In: The Oxford Handbook of Thinking and Reasoning. Ed. by K. Holyoak and R. G. Morrison. New York: Oxford University Press. Kazemnejad, Amirhossein, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy (2023). "The Impact of Positional Encoding on Length Generalization in Transformers". arXiv: 2305 9 66 [cs].
Kerg, Giancarlo, Bhargav Kanuparthi, Anirudh Goyal, Kyle Goyette, Yoshua Bengio, and Guillaume Lajoie (2020). "Untangling tradeoff's between recurrence and self-attention in artificial neural networks". In: Advances in Neural Information Processing Systems.
Kerg, Giancarlo, Sarthak Mittal, David Rolnick, Yoshua Bengio, Blake Richards, and Guillaume Lajoie (2022). "On Neural Architecture Inductive Biases for Relational Tasks". arXiv: 2206
05056 [cs . .N].
Mahowald, Kyle, Anna A. Ivanova, Idan A. Blank, Nancy Kanwisher, Joshua B. Tenenbaum, and
Evelina Fedorenko (2023). "Dissociating language and thought in large language models". arXiv: 0001.066220 [cs. [L]
Mondal, Shanka Subhra, Taylor Whittington Webb, and Jonathan Cohen (2023). "Learning to reason visual objects". In: International Conference on Learning Representations.
Pritzel, Alexander, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell (2017). "Neural episodic control". In:
International conference on machine learning. PMLR.
Santoro, Adam, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Tim Lillicrap (2017). "A simple neural network module for relational reasoning". Advances neural information processing systems.
Saxton, David, Edward Grefenstette, Felix Hill, and Pushmeet Kohli (2019). "A "Analysing Mathematical Reasoning Abilities of Neural Models". arXiv: 1904. 01557 [cs. LG]. Shanahan, Murray, Kyriacos Nikiforou, Antonia Creswell, Christos Kaplanis, David Barrett, and Marta Gameeo (2020). "An Explicitly Relational Neural Network Architecture". In: Proceedings of the 37th International Conference Machine Learning. Shaw, Peter, Jakob Uszkoreit, and Ashish Vaswani (2018). "Self-Attentie with Relative Position Representations" arXiv: 8033... 021 55 [cs CL.).
10

---

Published as a conference paper at ICLR 2024
Snow, R. E.,, P. C. Kyllonen, and B. Marshalek (1984). "The topography of ability and learning correlations", In: Advances in the psychology of human intelligence. Ed. by R. J. Stembberg. Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin (2017). "Attention is all you need". In: Advances in neural information processing systems.
Webb, Taylor, Keith J Holyoak, and Hongjing Lu (2023). "Emergent analogical reasoning in large language models". In: Nature Human Behaviour.
Webb, Taylor W., Steven M. Frankland, Awni Altabaa, Kamesh Krishnamurthy, Declan Campbell, Jacob Russin, Randall O'Reilly, John Lafferty, and Jonathan D. Cohen (2024). "The Relational Bottleneck as an Inductive Bias for Efficient Abstraction". arXiv: 2309 06629 [ca, AI Webb, Taylor Whittington, Ishan Sinha, and Jonathan Cohen (2021). "Emergent Symbols through Binding in External Memory". In: International Conference on Learning Representations.
Whittington, James C.R., Timothy H. Muller, Shirley Mark, Guifen Chen, Caswell Barry. Neil Burgess, and Timothy E.J. Behrens (2020). "The Miiman-Ecchenbaum Machine: Unifying Space and Relational Memory through Generalization in the Hippocampal Formation". In: Cell.
Wolf, Thomas, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. (2020). "Transformers: State-of- the-art natural language processing". In: Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations.
Xiong, Ruibin, Yunchang Yang. Di He, Kai Zheng. Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu (2020). "On layer normalization in the transformer architecture". In: International Conference on Machine Learning. PMLR.
Zhou, Shuheng, John Lafferty, and Larry Wasserman (2009). "Compressed and Privacy-Sensitive Sparse Regression". In: IEEE Trans. Information Theory.

---

Published as a conference paper at ICLR 2024
A
MULTI-ATTENTION DECODER
Algorithm 2: Multi-Attention Decoder
Input: Target sequence: = (Yo,.... U1, -1/-
Context sequences: x (x) ,...,,I, )1),, i 1....,,
(6)
D y for <- 1 to do
D +- CausalSelEAttention (D(I-))
residual connection and layer-normalzzaiio:
for i + 1 to K do
D +- CrossAttention (D
residual connection and layer-normalzattion
end
D + FeedForward (D())
end
Output: (
B UNIVERSAL APPROXIMATION OF RELATION FUNCTIONS
In this section, characterize the function class of the Abstractor and relational cross-attention. We will show that, in an appropriate sense, a single Abstractor layer can compute a sequence of abstract states A = -A:..... A,) such that A, approximates arbitrary function of object i's relations with the other objects in the input. Recall that relational cross-attention takes the form
A ree (xx(X) S,
where og. Øk : X + Ra are multi-layer perceptrons, and S is the matrix of symbols. In this analysis, we consider dell x + to be the linear activation function. Further, we consider positional symbols. For simplicity, we will assume the single-head variant of relational cross-attention. The function class results derived here would of course carry to the multi-head case, where each 'head' can approximate a function in the function class.
Recall that an Abstractor module comprises several layers, each composed of relational cross-attention and feedforward layer. Hence, the overall operation in one layer of an Abstractor module is
AbstractorLayer([ (X) FeedForward (X))x
s). T)
We will characterize the function class AbstractorLayer X Roxd induced by varying the parameters of &gg Φk, Feed and S.
Remark 1. In Equation (4), we formulate relational gross-attention with the maps Øq. 0b as linear projections, whereas thg dhe are multi-layer perceptrons in Equation (7). However, recall that the input to each Abstractor layer is the output of the preceding layer which ends with a multi-layer perceptron. The function class of a multi-lan layer perceptron followed by two different linear projections is the same the function class of two different multi-layer perceptrons. We focus Equation (7) for ease of presentation.
The following result shows that single-layer abstractor returns a sequence of abstract states which can approximate an arbitrary function of each object's 's relations with the other object in the input. The result is based the analysis in (Altabaa and Lafferty, 2024) which characterizes the function class of inner products of neural networks.
Theorem 1. Suppose is a compact Euclidean space. Letr : X - R be any continuous relation function, and R +R R any continuous function. Consider the function g X paxd defined by
(#1,-..,,x)) (J(RR),...J((A
12

---

Published as a conference paper at ICLR 2024
where R² (r(x,, (E) is the vector of 2; e.'s relations with the other objects in the inpus, according r. Then, for any O, there exists MLPs Qy. @k FeedForward and choice of symbols S such that the Abstractor layer approximates in the sup-norm
Fn) - AbstractorTayer((1.. Tn)ll S Lebesgue whmost-everywhere.
Proof. Let (Ai...., An) = Layyrai .n) be the abstract states returned by the abstractor layer. For each E [ni]. the abstract state A; takes the form,
A; FoedForwwrd Σ
(0,(2)),
where ($i...., Sn,) are the symbols assigned each object. Let the symbol dimension be and let ei the canonical basis vectors. Then, A, FeedForward(RR)) where R (Te) (x(()) E IR".. By (Altabaa and Lafferty, 2024, Theorem 3.1), there exists MI.ss Øg ox such that their inner product approximates any continuous relation function. In particular, for any eι O, there exists &y: Φu such that
rxx, v) (x), dx(1)) E1, Lebesgue almost-every 2, X
(8)
Let &g. 02 be such MLPs where is to be determined later. Note that (Altabaa and Lafferty, 2024) gives bound on the number of neurons needed terms of the continuity of r and the dimension of X. Similarly, by the universal approximation property of MLPs (e.g., Cybenko, 1989), f R" - Rd can be approximated by FeedForward uniformly in the sup-norm. That is, for any E2 there exists FoodForwwrd such that
sup [z) FeedForward(()]le S2.
=ER"
Let [R]y = xz) and [R] (x;), on (xz))) Then, the difference g(.1,...,,N) AbstractorLayer((1....... given by
[eedForvard((n))...... PeedForward(RR.))
(10)
Note that R, is close to R. by Equation (8)
R2 Ri
max (1) - - r1)
sE(n]
(11)
S E1
Lebesgue almost-everywhere.
Now consider the (i, ji,j)-th element of the difference in Equation (10)
FeedForward, (R,, f. (R,,) FeedForward f.(R,)) + -
The first term is bounded by due to Equation (9). Let E2 E/2. Recall that f - Rd is continuous, and hence for all 0 there exists 8y(e) 0 such that || 2211xx S 8(E) [lʃ(z)) [=2)) < Letting = 8(e/2) implies that the second term is bounded by 8/2 Lebesgue almost-everywhhre due to Equation (11).
This holds for all i, j. which completes the proof.
C EXPERIMENTAL DETAILS
In this section, we give further experimental details including architectures, hyperparameters, and implementation details. All models and experiments implemented in Tensorflow. The code is publicly available on the project repo along with detailed experimental logs and instructions for reproducing our results.
13

---

Published as a conference paper at ICLR 2024
C.I DISCRIMINATIVE TASKS (SECTION 4.1)
C.I.1 PAIRWISE ORDER
Each model in this experiment has the following form input + {-} + flatten MLP, where {-) is one of the modules below and MLP is an MLP composed one hidden layer with 32 neurons and ReLU activation.
Abstractor architecture. The Abstractor module used the following hyperparameters: number of layers L 1. relation dimension d, = 4, symbol dimension d, 64, projection (key) dimension dy = 16, feedforward hidden dimension dæ = 64, relation activation function Tee = sigmoid. No layer normalization residual connection. We use positional symbols as the symbol assignment mechanism, which are learned parameters of the model. The output of the Abstractor module is flattened and passed the MLP.
CoRelNet architecture. CoReINet has no hyperparameters. Given a sequence of objects, X ,....,,TTn,,, standard CoReINet (Kerg et al., 2022) simply computes the inner product and takes the Softmax. We also add learnable linear map. W E Rdxd_ Hence, R. = Softmax( R = [(Wx, x,, Waylley. The CoReINet architecture flattens R and passes it an MLP to produce the output. The asymmetric variant of CoRelNet is given by R Softmax (R), R = [ 24, W2 ky>lig> where W1. W2 E dddd are learnable matrices.
PrediNet architecture. We based our implementation of PrediNet (Shanahan et al., 2020) on the authors' publicly available code. We used the following hyperparameters: using heads, and 16 relations, a key dimension of 4 (see the original paper for the meaning of these hyperparameters). The output the PrediNet module is flattened and passed to the MLP.
MIP. The embeddings of the objects are concatenated and passed directly to an MLP. The MLP has two hidden layers each with 32 neurons and ReLU activation.
Training/Evaluation. We use the crossentropy loss and the Adam optimizer with learning rate of 10-22 31 = 0.9, 32 = 0.999, = = 10-7, We use a batch size of 64. We train for 100 epochs and restore the best model according validation loss. We evaluate on the test set.
C.1.2 SET
The card images are RGB images dimension 70 x 50 x 3. A CNN embedder processes the images separately and produces embeddings of dimension d = 64 for each card. The CNN is trained to predict the four attributes of each card and then an embedding for each card is obtained from an intermediate layer (i.e., the parameters of the CNN then frozen). Recall that the common architecture is CNN Embedder + {-) > Flatten + Dense (2), where {.} is an Abstractor, CoRelNet, PrediNet, or an MLP. We tested against the standard version of CoReINet, but found that it did not learn anything. We iterated over the hyperparameters and architecture to improve its performance. We found that removing the softmax activation in CoRelNet improved performance bit. We describe hyperparameters below.
Common embedder's architecture The architecture is given by Conv2D - MaxPoo12D + Conv2D > MaxPool 1 2D -> Fl Flatten Dense ( (64, 'relu' ) - Dense (64, 'relu') -} Dense (2). The embedding extracted from the penultimate layer. The CNN trained to predict the four attributes of each card until it reaches perfect accuracy and near-zero loss.
Abstractor architecture The Abstractor module has hyperparameters: number of layers = relation dimension d, = 4, symmetric relations (i.e.. W - W i E [d.]). linear relation activation (i.e., rel I + 2), symbol dimension d, 64, projection (key) dimension dy 16, feedforward hidden dimension dg 128, and no layer normalization or residual connection. We use positional symbols as the symbol assignment mechanism, which are learned parameters of the model.
CoReINet architecture Standard CoRelNet is described above. It simply computes, R Softmax (A)maax(A,44 A = [WWx, W Wxyley- This variant was stuck at 50% accuracy regardless of train- ing set size. We found that removing the Softmax helped. Figure 4b compares against both variants of CoRelNet.
14

---

Published as a conference paper at ICLR 2024
This finding suggests that allowing reel be a configurable hyperparameter is useful feature of the Abstractor. Softmax performs contextual normalization of relations, such that the relation between i and is normalized in terms of s's relations with all other objects. This may be useful at times, but may also cripple a relational model when it is more useful to represent an absolute relation between a pair of objects, independently of the relations with other objects.
PrediNet architecture. We used the following hyperparameters: using 4 heads, and 16 relations, a key dimension of 4 (see the original paper for the meaning of these hyperparameters), The output of the PrediNet module is flattened and passed to the MLP.
MIP. The embeddings of the objects concatenated and passed directly an MLP. The MLP has two hidden layers each with 32 neurons and ReLU activation.
Data generation The data is generated by randomly sampling a "set" with probability 1/2 and non-"set" with probability 1/2. The triplet of cards is then randomly shuffled.
TrainingEEvaluation We the crossentropy loss and the Adam optimizer with a learning rate of 10 10-3 3, 31 0.9, 02 = 0.999, & = 10-7. We use a batch size of 64. We train for 200 epochs and restore the best model according to validation loss. We evaluate on the test set.
C.2 RELATIONAL.
WOWNOONAASEQUUNN---O---S-- SEQUENCE-TO-SSQURRNC TASKS (SECTION 4.2)
C.2.1 SAMPLE-BFICIEECY IN RELATIONAL SEQ2SEQ TASKS
Abstractor architecture The Abstractor model architecture (b) of Figure 2. For each of the Encoder, Abstractor, and Decoder modules, we use L = 2 layers, 2 attention heads/relation dimensions, a feedforward network with dr 64 hidden units and a model/symbol dimension of dmodeel = 64. The relation activation function is drel = Softmax. We use positional symbols as the symbol assignment mechanism, which are learned parameters of the model. The number of trainable parameters is 386, 954.
Transformer architecture We implement the standard Transformer of (Vaswani et al., 2017). For both the Encoder and Decoder modules, use matching hyperparameters per-layer but increase the number of layers. We use layers, 2 attention heads, feedforward network with 64 hidden units and model dimension of 64. The number of trainable parameters is 469, 898. We increased the number of layers compared to the Abstractor in order to make a comparable size in terms of parameter count.
Ablation model architecture The Ablation model uses an identical architecture to the Abstractor, except that the relational cross-attention is replaced with standard cross-attentie at the Encoder- Abstractor interface (with Q < A, K + F, V < E). It has the same number of parameters as the Abstractor-hased model.
Training/Evaluation We use the crossentropy loss and the Adam optimizer with a learning rate of 10- 3, Bi 0.9, .32- = 0.999, E 10-7 We use batch size of 512. We train for 100 epochs and restore the best model according to validation loss. We evaluate learning curves by varying the training set size and sampling a random subset of the data at that size. Learning curves are evaluated starting at 100 samples up to 3000 samples in increments of 100 samples. Each 'sample' is a pair of input-output sequences. For each model and training set size, we evaluate 10 runs with different random seeds and report the mean and standard error of the mean.
C.2.2 GENERALIZATION TO NEW OBJECT-SORTING TASKS
Abstractor architecture The Abstractor model uses architecture (a) of Figure 2. The Abstractor module uses learned positional symbols, has L layer, model dimension of dmodel 64, relation dimension of d, = 4, a softmax relation activation r = Softmax, and a feedforward network with dæ 64. The decoder also has layer with 4-head MHA and 64-unit feedforward network.
Transformer architecture The Transformer is identical to the previous section.
Training/Evalusten The loss, optimizer, batch size, and learning curve evaluation steps are identical to the previous sections. Two object-sorting datasets are created based on "attribute-product
15

---

Published as a conference paper at ICLR 2024
structure" primary dataset and a pre-training dataset. As described in Section 4.2, the pre- pre-training dataset uses the same random objects as the primary dataset but with the order relation of the primary attribute reshuffled. The models trained on 3,000 labeled sequences of the pre-training task and the weights are used to initialize training on the primary task. Learning curves are evaluated with and without pre-training for each model.
C.3 MATH PROBLEM-SOLVING (SECTION 4.3)
Abstractor architectures. The Abstractor models use architecture (d) of Figure 2. The Encoder, Abstractor, and Decoder modules share the same hyperparameters: number of layers L = 1, relation dimension/number of heads d, = n4, 4, symbol dimension/model dimension d, dmodel = 128, projection (key) dimension d2 32, feedforward hidden dimension de 256. In the Abstractor, the relation activation function is rrel = softmax. one model, positional symbols are used, with sinusoidal embeddings. In the other model, symbolic attention is used with symbol library of n, = 256 learned symbols, and 4-head symbolic attention.
Transformer architecture. The Transformer Encoder and Decoder have identical hyperparameters to the Encoder and Decoder the Abstractor architecture.
Transformer+ architecture. In 'Transformer+', the model dimension is increased to dmodel 200 and the feedforward hidden dimension is increased to dæ = 400. The remaining hyperparameters are the same.
Training/Evaluation We train each model for 50 epochs with the categorical cross-cntropy loss and the Adam optimizer using learning rate of 10 10-4, B1 0.9, B3 0.995, 10-9. We use batch size 128.
C.4 ADDITIONAL EXPERIMENT: ROBUSTNESS AND OUT-OF-DISTRIBUTION GENERALIZATION IN THE
EXPERIMENTS
This experiment explores the Abstractor's 's robustness noise and outto--distributtin generalization as compared to a standard Transformer. We consider the models in Section 4.2 and the corresponding object-sorting task. We train each model on this task using 3,000 labeled sequences. We chose the fixed training set size of 3,000 because is large enough that both the Abstractor and Transformer are able to learn the task. Then, we corrupt the objects with noise and evaluate performance on sequences in the hold-out test set where objects replaced by their corrupted versions. We evaluate robustness to random linear map as well as to additive noise, while varying the noise level. We evaluate over several trials, averaging
the realizations of the random noise.
On the hold out test set, we corrupt the object representations by applying a random linear transfor- mation. In particular, randomly sample random matrix the entries of which are iid zero-ment Gaussian with variance a², t E IRd Dij ~ (0, o²). Each object in O is then corrupted by this random linear transformation, via 0, = 0g + E, ~ 00 o
Too; for each i [48). We also test robustness to additive noise
Robustness to Additive Noise
Robustness to Random Linear Transformations
1.00
Anserattrr
1.00
Metractor
Sorting Accuracy
Sorting Accuracy
0.75
0.75
0.50
0.50
0.25
0.00
0.00
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
1.0
Noise Level
Noise Level
(a) The Abstractor more robust to corruption by additive noise.
(b) The Abstractor more robust to corruption by random linear transformation.
Figure 8: Experiments robustness.
16

---

Published as a conference paper at ICLR 2024
The models evaluated the hold-out test set with objects replaced by their corrupted version. We evaluate the sorting accuracy of each model while varying the noise level a (5 trials at each noise level). The results are shown in figures 8a and 8b. We emphasize that the models are trained only the original objects in O, and are not trained on objects corrupted by any kind of noise.
This experiment can be interpreted in lights: the first is robustness noise. The second is form of out-of-distribution generalization. Note that the objects seen by the models post-corruption lie in different space than those seen during training. Hence the models need learn relations that are in some sense independent of the value representation. As a theoretical justification for this behavior, Zhou al., 2009 shows that (Øz, by) (a. y) in high dimensions, for a random matrix & with iid Gaussian entries. This indicates that models whose primary computations are performed via inner products, like Abstractors, may more robust to this kind of corruption.
17


Reviews:
[{"review_id": "M5iQ6geZzR", "reviewer": "Reviewer_eiad", "summary": "The authors introduce a modification of self-attention, referred to as \"Relational Cross-Attention\" (RCA), in which value projections are replaced with input-independent vectors, referred to as \"symbols\".  These symbols are either learned or implemented as (relative) positional embeddings. Additionally, the softmax function used in the attention operation can be replaced with element-wise activation functions for improved performance on some tasks. The authors replace the attention-layer in a transformer encoder layer with RCA, and refer to the resulting layer as an _Abstractor Module_. \n\nVarious experiments on toy-problems requiring relational reasoning are carried out. The authors show that Abstractors out-perform transformers and CoRelNets (another explicitly relational baseline) on these tasks, and that Abstractors are able to learn generalizable relations (demonstrated through pre-training experiments) and can easily be constrained to learn only symmetric relations (through symmetric inner products $\\phi_{\\theta}(x)^\\top \\phi_{\\theta}(x)$, rather than the default $\\phi_{\\theta}(x)^\\top \\psi_\\omega(x)$ in RCA). Additionally, they experiment with various forms of overall architecture, varying whether the Abstractor layer is used alongside Encoders (in parallel or in sequence) prior to decoding.", "questions": "Some minor questions:\n1. Multi-attention decoder: Perhaps this form of decoder is standard, but why is CausalSelfAttention used before Cross-Attention (which is not causally masked?)\n2. On the Math Problem experiments: It seems the dataset consists of 8 tasks, but only 5 of these are investigated in the experiments. Is there a particular reason for this omission?\n3. On the SET Comparison against a symbolic-input model: When pre-training the abstractor relations, is the input from the same pre-trained CNN as when training the multi-head abstractor, or also from the symbolic inputs?\n4. The authors argue that relations are well-modelled as inner-products. I am curious as to which differences this might impose on the learned relations when compared to the relational-form used in the PrediNet, in which a difference of projections (\"differential comparator\") is used?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Overall presentation is very clear and the method is well-explained, though the analogy to self-attention may be better mentioned earlier in the paper, rather than first \"deriving\" the relationship between the calculation of a project relation matrix and $QK^\\top$ Self-Attention. \n\nThe experiments serve to demonstrate the efficacy of Abstractors in toy settings, and some effort is made to demonstrate that they do indeed learn re-usable relations which are robust to noisy data. Alongside the ablations, these provide sufficient evidence that abstractors provide a strong relational inductive bias which may make them useful in more applied settings - though this remains to be demonstrated. \n\nThe experiments comparing Abstractors with pre-learned relations (per-head) against a symbolic input MLP are particularly well-thought-out and interesting, in that they go some way toward demonstrating the generality of the learned relations.", "weaknesses": "Methodologically there are no substantial weaknesses, as comparisons against baselines are made as fair as possible. One minor exception is the omission of a baseline transformer for the Order Relation and SET tasks, though this was presumably an intentional omission based on the already superior performance of CoRelNet compared to Transformers?\n\nThere is one consistently made claim which may be slightly overstated (this is a question to the authors) - namely that the output of Abstractors represents \"purely relational information\"; I believe this only holds if there are no residual connections in the abstractor module implementation being used (which is offered as an option); if there is a residual connection, then it seems the abstractor MLP could still learn to operate on information present in object-level representations.\n\nRelated to the above point, it would be very interesting to see an attempt to \"lift\" the learned relational information into a symbolic form, as was done e.g. in the cited PrediNet work (Shanahan)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3mwop6wpdd", "reviewer": "Reviewer_nkpJ", "summary": "This paper proposes a novel Abstractor module that uses relational cross-attention to enable explicit relational reasoning. such relational cross-attention disentangles relational information from object-level features. The experimental results show that such explicit relational reasoning greatly improves the sample efficiency of Transformers on relational tasks.", "questions": "1. Could you compare both the performance and sample efficiency of Abstractor with other relational structures like PrediNet?\n2. STSN (Mondal et al., 2023) has used Transformer for RAVEN and PGM problems, which involve relational reasoning, how much would the Abstractor improve over Transformer in those tasks?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow.\n2. The Abstractor module and relational cross-attention are novel and interesting, simple but effective.\n3. The method is well motivated and the proposed method is well justified on a variety of tasks.", "weaknesses": "1. When comparing sample efficiency on pure relational tasks, the Abstractor is only compared to the Transformer baseline. It would be better to show the comparison with other relational structures like PrediNet.\n2. There is no explicit section for related work and limitations are not discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FKQwFuogXn", "reviewer": "Reviewer_4fms", "summary": "The motivation of this paper lies in addressing the challenge of relational and abstract reasoning. Contemporary deep learning models excel at tasks involving semantic and procedural knowledge, but their abilities to infer relations, draw analogies, and generalize from limited data to novel situations remain limited. The popular Transformer architecture, through its attention mechanisms, has the capacity to model relations between objects implicitly. However, these standard attention mechanisms often create entangled representations, blending relational information and object-level features in a way that is not optimal for efficient learning of relations.\n\nThis paper proposes a new variant of Transformers, named the Abstractor. The core of the Abstractor is a variant of attention, named relational cross-attention. Compared to the vanilla attention in the Transformer, relational cross-attention disentangles the object-level feature and relational information by replacing the value vector in vanilla attention with a learnable symbol vector or positional function, which is independent of the object-level feature (X). The authors demonstrate the effectiveness of their Abstractor through \"pure relational\" and \"partial relational\" tasks, and compare the Abstractor to previous architectures designed for relational tasks.", "questions": "see weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The overall problem is interesting. As the author has noted, human brains can perform tasks involving analogy and abstraction with limited experience, whereas current models require vast amounts of data to acquire such abilities.\n\n2. The proposed Abstractor, a novel variant of the Transformer model, disentangles relational and objective features through relative cross-attention. \n\n3. The authors have proven, through comprehensive experiments, that the Abstractor is more sample-efficient than the vanilla Transformer and previous architectures used for pure or partial relational tasks.", "weaknesses": "1. The overly constrained setting limits the significance of the work. The authors conduct experiments under two settings: \"purely relational\" and \"partially relational\". As per the authors' definition, \"purely relational\" implies that object-level features are extraneous and that the statistics of relation/order are already sufficient for solving the task. This is an extremely restricted setting and may not fully represent the complexity of real-world tasks where both relational and object-level information are often important. The authors use math problem-solving to represent the \"partially relational\" setting, but the math problem here is arguably more relational/symbolic than object-level. It seems like the math problems here can be solved by symbolic rules.\n\n2. Scalability is one of the most significant advantages of Transformer architectures. The performance can increase with the model size and data size. Given that the Abstractor is a variant of the Transformer, it's essential to determine whether the scaling law still applies to the Abstractor. From the results provided, the performance can outperform the vanilla Transformer when the data size increases from the 1000 - 5000 range. But what about using more data and a larger model size? Will the Abstractor consistently improve and outperform the Transformer?\n3. Why did the authors choose to replace value vectors with input-independent vectors, while keeping the query and key vectors the same (Q -- X, K -- X, V -- S)? Would not the configuration (Q -- S, K -- S, V -- X) also disentangle object-level features (x) and symbolic vectors (s)? To me, the latter one is more intuitive: the relation weight R_{ij} between i, j is represented by the inner product of symbolic vectors, and then object-level features are weighted by R_{ij}.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8OS7M2t7Es", "reviewer": "Reviewer_qTc8", "summary": "This paper presents a modification on attention that replaces the value part of self-attention with an \"relational-only\" representation S. This representation serves as an \"information bottleneck\" for the model so that it can learn separate encodings for object relations and values (eg in a separate encoder). The experimental results are strong and offer good empirical support for the effectiveness of the Abstractor layer.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The proposed abstractor method outperforms all baselines and is more sample efficient in learning relational tasks.\n- The proposed method is interesting and a relatively simple modification on top of self-attention. However, I have some questions about how S works (below)", "weaknesses": "- It is not clear to me how S works. Section 2.2 says \"The symbols S can be either learned parameters of the model or nonparametric positional embeddings.\" How is S different from a positional embedding then? If S is unique per token, even if the token is a repeat of another token value-wise, then S is essentially a positional embedding. If S is unique per *value* of a token (ie all tokens of value v have the same s_i), then doesn't S implicity encode the value of the token?\n- If the outputs of an abstractor layer are \"abstract states that represent purely relational information\" (section 2.3), then how are features associated with objects passed along/learned in models that use abstractor layers? In all the examples presented in Figure 2, abstractors are used in conjunction with regular encoders. How can you ensure that the abstractor is learning meaningful information and that all the \"information flow\" does not happen through the encoder layers in parlalle setups (architectures c, d, and e)? \n- Section 4.1 claims that the authors hypothesize that \"the ability to model relations as multi-dimensional is also the reason that the Abstractor can learn the order relation better than ...\" Are there any experiments to test or verify this hypothesis? The paper also makes many claims about an information bottleneck in S, but there is no analysis on what is actually learned in S.\n- Is there an ablation on model size? The baseline transformer is not very large, but it may be that a smaller transformer can learn with fewer training data points, or a larger transformer may converge faster.\n- The experiments seem to be on simple problems with small models. Simple problems may not necessarily be an issue since they are relatively diverse problems, but it would be nice to see larger, more complicated problems. For example, the partial order task training set is small enough that one could consider using in context learning with a large LLM, which may offer comparable performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)