Attention as a Hypernetwork
Authors: Simon Schug
Seijin Kobayashi
Yassir Akram
Joao Sacramento
Razvan Pascanu
Venue: ICLR 2025 Oral (2025)
Keywords: attention
compositional generalization
abstract reasoning
in-context learning
transformer
mechanistic interpretability
Primary area: interpretability and explainable AI
TL;DR: Multiple heads in attention create reusable operations that support compositional generalization in abstract reasoning.

Abstract: Transformers can under some circumstances generalize to novel problem instances whose constituent parts might have been encountered during training, but whose compositions have not.
What mechanisms underlie this ability for compositional generalization?
By reformulating multi-head attention as a hypernetwork, we reveal that a composable, low-dimensional latent code specifies key-query specific operations.
We find empirically that this latent code is predictive of the subtasks the network performs on unseen task compositions, revealing that latent codes acquired during training are reused to solve unseen problem instances.
To further examine the hypothesis that the intrinsic hypernetwork of multi-head attention supports compositional generalization, we ablate whether making the hypernetwork-generated linear value network nonlinear strengthens compositionality.
We find that this modification improves compositional generalization on abstract reasoning tasks.
In particular, we introduce a symbolic version of the Raven's Progressive Matrices human intelligence test, which gives us precise control over the problem compositions encountered during training and evaluation.
We demonstrate on this task how scaling model size and data enables compositional generalization in transformers and gives rise to a functionally structured latent space.

Published as a conference paper at ICLR 2025
ATTENTION AS A HYPERNETWORK
Simon Schug
Seijin Kobayashi ETH Zürich
ETH Zürich sschugeethz.ch
beijink@ethz.co
Yassir Akram
João Sacramento'
Razvan Pascanu' Google DeepMindd razp@google.com
ETH Zürich
Google, Paradigms Intelligence Team joaosacramen: com
yakram@ethz.co
ABSTRACT
Transformers can under some circumstances generalize to novel problem instances whose constituent parts might have been encountered during training, but whose compositions have not. What mechanisms underlie this ability for compositional generalization? By reformulating multi-head attention as a hypernetwork, we reveal that a composable, now-dimensional latent code specifies key-query specific operations. We find empirically that this latent code is predictive of the subtasks the network performs on unseen task compositions, revealing that latent codes acquired during training are reused to solve unseen problem instances. To further examine the hypothesis that the intrinsic hypernetookk of multi-head attention supports com- positional generalization, we ablate whether making the seannettorr--geerratee linear value network nonlinear strengthens compositionality. We find that this modification improves compositional generalization on abstract reasoning tasks. In particular, we introduce a symbolic version of the Raven's 's Progressive Matrices human intelligence test, which gives us precise control over the problem composi- tions encountered during training and evaluation. We demonstrate on this task how scaling model size and data enables compositional generalization in transformers and gives rise to functionally structured latent space.
INTRODUCTION
Abstract reasoning is a cornerstone of human intelligence. Many intelligence tests in psychology measure intelligence using abstract reasoning tasks (e.g. Wechsler 1958 Binet & Simon 1961
Raven 1962). Arising from the idea that combinatorial structures are central to human cognitive abilities, the ability of connectionist models to reason has been debated for decades (e.g. Rumelhart
& McClelland 1986; Fodor & Pylyshyn| 1988; Smolensky [991). However, with the success of neural networks scale, the capacity for reasoning has scemingly come into reach (Huang & Chang 2023 Arguably key to this success is the ability for in-context learning paired with the capacity to flexibly recombine learned skills and knowledge unseen settings. While these capabilities emerge at least partially from training large models on huge datasets (Brown et al.) 2020). it is not clearly understood how the resulting networks implement them and why they still often fail.
A number of recent studies that aim to illuminate in-context learning abilities have found that transformers can learn to perform gradient- optimization within sequence (Dai et al. 2023 Akyürek al. 2023 von Oswald al.) 2023). They build on the insight that linear attention can be understood as a faight programmer (Schmidhuber 1992), that learns to construct an implicit model in-context, with weights given by a of input dependent outer products (Schlag al.) 2021). Complementing this perspective, (Hendel et al. 2023 Todd et al., 2024) have found that in-context learning creates task vectors in the hidden activations of deeper layers of the network, which summarize in-context information and modulate the forward computation. What is striking is that learning in-context can under some circumstances lead to compositional generalization. that is
I Shared senior authors
Code available

---

Published as a conference paper at ICLR 2025
A
Curput
Linear hypernetwerk
Multi-head attention
Dutpur
Modue
Fix Wi
Linear hypernetwork
Latent code
M
peramelariin Liees
Wailin
Wg
Value network
Waiter
Moduic
k2:00 =T
Whilen
Value network
B
W
Ameettian
Headt
Dor productuct
Head indeo/
latere
Kay iodex
Query
Toker
Figure 1: Hypernetwork attention. A A linear hypernetwork maps a latent code to a set of
parameters that configure a value network process the input. B The attention scores along the head index form the latent code of the hypernetwork. C Multi-head attention can be equivalently expressed linear hypernetwork that configures key-query specific computations of linear value network.
generalization to unseen combinations of previously observed constituents (An et al. 2023 Lake
Baroni 2023 Hosseini al., 2202)), ) capacity that neural networks notoriously struggle with (Srivastava ct al. 2023 Press et al. 2023 Dziri et al., (023).. The mechanism behind this ability
however is not well understood. Here we to shed light on this question.
Our key finding is that through the use of multiple heads, the attention mechanism of transformers is mathematically equivalent to a hypernetwork, neural network that reconfigures the computation of another neural network as shown in Figure [ This results in novel interpretation of the attention scores along the head index as a compact latent code that specifies the operations the attention layer applies to each input. Notably, we hypothesize that because multi-head attention shares the same hypernetwork per layer, it encouraged to reuse and recombine previously acquired operations.
In order to investigate whether this constitutes faithful characterization of what transformers learn in practice, we study in-context learning tasks that require recomposing knowledge obtained during training. As one such task, develop challenging abstract reasoning task based the Raven's Progressive Matrices human intelligence test (Raven 1962) and show how multi-head attention develops structured latent code that captures information about the operations implemented.
Our main contributions can be summarized as follows:
We reformulate standard multi-head attention from a hypernetwork perspective, revealing that network computations are specified by low dimensional latent code.
We show that scaling up model size and data enables transformers to compositionally generalize on abstract reasoning tasks and leads to structured latent space that is predictive of network function.
To test the hypothesis that the hypemetwork mechanism supports compositionality. introduce simple modification to multi-head linear attention that makes the value network nonlinear without introducing additional parameters, finding that it improves compositional generalization.
We introduce SRAVEN, challenging symbolic, abstract reasoning task whose difficulty can be controlled parametrically and which offers precise control over the problem compositions encountered during training in order to test for compositional generalization.
2 ATTENTION AS HYPERNETWORE
In this section, we will first briefly recapitulate hypernetworks. (Ha et al. 2017) and standard multi- head dro-product attention (Vaswani et al. 2017) before showing how attention can equivalently be expressed hypernetwork. We will then introduce Hypernetwork Linear Attention (HYLA) as a simple modification of linear attention that renders the hypernetwork mechanism inside attention more expressive.
2

---

Published as a conference paper at ICLR 2025
For notation, will use bold lower-case letters for vectors (e.g. 2), bold upper-case letters for matrices (e.g. A) and bold, italic upper-case letters for learnable parameter matrices (e.g. Wy).
2.1 MULTI-HEAD ATTENTION
hypernetwork a neural network h(z; 0) parameterized by 0 that takes as input latent code and outputs parameters W that parameterize a value network f(;; W). which then processes the inputs (see Figure IA]. The latent code typically now-dimessional and can be interpreted as specification of the computation performed by the value network. As we will see in the following, multi-head attention can viewed as a linear hypernetwork producing the weights for a linear value network. Note that this is different from the observation that linear attention can be understood as a fast weight programmer (Schlag al. 202T), where the fast weights of a value network are constructed sum of outer products over the key indices instead of the head indices.
Self-attention maps a sequence of inputs X DxT to a sequence of outputs E RDxT For each attention head E {1,..., .....,,), H). the inputs are projected into keys K = Wh X and queries
Q2 W X using head-specific projection matrices Whey W W E R ny x with
D The attention matrix can then be obtained
A A=o A A=
An) with A
Q Kr
Dhhead
where A- = (an.g.k)....k E R stacks the head-specific unnormalized attention matri- ces (Aa), into a three-dimensional array and applies a normalization operation a(-). For linear attention, a..) = Id(-) is simply the identity whereas for softmax attention a.)) Softmax() applies the softmax operation for each head and query index independently across the key indices.
Considering a single query index q E (1...., T} and x; E TR. the k-tth column X, multi-head attention (MHA) can be equivalently expressed as hypernetwork:
H T
MXA,(X) = WWo
Obggk W x4
(2)
k=1k= k=
Wout
k Wyauue
(3)
1=1
Ob.g.k
Wout Wraaue
X6
(4)
Intent code
hypernewook
X
(5)
value network
where @ denotes concatenation and the Wout hau IR DxD are defined the slices of Wouu [R DDD
such that Woutt = + H
h= Wout See Figure Ic for diagram of the computation.
As result, the key-query specific attention scores over the multiple heads form latent code, where the number of heads // corresponds to the latent dimension of hypernetwork. Accordingly, the dot products between any key-query pair over all heads can be interpreted as amortized inference process to configure the computation implemented in the linear value network.
The key significance our finding is that suggests that the use of multiple heads in attention allows it to implement specialized, reusable operations through the latent code of hypernetwork. It has been theoretically and empirically shown that hypernetworks support compositional generalization (Schug et al. (024). providing possible explanation for similar capabilities observed in transformers (Lake & Baroni] 2023) An al. 2023 Hosseini al. (2022). Note that multi-head attention reuses the same hypereetwork for every key-query index. Specifically, the hypernetwork linearly combines
"e focus self-attention, but the same derivation holds for cross-attention, where the keys and values obtained the output of encoder and the queries come from decoder.
3

---

(without adding Published as a conference paper at ICLR 2025
the matrices ht Wy according to the weights given by ah.g.k This incentivizes reuse of latent codes and their associated operations implemented through the value network.
2.2 HYPERNETWORK LINEAR ATTENTION (HYLA)
If the perspective of multi-head attention hypernetwork indeed reflects how it operates in practice, it should be possible to modify multi-head attention in way that further reinforces the hypernetwork mechanism. For instance, both the hypemeetwrk and the value network in multi-head attention are simple linear networks and could be replaced by nonlinear networks. Furthermore, the normalization operation (-) could be used to encode inductive biases such as competition or sparsity in the latent code by imposing structural constraints.
Here, we focus on a simple modification to linear attention, where we make the value network nonlinear and normalize the latent code along the head index (instead of along the key index in standard softmax attention). Specifically, we define Hypernetwork Linear Attention (HYLA) to use a single hidden layer value network by adding nonlinearity. This should increase the expressivity of the subfunctions that are learned and recomposed by the layer. Noting that the output and value projections W Waauue already form deep linear network in standard multi-head attention (c.f. Eq. 4 this can be achieved without adding additional parameters as follows:
T
HYLA,
Σ
Σ
aR.o. Wout
Σ
an.o. Wraaue X
(6)
A=
T
W
a.k xx)x),,
where o(x) is an element-wise nonlinearity; here we set it to be ReLU, Φ(x) Additionally, we use (-)) RMSHead(-). to normalize the attention scores for each key and query index independently across the head indices using RMSNorm(z)
max(0, x))
V >
learnable scalar in Zhang & Sennrich (019)). This ensures that the parameters of the value network generated by the hypernetwork maintain the variance preserving properties neural network initializations that have been found to be important for the stability of gradient-based training (Glorot & Bengio 2010). The resulting HYLA-layer is simple drop-in replacement to standard attention. Notably, since the normalization operation (--) is local the query and key index, it requires no communication across keys as opposed to softmax attention. By making the value network more expressive, we allow the network to implement more complex operations for a given latent code which we hypothesize strengthens the ability for compositional generalization.
3 COMPOSITIONAL GENERALIZATION ON FUZZY LOGIC FUNCTIONS
We start by developing a fuzzy logic task with compositional structure, inspired by similar task previously considered by Rahaman al. (202T). Learning such tasks an in-context learning setting will allow us to study the ability of multi-head attention-based models to compositionally generalize and analyze the structure of their latent code when solving each task.
Given set of L scalars {I1,..., ] with E [0, 1] we define fuzzy logic operators on the elements of this set as the Zadeh operators (Zadeh 1965):
A æij == min(
Vaj == max( xy), Tr. =
(8)
These operators have the property that for boolean inputs x; E (0. 1 the outputs coincide with the outputs of the corresponding Boolean operation. Each task is then defined as function in disjunctive normal form consisting of K terms such as for example for = 4 and K =
f(11,
(x1 A 22 Ax ZA)
(11 x2 Tz AA)
Term
where each term is conjunction of all variables with negations applied to some of them. In order to test for compositional generalization, we index each of the 2 possible terms, generate all possible combinations of K terms and hold out fraction of the combinations for -uttr--distriuution (OOD)
4

---

Published as a conference paper at ICLR 2025
A
Eoempii test i le, Vay MayAng
B
Liteer atcertion Sotrmax sttention
Target
Temm
HYLA
Term decoder (F scord)
Softmaa aor attention
Uneer attention
IMA
Topel abc
100
1000
Bercnns
Layer
Layer
C
Softmax attercion
Unear sttention
HYLA
000 oracsion 0.3
OOD fractio 0.5
OCO fraction 0.7
Softmma attention
Uneer attention
HYA
000
0.3
100
0.6
-100-
0.4
In-contect exampless
examples
examples
Figure 2: Compositional generalization on fuzzy logic functions. A We split fuzzy logic functions according to their constituent terms into train and out-of-distribution (OOD) sets to measure composi- tional generalization in a sequence model that learns these functions in-context. B The latent code of the response token is predictive of the constituent terms underlying each task. Shown is the F1 score on unseen tasks of logistic regression classifiers for each layer and term, trained to predict the terms underlying each task based on the attention scores across the head index for the response token attending to itself. C Task performance on unseen tasks reported as OOD R² for varying number of in-context examples and fraction of tasks held-out during training. D ISNE visualization of the latent codes (attention scores across the head index for the response token attending to itself) colored according to the target label (top) and colored according to the first term of the fuzzy logic function of each task (bottom).
testing, ensuring that their constituent terms have been encountered in some tasks during training. For each task, we present N examples in-context (see Figure ZA) where each token has dimension L ++1 1, consisting of the concatenated inputs sampled uniformly from [0, I] and the corresponding target of the function for the current task. For the final token, mask out the target and use the output logits of the final token to predict the target, measuring the loss using the mean-squared error.
3.1 COMPOSITIONAL GENERALIZATION
We first investigate how well multi-head attention based transformer models are able to learn fuzzy logic tasks in-context and compositionally generalize to held-out tasks. We vary the number of examples presented in context, as well as the fraction of combinations of fuzzy logic terms held-out from the training set. Figure 2 compares standard multi-head softmax attention and linear attention, as well as our HYLA model in terms of their coefficient of determination (It2) on held-out tasks. We find that all models are able to solve the fuzzy logic task given sufficient examples are presented in context. As the fraction of held-out tasks increases, compositional generalization performance starts decreasing for all models. Interestingly, HYLA is less affected by this decline than softmax and linear attention. We hypothesize that the more expressive value network allows to more efficiently represent the nonlinear terms being composed and thereby strengthens the inductive bias towards compositionality. Consistent with this interpretation, in-distribution the models show comparably small differences in performance (see Figure A3B+D). In Appendix A we conduct additional model ablations that separate the contributions of the RMSHead normalization and the nonlinearity added to the value network, showing that both components in combination are required obtain the observed performance improvements.
Our compositional split only contains novel combinations of known terms. A priori. it might be possible that the system learns to compose the basic fuzzy logic operations in a way that allows it to generalize to any fuzzy logic function, including those that contain novel combinations of unknown
5

---

Published as a conference paper at ICLR 2025
A
Split rule combinations
Sample task rules
Sample task instance
In-context learning
Rules
Splits
Combination rules
Predect omppante poret
9 pragression [+1]
AA
B progression (+1)
constant
Connect parals
3
progression (+2)
Permutation
8
3
Sequainee model
D pragression (-1)
dosh column penmate fnauurs ts mackt
progression
F addition
000
2
G difference
H. distribute-hhree
Calumn1
Calumn
Columi
Co
Respoise
panel
B
Finding correspondences
Hyyotthesis Group shape
Hypothesis Group orientation
H
shape
ertentatto
Context panels
creudinn
rumber
number
3
2
criertation
encoting
crtentation
22
cilertation
StAze
?
crlerratio
number
orientation
shops
crientation
shape
Response panel
Symbbels extrelngg tird pame! pootiii engerings
Pooel
Porel:2
Pann
Pane
Porel
Figure 3: SRAVEN. A Illustration of SRAVEN task generation and the construction of the compositional generalization split. B Example problem instance illustrating key challenge of the original Raven's Progressive Matrices of finding correspondences (adapted from Carpenter et al.) (1990)). When attempting to solve this instance, different hypotheses over which figural elements are governed by consistent rule across rows are possible. This is akin to different orderings of the symbolic features.
terms. However, testing on functions obtained as novel combinations of K unknown terms, we find that none of the models considered here is able to solve such a task (see Figure A3C).
3.2 LATENT CODE STRUCTURE
Next, analyze the structure of the latent code for the different models solving the task. As suggested by Eq. 4 the attention scores for a given key-query pair across the head dimension configure the computation performed in the value network, and accordingly to reflect the fuzzy logic function that needs to be implemented to correctly predict the output of the response token. To test this hypothesis, we collect the attention scores for the response token attending to itself, obtaining a vector of dimension H for each layer and task. For this purpose, we use the tasks held-out from training, varying the response token while keeping inputs shown in-context fixed. In order to visualize the H-dimensional points we thereby obtain for each layer, we reduce the dimensionality using tSNE (Maaten & Hinton 2008). The results of this analysis are shown in Figure[p (for an extended figure showing all layers, see Figure A2).. In line with our hypothesis, the latent codes for each model form clusters. The structure of these clusters is only partially explained by the specific output value of each function. Strikingly, coloring data points according to the fuzzy logic terms that make up each function reveals that clusters the latent codes correspond to the terms. Indeed, it is possible to decode the terms underlying each function from the latent code using a logistic regression classifier which is trained to predict the function terms given the latents and is evaluated on held-out tasks shown in Figure 2
would expect it
SRAVEN: SYMBOLIC RAVEN
Next, we introduce an abstract reasoning task based on Raven's Progressive Matrices (Raven 1962), we refer as SRAVEN. We develop this task to provide a challenging benchmark reasoning capabilities while giving us fine-graine control over the constituent parts of each task in order assess compositional abilities and analyze latent code structure.
that probes symbolic
4.1 ABSTRACT REASONING BASED ON SYMBOLIC RAVEN'S PROGRESSIVE MATRICES
Raven's 's Progressive Matrices are a classic human intelligence test that probes the ability to induce abstract relations (Raven 1962.. It requires test subjects generate and verify potentially large

---

Published as a conference paper at ICLR 2025
A
B
Softma omtendion
Unear attention
HIL
Softmoo attention
Uneo
HYLA
6.9
Ta witth
2a withh
mech
0.9
layers
layars
0CO 000 accurayy
0.-
CCD accurocy
c.-
6.-
c.-
65
100
P0L
Number of trtriingg tasks
Number of training tasks
Figure 4: Scaling data and model size on SRAVEN. A Compositional generalization measured by OOD accuracy function of the number of training problem instances for different widths (scaling embedding and key-query-vaie dimensions). B Same as A but increasing depth instead of width.
number of hypotheses over rules that can parsimoniously explain given problem instance CCarpenter ctt al., 1990). While there have been previous machine learning datasets inspired by Raven's s Progres- sive Matrices (Wang & Su 2015 Zhang et al. 2019 Barrett et al. 2018 Hu et al. 202I). we seek to study a setting that is both symbolic, compositional and models a key difficulty of Raven tasks referred to as finding correspondences by Carpenter al.] (1990) which requires searching through
large number of possible hypotheses. We discuss existing Raven-based benchmarks in Section 6
4.2
TASK DESCRIPTION
Figure A illustrates the generative process of SRAVEN. For each task, a matrix of eight context panels presented, and the goal is to predict the contents of the final response panel. In the original Raven's s Progressive Matrices Test, each panel typically displays a combination of geometrical shapes, as for instance shown in Figure BB. For our symbolic version, each panel is defined as a tuple of K integers that symbolically encode features (in our experiments we set K = unless stated otherwise). The K features within each row evolve according to K rules. Every rule is function that takes two integers as input and outputs an integer. Figure 33 lists all rules, and Appendix BB.3 describes each rule in detail. In order maintain finite set of possible symbols, we limit all tasks to contain only integers {0, 1, 2.... F - }} (in our experiments, we set F = 8). Arithmetic rules such as progression, addition or difference are therefore defined using modular arithmetic modulo F.
To create a task, we first sample K out of the R = 8 possible rules and then use each of the K rules, to generate three sequences of length three given random inputs to each rule. As result, we obtain nine panels, where each row can be explained by the same set of K underlying rules. We use the first eight panels as context panels and present them sequentially to a sequence model, which needs to predict each feature of the response panel independently. A prediction is considered correct if and only all subpredictions were correct. In the original Raven task, a number of possible answer panels are presented alongside the context panels and test subjects are asked to pick the one which contains the correct answer (Raven) 1962). Since in our case, the correct answer will consist of combination of finite set of known symbols, predict the missing symbols. Conveniently, this helps avoid potential shortcuts for solving a task that might arise from a biased generative process create multiple-choice answers
can ask the agent solving the task to directly
popular prior
Raven-based benchmark has been found to be affected by (Zhang et al. 2019 Hu et al. 2021).
One of the key features of SRAVEN is its compositional structure. Each task consists of combination of a finite set of rule combinations. To systematically test whether a system trained on a subset of rule combinations can generalize to unseen combinations, we split all rule combinations into a train set and a test set. Unless noted otherwise, we hold out 25% of all possible rule combinations for evaluation in our experiments. Note that permutations of rule combinations (e.g. AB and BA) are considered to be the same and will only appear in one of the two sets. This ensures that at test time, the agent is confronted with instances of tasks whose underlying combination of rules has never encountered during training.
4.3 FINDING CORRESPONDENCER
Making the task symbolic instead of presenting images of geometric shapes comes with the advantage of smaller input space, allowing us to scale to larger model and data sizes in our experiments. How-
7

---

Published as a conference paper at ICLR 2025
ever, it removes the visual representation learning step of discovering the right symbolic abstractions. We argue that while this is a difficult problem worth studying by itself, there is a major source of difficulty that arises even with access to symbolic abstractions. As Carrpentrrr et al. (I990) argue in their theoretical account of the original Raven test, what makes many tasks difficult is the problem of finding correspondences of rules feature dimensions.
This is best illustrated when considering the example in Figure B showing an adapted version of Figure 5 of Carpenter et al. (1990). When first attempting to solve this task, a test subject might hypothesize that there is a rule guiding the numerosity and orientations per shape. For example, one might try to find rule that governs all curved shapes, one rule that governs all lines and one rule that governs all rectangles. Translated into a symbolic encoding, this is akin to ordering symbols by the shape and trying to find rules along each dimension in this ordering. However, this approach turns out to not provide a parsimonious answer. The correct answer can instead be obtained by grouping the objects by their orientations, i.e. by finding a rule that applies to all horizontal objects and rule that applies to all vertical objects. In the symbolic encoding this new correspondence of features amounts to new ordering of the same symbolic encoding where symbols are now ordered by orientation and one needs to find rules guiding the numerosity and shape per orientation.
To model the property of finding correspondences in SRAVEN we similarly allow for permutations of the features. Specifically, we sample and apply column -peecf--peccfic permutation of the features when generating each task, such that the entries of the tuples encoding each panel are permuted with a consistent permutation for given column across all rows. This makes finding the correct solution more difficult, since a task can no longer be treated as K independent tasks with a single feature and as result the number of possible hypotheses grows significantly (see Appendix B. for more details). Note that in principle. a problem instance might have multiple valid but divergent answers, making the task ill- poo-po.ed. In the settings considered here, this happens for less than half a percent of the problem instances. See Appendix B.2 for more details on handling such ambiguities.
4.4 RESULTS
Scaling model size and data enables compositional generalization. We first evaluate to what extent transformers can compositionally generalize on SRAVEN as we increase the number of problem instances trained on and scale up the models both in terms of width and depth. Figure Aigure4ssher shows the results of these scaling experiments. Given sufficient data and model size, all model classes are able to solve 80% of problem instances created from tasks held-out during training. The inductive bias of making the value network more expressive by introducing a nonlinearity in HYLA seems to be beneficial for this task. Especially when trained on fewer problem instances and for small model sizes, it outperforms both linear and softmax attention.
Disrupting the hypernetwork mechanism hurts compositional generalization. Varying the number of heads offers another way of manipulating the hypernetwrr mechanism of meati-head attention. Specifically in the case of a single head the mechanism degenerates, only allowing the network to rescale the weights of the sole remaining value projection but no longer allowing it to compose multiple value projections, c.f. equation (4. In line with this argument, we observe a noticeable decrease in OOD performance in Figure 5 for the single head case.
The latent code is structured according to the rules of the task. We next analyze the latent code for the different models solving the task. To this end, we collect the attention scores of the response token attending to itself. Following the perspective of attention as a hypernetwork, we expect latent codes corresponding strikingly structured code for the final layer across models, with clusters closely matching the rule to be inferred to correctly predict the output for each given task. In comparison, Figure 55A colors the points in latent space by the magnitude of the inputs the response token needs to attend to, to make the correct prediction. This tends to explain clusters more prominently in early layers, as shown in the appendix in Figure and for 16 layer softmax transformer in Figure A7
the same rule to form clusters. Indeed, Figure (53 reveals a
To obtain a better understanding of the semantics of the latent code, we explore the pairwise cosine similarity between average latent codes for each rule across tasks for HYLA. Figure shows that semantically related rules like sum/difference can form clusters, suggesting that latent codes might even be reused to solve related task operations. Finally, we train logistic regression classifier
8

---

Published as a conference paper at ICLR 2025
A
Input
C
Sotonak
D
actendinn
Unear attention
HYLA
HYLA
Sollmme altertionn
Lineat altention
100-
accurscy
acurscy
c.-
50-
0.5
0.5-
-50-
0.4
Number
Number features
B
E
Rule
overgge latent coulr rach roll
Poinetre cosine simlenty between
Softmac atbention
Unear
Softiax attendon
Linear attendon
HYLA
-0
Ruee decoder oo'e core
0.8-
HYLA
50-
Rule
0.6-
10pw
G-
-50-
-50
Rule
Layer
Figure 5: Latent code structure of SRAVEN. A ISNE visualizations of the final layer latent codes for the final response token colored by the magnitude of the predicted target value. B Same as A but colored by the ground -ruth rule the model needs to apply to generate the correct prediction. C OOD accuracy for varying numbers of attention heads. For single head, the hypemnetorr mechanism is absent, which hampers OOD generalization. D The difficulty of SRAVEN can be parametrically controlled by varying, K, the number of features per panel E. Heatmap showing the pairwise cosine similarity between the average latent code of each rule for HYLA revealing how semantically related rules form elusters. For instance, rule F (addition) and G (difference) are implemented with very similar latent code, indicating that the same code might be reused by flipping the sign of the operands. F Decoding performance of a logistic regression classifier trained to predict the ground truth SRAVEN rule based on the latent code at the final response token of training tasks and evaluated on unseen OOD tasks, revealing that the latent code is predictive of the implemented rule.
to predict the ground-truth SRAVEN rule based on the latent code at the final response token of in-distribution SRAVEN tasks and evaluate it on unseen OOD tasks. This reveals that, especially for later layers, the latent code is strongly predictive of the subfunctions performed by the network.
5 LANGUAGE MODELING
Multi-head attention has proven especially effective for language modeling. With language itself being highly compositional, stands to question whether the implicit hypernetwork mechanism in multi- head attention might play a part in explaining its success. To shed light this question, conducted language modeling experiments and evaluate the performance of HYLA in comparison to linear and softmax attention. We train decoder-only transformer models with 50M parameters autoregressively for 130 Billion tokens on the C4 dataset (Raffel et al.) 2020). We find that strengthening the hypernetwork mechanism via HYLA improves performance over linear attention and performs elosely to softmax attention despite being a linear attention variant itself. This is noteworthy considering the importance of softmax in language modeling hypothesized be due to its role in binding and associative recall problems on which linear attention methods typically struggle (Arora et al. 2023 Olsson et al., 2022 Schlag et al. 2021). The fact that reinforcing the hypernetwork mechanism implicit in multi-hear head attention as done by HYLA helps close the gap of linear attention to softmax attention suggests that the hypernetwork mechanism might be of practical relevance for understanding large-scale transformer models.
40
Softmax attention
Linear attention
Evaluation perplexity
HYLA
35
30-
25-
50B
1008
Tokens
Figure 6: LLanguage mod- eling. Performance of au- toregressively trained decoder- only transformer models with 50M parameters over the course of training on the C4 dataset with 1 130B tokens.

---

Published as a conference paper at ICLR 2025
6 RELATED WORK
The role of multiple heads in attention. Prior work that has studied the question of why the attention mechanism benefits from multiple heads has counterintuitively found that it possible to prune almost all but one head in some layers after training, sacrificing only relatively small amounts of performance ( (Voita et al. 2019; Michel et al. 2019). It has therefore been speculated that equipping the attention mechanism with multiple heads primarily aids stability during training (Liu et al. (021). The hypernetwork perspective of attention offers an alternative account, suggesting that multiple heads create way of configuring compositional computations in the value network. Seen from this perspective, the observation of singular heads dominating could point to a connection to the phenomenon of module collapse often observed when training modular systems in practice (Shazeer ct al. 2017| Kirsch et al., 2018: Rosenbaum et al., 2018 Mittal et al. (022).
Compositional generalization. Consistent with our observation scaling model size the SRAVEN task, Hosssiiii et al. (2022) find that as pretrained large language models are scaled, their ability compositionally generalize in-context semantic parsing tasks improves. Similarly, outside the in-context learning regime, Furrer et al. (202T) have demonstrated that pretrained large language models outperform many architectures specialized towards compositional generalization. Still, to
what extent the ability for compositional generalization extends beyond the training distribution
remains openly debated (Srivastava et al. 2023 Press et al. 2023 Dziri et al. 2023
Raven-based tasks. A number of Raven-inspired tasks have been introduced to assess abstract reasoning capabilities of neural networks (Wang & Su 2013 Zhang et al., 2019, Barrett et al. 2018
Hu et al. 2021). Common to all these variants, including our version is an underlying generative task model based on finite number of possible rules that are applied on one or more of the feature dimensions for each panel. Different from our version, these variants render the problem instances as images using simple geometrical shapes resulting larger datasets and computational demand. Still, argue that the generated tasks not necessarily harder than the tasks of SRAVEN, given that SRAVEN models an additional difficulty of finding correspondences as detailed in Section 4.3 and its difficulty can be controlled parametrically; for instance by increasing the number of features.
DISCUSSION
We have proposed novel decomposition of multi-head attention, revealing that it can be interpreted as a hypemetwork that composes value networks specific to each key-query pair. Consistent with this perspective, we find empirically that the attention scores over the heads for each key-query pair form a functionally structured space, identifying reusable subfunctions in two abstract reasoning tasks. Furthermore, modifying multi-head attention to strengthen the hypernetwork mechanism improves compositional generalization on these tasks.
Adopting the hypernetwork perspective more broadly, multi-head attention can be interpreted as particular choice for the level of granularity at which the hypernetwork parameterizes the value network: The value networks are key-query specific and are followed by a pooling step that sums the key-query specific outputs over the key index. In principle, other levels of granularity are possible. For instance, the hypernetwork could parameterize a query specific value network which subsumes the aggregation over the key index. In the extreme case, it could directly parameterize the full sequence operation potentially facilitating the re-use of sequence-level operations. This also offers an interesting connection to attention- et al. 2023). Classically, a non-causal single layer transformer is seen as a fully connected graph where each token corresponds to node and aggregation is done via attention (Battaglia et al.) 2018). Our derivation suggests an alternative interpretation where the message function is a hypernetwork that subsumes the attention weights while the aggregation becomes the sum operator. Given the crucial role that aggregation plays in graph neural networks (Rosenbluth et al. 2023 Dudzik et al. 2023), a natural question is whether different pooling operators should be considered more generally for multi-head attention.
graph neural networks (Veličković et al. 2018 Shirzad
We have focused our analysis on models trained from scratch in order to give us fine -rainnd control over what data is encountered during training. Many interesting behaviors emerge in large-scale models pretrained on diverse tasks. Investigating whether the resulting models form a similarly structured latent code is interesting avenue for future work.
Limitations.
10

---

Published as a conference paper at ICLR 2025
Ethics statement. This paper conducts foundational research aiming illuminate the mechanisms of the widely used multi-head attention layer. While we foresee no immediate negative societal impact, we hope that it may improve understanding of this widely deployed technology.
Reproducibility statement. To ensure the reproducibility of our work, we are providing the code for all our experiments as part of the supplementary material. We further expand upon the experimental details explained in the main text in Appendix C
ACKNOWLEDGEMENTS
We would like to thank Angelika Steger as well as Alexander Meulemans, Johannes von Oswald, Maciej Wolczyk and Owen He for fruitful discussions throughout the development of the project. This research was supported by an Ambizione grant (7200P3_18602) from the Swiss National Science Foundation and ETH Research Grant (ETH-23 21-1).
REFERENCES
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023. URL https: httl:tttp/::pppnooee /oopnre////eeeeiee.. wwt. net/forum? forum?
1d=0g004HByN411
Shengnan An, Zeqi Lin, Qiang Fu, Bei Chen, Nanning Zheng. Jian-Guang Lou, and Dongmei Zhang. How Do In-Context Examples Affect Compositional Generalization? In Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 11027-11052, Toronto, Canada, July 2023. Association for Computational Linguisties. doi: 10. 10.188533vv1220 1/2p5/3//1222..aae 20w.0/222.2.2..o.....11111 URL
618
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and Improving Recall in Efficient Language Models, December 2023. URL httth::t:://// //axxiv.//rxiv../// org abs/ /211.. 04
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: /ithhub.gittuttu./. com/ deepmind
David Barrett, Felix Hill, Adam Santoro, Ari Morcos, and Timothy Lillicrap. Measuring abstract reasoning in neural networks. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 511-520. PMLR, July 2018. URL https:: https::/prrccee:: ngs. ml press/ 1r00/aarrettla.ht..
Peter W. Battaglia. Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Caglar Gulcehre, Francis Song, Andrew Ballard, Justin Gilmer, George Dahl, Ashish Vaswani, Kelsey Allen, Charles Nash, Victoria Langston, Chris Dyer, Nicolas Heess, Daan Wierstra, Pushmeet Kohli, Matt Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks, October 2018. URL http:ttp://axx.v.rrr//////////// arXiv:1806.0 806.01266 [cs, stat].
01261
Lukas Biewald. Experiment Tracking with Weights and Biases, 2020. com/
Alfred Binet and Theophile Simon. The development of intelligence in children. 1961. Publisher: Appleton-Century-Coon
11

---

Published as a conference paper at ICLR 2025
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
wig.....bb...... com coom/goollo. eroogle/ipg jax
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger. Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL UBLhhttpsss/prroceisssse.rrrrrre 2020/ fi//iile
cC wwtppppp.ppp.peeeeeee. files/paper pdf
P. A. Carpenter, M. A. Just, and P. Shell. What one intelligence test measures: theoretical account of the processing in the Raven Progressive Matrices Test. Psychological Review, 97 (3):404 431, July 1990. ISSN 0033-295X. URL httpst:/psyccett......mm apa. opg.er//eeeoodd
1990- 1990-27436-000 001]
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why Can GPT Learn In-Context? Language Models Secretly Perform Gradient Descent as Meta- Optimizers. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 4005-4011 Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 20031/6633/22233 http:///22....innnaall.....
findings--cll.27, URL
htt.:22.......nneeeeeeee findings-acil 247
Andrew Joseph Dudzik, Tamara von Glehn, Razvan Pascanu, and Petar Veličković. Asynchronous Algorithmic Alignment with Cocycles. In The Second Learning on Graphs Conference, 2023.
URL
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang. Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and Fate: Limits of Transformers on Com- positionality, June 2023. URLhtt / wrt. http://aa///....... 2300 18654 arXiv:2305. 18654 [cs].
Jerry A. Fodor and Zenon W. Pylyshyn. Connectionism and cognitive architeccture A critical analysis. Cognition, 28(1-2)33-71, March 1988. ISSN 00100277. doi: 0600077 10.101600000000077777/9909111111 URL
e sevier.ee.....e http::/hww..a..ea..oooo.....aaaa
Daniel Furrer, Mare van Zee, Nathan Scales, and Nathanacl Schärli. Compositional Generalization in Semantic Parsing: Pre-training vs. Specialized Architectures, September 2021. URL http:
w MtXivv2007.08970 [es].
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-2569 JMLR Workshop and Conference Proceedings, 2010.
David Ha, Andrew M. Dai, and Quoc Le. HyperNetworks. In International Conference on Learning Representations, 2017, URL httpttpps///ppeeiieeeee www.nneetoooddd-eeeeee http://ppnereeeeeeeeeeooroaiieee
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL
httt:://iithu. ccm/ wwtp::cww...ccccctot wwwgoooooogge. ffaxx
Roee Hendel, Mor Geva, and Amir Globerson. In-Context Learning Creates Task Vectors, October 2023. htRlhttp::/ http//awwwwv...aa//////a/3/ 15916 16X1 [cs].
Arian Hosseini, Ankit Vani, Dzmitry Bahdanau, Alessandro Sordoni, and Aaron Courville, On the Compositional Generalization Gap of In-Context Learnigg In Jasmijn Bastings, Yonatan Belinkov,
12

---

Published as a conference paper at ICLR 2025
Yanai Elazar, Dieuwke Hupkes. Naomi Saphra, and Sarah Wiegreffe (eds.), Proceedings of the Fifth BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP. pp. 272-280. Abu Dhabi, United Arab Emirates (Hybrid), December 2022. Association for Computational Linguistics. doi: 10. http:/33/w//222222accooooo 0022.blackboxmmpp-111
1.22. URL / org/
Sheng Hu, Yuqing Ma, Xianglong Liu, Yanlu Wei, and Shihao Bai. Stratified Rule-Aware Network for Abstract Visual Reasoning. Proceedings of the AAAI Conference on Artificial Intelligence, 35 doi: 10..1609/aai..33552.... 16248. URL https://ojs.aant aaai. org hit///ndxx...o /ttp:aa.....aaa.aaaaaaaaaaaaaaa. http:///pppa..aaaaaaaaaaa///////
(2):1577-1574, May 2021. ISSN 2374-3468, 2159-5399.
Jie Huang and Kevin Chen-Chuan Chang. Towards Reasoning in Large Language Models: A Survey. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, pp. 1049-1065, Toronto, Canada, July 2023. Association for Computational Linguisties. doi: 10. 111816////1122222...i2ngg-////// https///eeLanthoooo.ooo///////// hing:e:3...iiiing---eaa--
URL
Plotly Technologies Inc. Collaborative data science, 2015. URL https://plot.1. Place:
Montreal, QC Publisher: Plotly Technologies Inc.
Louis Kirsch, Julius Kunze, and David Barber. Modular Networks: Learning to Decompose Neu- ral Computation. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018 URL UtLBitpps./prrceeiings...uu oatpppeeeeeeccttcoroeeaeearrrrrr Pap.13-aappppppppp..
cc/ wwpppperppppppp..eeeppeeeeee... files/paper/
Taku Kudo and John Richardson. SentencePiece: simple and language independent subword tokenizer and detokenizer for Neural Text Processing. In Eduardo Blanco and Wei Lu (eds.), Pro- ceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 66-71, Brussels, Belgium, November 2018. Association for Computational Linguistics, doi: 10. 10.18853/v1/// L//D//I-22011
Brenden M. Lake and Marco Baroni. Human Mluman-liee systematic generalization through a meta- learning neural network. Nature. 623(7985); 15-121, November 2023. ISSN 1476- 4687. doi:
156--030-6--0--666------ URL /tttt::/t/w//tt http:///wwww.aau..eea.eeamm nawure. com wwtoo//rticlessss 941586-023-06668
Publisher: Nature Publishing Group.
Liyan Liu, Jialu Liu, and Jiawei Han. Multi-head Single-head? An Empirical Compari- son for Transformer Training, June 2021. URL httt:::/arxii..og/aas///// 09650 arXiv:2106.09650 [cs].
Peter J. Liu, Roman Novak, Jachoon Lee, Mitchell Wortsman, Lechao Xiao, Katie Everett, Alexan- der A. Alemi, Mark Kurzeja, Pierre Marcenac, Izzeddin Gur, Simon Kornblith, Kelvin Xu, Gamaleldin Elsayed, Ian Fischer, Jeffrey Pennington, Ben Adlam, and Jascha-Sohl Dickstein. NanoDO: minimal Transformer decoder- language model implementation in JAX., 2024.
URL htL::tt::// thub.
ind nanoood
Ilya Loshchilov and Frank Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, May C3983 arXiv:1608.cm N0w:1668.03983 03983 [cs, math].
2017. hOtp:U///LEEp///aaatoooo////////
Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. January 2019. URL arXiv:171i.005 1.01011 [cs, math].
http::/arxiv.cn http://v.wogg/a///1.1./
Laurens van der Maaten and Geoffrey Hinton. Visualizing Data using t-SNE. Journal of Machine Learning Research, 986))2579-2205, 2008. ISSN 1533-7928. URL http:://mmr: org/
aaperss/t//vanderraateeeet//
html
Paul Michel, Omer and and Graham Neubig, Are Sixteen Heads Really Better than One?, November 2019. U
10650 arXiv:1905.1066 [cs].
Sarthak Mittal, Yoshua Bengio, and Guillaume Lajoie. Is a Modular Architecture Enough? In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https:://oppeeee.
net htt./oorrmmmom
3- 3XMModt.xx
13

---

Published as a conference paper at ICLR 2025
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds Danny Heraandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-Context Learning and Induction Heads, September 2022. URL http:
11895 arXiv:2209.1 1895 [cs].
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A. Smith, and Mike Lewis. Measuring and Narrowing the Compositionality Gap Language Models, May 2023. URL http:///rxxxii or org/ 2210. 03350 arXiv:2210.03350 [es].
Colin Raffel, Noam Shazeer, Adam Roberts. Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. Journal of Machine Learning Research, 21(140):1667, 2020. URL
Nasim Rahaman, Muhammad Waleed Gondal. Shruti Joshi, Peter Vincent Gehler, Yoshua Bengio, Francesco Locatello, and Bernhard Schölkopf. Dynamic Inference with Neural Interpreters. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL
net nwt/ferummfdiie
IUJt25DtqC4
John C. Raven. Advanced Progressive Matrices, Set II. H. K. Lewis, London, 1962.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing Networks: Adaptive Selection of Non-Linear Functions for Multi- Learning. In International Conference on Learning Representations, 2018
Eran Rosenbluth, Jan Toenshoff, and Martin Grohe Some Might Say All You Need Is Sum, May 2023. (tRt:ttt://aaatt
1xXiv:2302.11600 [cs].
David E Rumelhart and James L McClelland. On learning the past tenses of English verbs. 1986. Publisher: Cambridge, MA: MIT Press.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers, June 2021. URL http::/t///aaa....cr//////221111 [cs].
01111744riiv22221 111
Jürgen Schmidhuber. Learning to Control Fast-Weight Memories: An Alternative to Dynamic Recurrent Networks. Neural Computation, 4(1):131-139 1992. doi: 199.2/4....99..4.....
Simon Schug, Seijin Kobayashi, Yassir Akram, Maciej Wolczykk, Alexandra Proca, Johannes von Os- wald, Razvan Pascanu, João Sacramento, and Angelika Steger. Discovering modular solutions that generalize compositionally. In The Twelfth International Conference on Learning Representations, 2024. URL
n wwtp:eeeeoooo.aaaaaaaooe...
Noam Shazeer, * AAzalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely- Mixture-of-Experts Layer. In International Conference on Learning Representations, 2017. URL htttps://openreee
net / Eorummidd ckMDqlg
Hamed Shirzad, Ameya Velingker, Balaji Venkatachalam. Danica J. Sutherland, and Ali Kemal Sinop. Exphormer: Sparse Transformers for Graphs. In Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 31613-31632. PMLR, July 2023.
Paul Smolensky. Connectionism, Constituency and the Language of Thought. In Barry M. Loewer and Georges Rey (eds.), Meaning in Mind: Fodor and His Critics. Blackwell, 1991.
Aarohi Srivastava, Abhinay Rastogi. Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R. Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonso, Agnieszka Kluska, Aitor Lewkowyez, Akshat Agarwal, Alethea Power, Alex Ray, Alex Warstadt, Alexander W. Kocurek, Ali Safaya, Ali Tazary, Alice Xiang, Alicia Parrish, Allen Nie, Aman Hussain, Amanda
14

---

Published as a conference paper at ICLR 2025
Askell, Amanda Dsouza, Ambrose Slone. Ameet Rahane, Anantharaman S. Iyer, Anders Johan Andreassen, Andrea Madotto, Andrea Santilli, Andreas Stublmüller, Andrew M. Dai, Andrew La, Andrew Lampinen, Andy Zou, Angela Jiang, Angelica Chen, Anh Vuong, Animesh Gupta, Anna Gottardi, Antonio Norelli, Anu Venkatesh, Arash Gholamidavoodi, Arfa Tabassum, Arul Menezes, Arun Kirubarajan, Asher Mullokandov, Ashish Sabharwal, Austin Herrick, Avia Efrat, Aykut Erdem, Ayla Karakaş. B. Ryan Roberts, Bao Sheng Loe, Barret Zoph, Bart'lomiej Bojanowski, Batuhan Özyurt, Behnam Hedayatmia, Behnam Neyshabur, Benjamin Inden, Benno Stein, Berk Ekmekci, Bill Yuchen Lin, Blake Howald, Bryan Orinion, Cameron Diao, Cameron Dour, Catherine Stinson, Cedrick Argueta, Cesar Ferri, Chandan Singh, Charles Rathkopf, Chenlin Meng, Chitta Baral, Chiyu Wu, Chris Callison-Burch, Christopher Waites, Christian Voigt, Christopher D. Manning, Christopher Potts, Cindy Ramirez, Clara E. Rivera, Clemencia Siro, Colin Raffel, Courtney Ashcraft, Cristina Garbacea, Damien Sileo, Dan Garrette, Dan Hendrycks, Dan Kilman, Dan Roth, C. Daniel Freeman, Daniel Khashabi, Daniel Levy, Daniel Moseguí González, Danielle Perszyk, Danny Hernandez. Danqi Chen, Daphne Ippolito. Dar Gilboa, David Dohan, David Drakard, David Jurgens, Debajyoti Datta, Deep Ganguli, Denis Emelin, Denis Kleyko, Deniz Yuret, Derek Chen, Derek Tam, Dieuwke Hupkes, Diganta Misra, Dilyar Buzan, Dimitri Coelho Mollo, Diyi Yang. Dong-Ho Lee, Dylan Schrader, Ekaterina Shutova, Ekin Dogus Cubuk, Elad Segal, Eleanor Hagerman, Elizabeth Barnes, Elizabeth Donoway, Ellie Pavlick, Emanuele Rodolå, Emma Lam, Eric Chu, Eric Tang, Erkut Erdem. Ernie Chang, Ethan A. Chi. Ethan Dyer, Ethan Jerzak, Ethan Kim, Eunice Engefu Manyasi, Evgenii Zheltonozhskii, Fanyue Xia, Fatemeh Siar, Fernando Martúnez-Plumed, Francesca Happé, Francois Chollet, Frieda Rong, Gaurav Mishra, Genta Indra Winata, Gerard de Melo, Germán Kruszewski, Giambatrista Parascandolo, Giorgio Mariani, Gloria Xinyue Wang, Gonzalo Jaimovitch-Loppe Gregor Betz, Guy Gur-Ari, Hana Galijasevic, Hannah Kim, Hannah Rashkin, Hannanch Hajishirzi, Harsh Mehta, Hayden Bogar, Henry Francis Anthony Shevlin, Hinrich Schuetze, Hiromu Yakura, Hongming Zhang, Hugh Mee Wong, Ian Ng, Isaac Noble, Jaap Jumelet, Jack Geissinger, Jackson Kernion, Jacob Hilton, Jaehoon Lee, Jaime Fernández Fisac, James B. Simon, James Koppel, James Zheng, James Zou, Jan Kocon, Jana Thompson, Janelle Wingfield, Jared Kaplan, Jarema Radom, Jascha Sohl-Dickstein, Jason Phang, Jason Wei, Jason Yosinski, Jekaterina Novikova, Jelle Bosscher, Jennifer Marsh, Jeremy Kim, Jeroen Taal, Jesse Engel. Jesujoba Alabi, Jiacheng Xu, Jiaming Song, Jillian Tang, Joan Waweru, John Burden, John Miller, John U. Balis, Jonathan Batchelder, Jonathan Berant, Jörg Frohberg, Jos Rozen, Jose Hernandez-Orallo, Joseph Boudeman, Joseph Guerr, Joseph Jones, Joshua Tenenbaum, Joshua S. Rule, Joyce Chua, Kamil Kanclerz, Karen Livescu, Karl Krauth, Karthik Gopalakrishnan, Katerina Ignatyeva, Katja Markert, Kaustubh Dhole, Kevin Gimpel, Kevin Omondi, Kory Wallace Mathewson, Kristen Chiafullo, Ksenia Shkaruta, Kumar Shridhar, Kyle McDonell, Kyle Richardson, Laria Reynolds, Leo Gao. Li Zhang, Liam Dugan, Lianhui Qin, Lidia Contreras-Ochandae Louis-Philippe Morency, Luca Moschella, Lucas Lam, Lucy Noble, Ludwig Schmidt, Luheng He, Luis Colvéros-Colón, Luke Metz. Lütfi Kerem Senel, Maarten Bosma, Maarten Sap, Maartje Ter Hoeve, Maheen Farooqi. Manaal Faruqui, Mantas Mazeika, Marco Baturan, Marco Marelli, Marco Maru, Maria Jose Samirez-Quintane Marie Tolkiehn, Mario Giulianelli, Martha Lewis, Martin Potthast, Matthew L. Leavitt, Matthias Hagen, Mátýás Schubert, Medina Orduna Baitemirova, Melody Arnaud, Melvin McElrath, Michael Andrew Yee, Michael Cohen, Michael Gu, Michael Ivanitskiy, Michael Starritt, Michael Strube, Michal Swedrowski, Michele Bevilacqua, Michihiro Yasunaga, Mihir Kale, Mike Cain. Minee Xu. Mirac Suzgun, Mitch Walker, Mo Tiwari, Mohit Bansal, Moin Aminnaseri, Mor Geva, Mozhdeh Gheini, Mukund Varma T. Nanyun Peng. Nathan Andrew Chi. Nayeon Lee, Neta Gur-Ari Krakover, Nicholas Cameron, Nicholas Roberts, Nick Doiron, Nicole Martinez, Nikita Nangia, Niklas Deckers, Niklas Muennighoff, Nitish Shirish Keskar, Niveditha S. Iyer, Noah Constant, Noah Fiedel, Nuan Wen, Oliver Zhang, Omar Agha, Omar Elbaghdadi, Omer Levy, Owain Pablo Antonio Moreno Casares, Parth Doshi, Pascale Fung, Paul Pu Liang, Paul Vicol, Pegah Alipcormolabahhi, Peiyuan Liao, Percy Liang, Peter W. Chang, Peter Eckersley, Phu Mon Htut, Pinyu Hwang, Piotr Milkowski, Piyush Patil, Pouya Pezeshkpour, Priti Oli, Qiaozhu Mei, Qing Lyu, Qinlang Chen, Rabin Banjade, Rachel Etta Rudolph, Raefer Gabriel, Rahel Habacker, Ramon Risco, Raphaël Millière, Rhythm Garg, Richard Barnes, Rif A. Saurous, Riku Arakawa, Robbe Raymaekers, Robert Frank, Rohan Sikand, Roman Novak, Roman Sitelew, Ronan Le Bras, Rosanne Liu, Rowan Jacobs, Rui Zhang, Russ Salakhutdinov, Ryan Andrew Chi, Seungjac Ryan Lec, Ryan Stovall, Ryan Techan, Rylan Yang, Sahib Singh, Saif M. Mohammad, Sajant Anand, Sam Dillavou, Sam Shleifer, Sam Wiseman, Samuel Gruetter, Samuel R. Bowman, Samuel Stern Schoenholz, Sanghyun Han, Sanjeev Kwatra, Sarah A. Rous, Sarik Ghazarian, Sayan Ghosh, Sean Casey,
Evans,
15

---

Published as a conference paper at ICLR 2025
Sebastian Bischoff, Sebastian Gehrmann, Sebastian Schuster, Sepideh Sadeghi, Shadi Hamdan, Sharon Zhou, Shashank Srivastava, Sherry Shi, Shikhar Singh, Shima Asaadi, Shixiang Shane Gu. Shubh Pachchigar, Shubham Toshniwal, Shyam Upadhyay, Shyamolima Shammie Debnath, Siamak Shakeri, Simon Thormeyer, Simone Melzi, Siva Reddy, Sneha Priscilla Makini, Soo- Hwan Lee, Spencer Torene, Sriharsha Hatwar, Stanislas Dehaene, Stefan Divic, Stefano Ermon, Stella Biderman, Stephanie Lin, Stephen Prasad, Steven Piantadosi, Stuart Shieber, Summer Misherghi, Svetlana Kiritchenko, Swaroop Mishra, Tal Linzen, Tal Schuster, Tao Li, Tao Yu, Tariq Ali, Tatsunori Hashimoto, Te-Lin Wu, Théo Desbordes, Theodore Rothschild, Thomas Phan, Tianle Wang, Tiberius Nkinyili, Timo Schick, Timofei Kornev, Titus Tunduny, Tobias Gerstenberg, Trenton Chang, Trishala Neeraj, Tushar Khot, Tyler Shultz, Uri Shaham, Vedant Misra, Vera Demberg, Victoria Nyamai, Vikas Raunak, Vinay Venkatesh Ramasesh, vinay uday prabhu, Vishakh Padmakumar, Vivek Srikumar, William Fedus, William Saunders, William Zhang, Wout Vossen, Xiang Ren, Xiaoyu Tong, Xinran Zhao. Xinyi Wu, Xudong Shen, Yadollah Yaghoobzadeh, Yair Lakretz, Yangqiu Song, Yasaman Bahri, Yejin Choi, Yichi Yang. Yiding Hao, Yifu Chen, Yonatan Belinkov, Yu Hou, Yufang Hou, Yuntao Bai, Zachary Seid, Zhuoye Zhao, Zijian Wang. Zijie J. Wang, Zirui Wang, and Ziyi Wu. Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL http:/ts://ppeereriiieeeee// www.:/netrfoouruu furum http:///ooouuuiiiSSSSSSSo d=
net
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. RoFormer: Enhanced transformer with Rotary Position Embedding. Neurocomputingg 568: 127063, February 2024. ISSN 0925-2312. doi: 10.
2023. 27063. URL https: www sci encedi com/ wwssscieencc.. www.encc/arcccllle/1////////
Eric Todd, Millicent Li, Amab Sen Sharma, Aaron Mueller, Byron C. Wallace, and David Bau. Function Vectors in Large Language Models. In The Twelfih International Conference Learning Representations, 2024. URL httptt/ss////ppeeeeiiiie
net/ http://frrmmiidi do yMwaG
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http://///pppprrsssss.....//o/p files wwt.e/////app/////////aaaa/ 3 35554445777699111003cccaaaaaaaa Awstract..a........... .htl
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph Attention Networks. In International Conference Learning Representations, 2018.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. In Anna Korhonen, David Traum, and Lluís Mårquez (eds.). Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics. pp. 5797-5808, Florence, Italy, July 2019. Association URL https: httpt::://aaaanthlllloo
for Computational Linguistics. doi: 10. 10.1 16553/
org/ PP1 1580
Johannes Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agiler Arcas, Max Vladymyrov, Razvan Pascanu, and João Sacramento. Uncovering mesa-optimizattion algorithms in Transformers, September 2023. ULL_pttp://arav..orr///////2...2 0585 arXiv:2309.05858 [cs].
Ke Wang and Zhendong Su. Automatic generation of Raven's progressive matrices. In Proceedings of the 24th International Conference on Artificial Intelligence, IJCAI'15, pp. 903-909, Buenos Aires.
Argentina, July 2015. AAAI Press. ISBN 978-1-57735-738-4 URL https:://///. htops:///www.. jjcai
orq Drc//rroceeaingg/////eeerr
pdf
David Wechsler. The measurement and appraisal of adult intelligence. Academic Medicine, 33(9): 706, 1958. Publisher: LWW.
L. A. Zadeh. Fuzzy sets. Information and Control, 8(3):33 June 1965. ISSN 0019-9958. doi: 10. 10.1016/500019999586659990008009 wtt.:tssse.....cceeeeeeeeee..... science/ F
16

---

Published as a conference paper at ICLR 2025
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Chi Zhang, Feng Gao, Baoxiong Jia, Yixin Zhu, and Song-Chun Zhu. RAVEN: A Dataset for Relational and Analogical Visual REasoNing. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 5312-5322, 2019. doi: 10.1 201.///C....1..0000.
17

---

Published as a conference paper at ICLR 2025
Algorithm 1 Multi-head softmax attention
Algorithm 2 Hypernetwork linear attention
Input:
Input: x Require: emb_dim, beads, h_dim
Require: emb_dim, heads, h_dim
1: Denseishape=theads, h_dii), ax
1: g Denselshppestzeade. h dim), axis=-1)(2)
2: k Denselshape=(hhaadd. h_dim), axis== 1)(x)
2: Dense d_dim,, axis=-
3: Punscchupe-theadd..__d uxis=-1Xx)
3: Densessuppe dim). axis=-1Xx)
4:
4:
5: einsum(" ...amd............. .hhd.................. h.g...
5: a cinsum(" ....gaad..... wwd->....hhh............ ".jx".g..
6. a softmax(a, axis= 1)
6: 0 rms norm(a, axis= 3)
7:
8:
cinsum(" .....hgk, .e.k,....k.d.>.... whd.
2 cinsum(" ... hhg.....ddd...qqqdd 0,
9:
relu(2)
10:
einsum( ... .g.. h.............d............ gkd->..kdd...... gkd .hd"., a,
11:
Denseffentrrseeemb_dl axis=(-2, -1))=)
11: y Dense(feutureseebb dim, (xis=
12:
12:
13: return
13: return
Figure A1: Pseudocode comparing multi-head softmax attention to hypernetwork linear attention. Differences between the two are highlighted in yellow.
A ADDITIONAL RESULTS
A.1 MODEL ABLATIONS
Compared to linear attention, HYLA introduces three modifications: (1) normalization of the attention scores across the head indices using RMSHead, (2) applying the attention- weighted operation also to the output projection and (3) inserting a nonlinearity the value network (c.f. the pseudocode for HYLA shown in Figure [ To delineate the contributions of these modifications, we run an ablation study on the fuzzy logic task and SRAVEN shown in Table AT We find that while simply applying RMSHead linear attention (Linear Attention RMSHead) and also the attention attention--eighted weighted sum operation without nonlinearity (HYLA nonlinearity) both by themselves can slightly boost performance, the full HYLAA model generally performs best.
Increasing depth of the value network. To further test our hypothesis that it is the hyperneworr mechanism inside multi-head attention which improves compositional generalization, we explore making the value network configured by the hypernetwork deeper, adding additional layer as follows.
)) (10)
HYLA (X)
Σ
Σ ah.gk Wont
W
Σ On.g.
Wratic x
ah.g.&
where W E R 12 xike is second layer value matrix which increases the number of parameters of this model. Consistent with our hypothesis, this model also demonstrates strong OOD performance as reported in Table A1
Replacing RMSHead normalization with the softmax. We perform an additional ablation where replace the RMSHead normalization of HYLA with the standard softmax that normalizes across keys but not across heads. As shown Table Al this model performs poorly.
Replacing feedforward layers with mixture of experts. Sparsely gated mixture of experts (MoE) have been proposed as a replacement for the feedforward layer in the transformer block (Shazeer al. 2017). In this model, each token is routed to sparse subset of out of E experts
E
fMOE(*) Σ Softmax(Top k( w)))))) x)) W.i of W down
Intuitively, this could allow individual experts to specialize on specific subtasks that can be composed by the router. However, it is not well understood if MoEs improve compositional generalization of
18

---

Published as a conference paper at ICLR 2025
transformers in practice. Here, we perform additional experiments where replace the feedforward layers with MoEs on the SRAVEN and the fuzzy logic task reported Table While we find moderate general improvements, these improvements seem to be similar to other increases of model capacity such as increasing the width/depth as discussed in Section 4.4
Table A1: Model ablations. Left column reports the OOD R2 score on the fuzzy logic task studied in Figure for sequence length of 32 and 70% of tasks held from training. Right column reports the OOD accuracy on the SRAVEN task studied in Figure for 4-layer transformer trained on 20M problem instances. We report the standard error over 3 seeds.
Fuzzy logic:
SRAVEN:
OOD R2
OOD accuracy
Softmax attention
63.28 + 2.31
56.56 + 1.05
Linear attention
59.89 5.22
56.30 1.11
HYLA
81.13 + 7.77
69.13 = 1.90
Linear attention + RMSHead
68.93 5.10
59.01 2.50
Linear attention + RMSHead + nonlinear value
56.91 + 4.25
55.38 + 0.32
HYLA RMSHead
81.27 7.88
62.53 + 5.73
HYLA nonlinearity
84.54 5.32
56.54 + 0.21
HYLA nonlinearity RMSHead
79.42 7.43
56.33 + 1.41
HYLA RMSHead + softmax
70.98 +: 7.08
38.21 + 2.98
HYLA deep value network
87.65 + 4.54
66.59 1 0.08
Softmax attention MoE
63.83 6.07
57.68 + 1.31
Linear attention + MoE
60.09 + 6.16
44.4443
HYLA + MoE
82.37 6.55
66.94 3.94
19

---

Published as a conference paper at ICLR 2025
B
Target
Softmak
HML
25
100
1C0-
-03
1000
100-
422
90-
-50
-50-
-103
10
110
-100
100
120
102
100
100
-000
100
100
100
100
100
C
Softmae ottonion
Unnor attoniinn
HNLA
Term
Softnax attention
Lincar astontion
HBA
100-
-100-
-100-
100-
1004
-50
-100
-100
-102
100
Figure A2: tSNE visualization the attention scores across the head index for the query token on the fuzzy logic task for sequence length of 32 and 70% of tasks held out from training colored according to (A) target label (B) fuzzy logic function, (C) term and (D) term 2.
A
B
Softmae attertion
Uhwer atlention
IMA
Soltire attention
Linee ailent on
IIAAA
000 fracion0.3
C00 fLaction.s
DOD frection 0.7
000
fraction..3
0C0 fraction 0.5
DOD fraction 0.7
R2
0.6
0.6
0.4
C.4
16
96
15
32
95.
16
06
12
95
Sequence length
Sequence length
Sequence longth
Sequence length
Sequence length
Sequence length
C
Sofimas attention
Limser atcention
ITLA
Seflmee attention
Unmer attantion
1YLA
000 fracsion 0.3
DCD froction 0.5
000 fraztion 0.7
COD fraction 0.3
0C0 froction 0.5
000 fraction 3.7
R2 on tooks with
unisee tarns
Traininglos
100
05
-22
26
22
10
Sequence
lnggth
Sequence langth
Secuence length
Sequenoe longth
Sequencee length
Sequence length
Figure A3: Fuzzy logic task performance metries. Various task performance metrics on the fuzzy logic task for varying number of in- in-context examples and fraction of tasks held-out during training. Coefficient of determination (R²) on tasks consisting of unseen combinations of seen terms (same as Figure 2C) B Coefficient of determination (R²) training tasks. C Coefficient of determination (R2) on tasks containing terms not encountered during training. As expected all methods fail in this setting. D Training mean squared error loss.
20

---

Published as a conference paper at ICLR 2025
B
Softmar athertion
Uneat sttention
IMIA
Sottman attention
Linter artention
IMA
Ta widh
25 witth
4x medh
4 Layers
8 Blayers
15 laprrs
Traiing loos
6.3
Traiiig
022
6.1-
EEEEE Problem instarcees
Problem
Probleminsiances
Problem instances
Problem
Probiem instarcee
Figure A4: Training loss when scaling data and model size on SRAVEN. A Training loss as function of the number of training problem instances for different widths (scaling embedding and key-query-valle value dimensions). B Same as A but increasing depth instead of width.
Softmax attention
Linear attention
0.9
HYLA
0OD accuracy
0.8
0.7-
0.6
0.5
0.1
0.3
0.5 0.7 0.9
Fraction of combinations held-out
Figure A5: Varying the fraction of
OOD tasks on SRAVEN. OOD accuracy on SRAVEN as varying fractions of tasks specified by their rule combinations are held-out during training.
A
B
Input
Rule
attertiinn
Urrer attertion
attertattertion
Uine aftention
102-
-90
-50-
50-
52-
100
100-
-50
0-
-02
-00
-30
50
-50
50
Figure A6: Latent code on SRAVEN. A tSNE visualization of the latent code predicting the final query token colored by the magnitude of the predicted target value for 4-layer transformers trained on 40M problem instances. B Same as A but colored by rule.
21

---

Published as a conference paper at ICLR 2025
A
Rule
Layer
Layer
Layer
Layer
X2
-50
Layer
Layer
Layer 6
Layer
50
-50
Layer
Layer 9
Layer 10
Layer 11
50
x2
-50
Layer 12
Layer 13
Layer 14
Layer 15
50
x2
-50
-50 0
50 100
-50 50100
-50 50.000
-50 0 50100
x1
x1
B
Input
Layer
Layer
Layer 2
Layer 3
50
-50
Layer
Layer 5
Layer 6
Layer
-50
Layer
Layer
10
Layer 11
X2
-50
Layer
Layer 13
Layer 14
Layer
-50
-50 0 50 100
-50 50100
-50 50100
-50 0 50100
x1
x1
x1
x1
Figure A7: Latent code visualization in SRAVEN for a softmax transformer of depth 16. ISNE visualization of the latent code predicting the final query token colored by the rule to inferred for softmax transformer of depth 16 trained on 40M problem instances. B Same plot as in A but colored by the magnitude of the predicted target value.
22

---

Published as a conference paper at ICLR 2025
B
Average
Softmax attention
silention
Smilarit
[eead
Foad
--55
-45
15+
Rule
Example
0-11
Average silention
Attention
Linear attention
Smilanit
04
Hiead
-05
Foad
10
--02
15-
19
Rule
Example
0-
Average
Attention
ailention
Smilanity
HYLA
Head
Foad
15+
19
Rule
Example
Figure A8: Detailed attention scores on SRAVEN. Random sample of the attention scores / latent codes for the final key-query pair on OOD SRAVEN tasks in the final layer of transformer using softmax attention, linear attention or HYLA. B Average attention scores latent codes per rule. C Pair-wise cosine similarities between all average latent codes per rule, Data from the same models as shown in Figure A6
23

---

Published as a conference paper at ICLR 2025
B SRAVEN
B.1 NUMBER OF POSSIBLE TASKS AND INSTANCES
In the general case, the number of possible tasks in SRAVEN is a product of the number of possible
rule combinations ( (R+ 11+K--) ()) and the number of possible unique permutations ((K!)--1 K) where
K
G' is the number of columns of the grid (in Raven-based tasks typically G 3). Figure A9A shows
the number of possible SRAVEN tasks and Figure|A9p shows the number possible SRAVEN problem
instances. Note specifically that the latter is orders of magnitudes larger than the number of training examples our experiments for K 4, 8.
B.2 AMBIGUITY
In principle, problem instance generated given set of rules and permutations might have ambiguous answer. That is, there is at least one other set of rules and permutations that fit all context panels but prescribe different final query panel. One possibility is check for such examples when generating task instances and reject ambiguous ones. This process is prohibitively expensive, as it requires a search over the exponentially many possible permutations. We therefore resort to Monte-Carlo estimation of the fraction of ambiguous instances. Table various parameterizations of the SRAVEN task. In the setting we consider here, (K F 8), the probabilities are negligible, and we do not take measures to filter out ambiguous examples given the computational overhead it produces the data generation process.
the results for
B.3 RULES
sraven contains the following rules:
1. constant: Each row consists of a random but fixed integer in the range 1 to F.
2. progression (+1): The first element of each row is sampled uniformly at random and incremented by modulo for each successive column.
3. progression (+2): The first element of each is sampled uniformly at random and incremented by 2 modulo F for each successive column.
progression (-1): The first element each row is sampled uniformly at random and decremented by modulo F for each successive column.
5. progression (-2): The first element of each row is sampled uniformly at random and decremented by modulo F for each successive column.
addit Two elements sampled uniformly random for each row and added modulo F to obtain the last column.
subtractions Two elements are sampled uniformly at random for each row and subtracted modulo F to obtain the last column.
8. distribute three: Three elements sampled uniformly at random and presented in three independently sampled random permutations for each row.
24

---

Published as a conference paper at ICLR 2025
A
B
Number features - - 6
1T-
70
Number of tasks
Number instances
10
18-
104s
1M
10
1015
1000
Number of features
Number of feature values
Figure A9: A Number of possible SRAVEN tasks as function of the number of features (K). B Number of possible SRAVEN problem instances function of the number of features (K) and the number of feature values (F)).
Table A2: Monte Carlo estimation of fraction of ambiguous examples over 4096 problem instances standard error. Bold used to denote the setting used for all experiments in this paper.
Number of features (K)
Number of feature values (F)
Estimated fraction of ambiguous examples
4
0.0642 + 0.0038
4
8
0.0032 + 0.0009
4
16
0.0005 + 0.0003
25

---

Published as a conference paper at ICLR 2025
C EXPERIMENTAL DETAILS
Architecture. For all models we use standard decoder-only transformer architecture where each block structured
Z hut::Headdttenntioo((yyyattttt (X)) X
Y Feedforraard(Laaertoorrf/
MultiHeadhttention is either linear attention, softmax attention or hypemetwork linear at- tention (HYLA) and uses 15-style relative positional embeddings (Raffel et al. 2020). For the seedForward layer, we employ a standard single-hidden layer MLP with GeLU nonlinearity applied to each position in the sequence independently. For our language modeling experiments we use rotary positional encodings (Su et al. 2024).
Tokenization. For the fuzzy logic task directly the concatenated input examples tokens. Similarly for SRAVEN we use one-hot encodings of the integer inputs as the input tokens. In both cases the input tokens mapped into the embedding space of the transformer using fully connected dense layer and another fully connect dense layer is used to produce the output logits. For the use the default sentencepiece tokenizer (Kudo & Richardson 2018) of NanoDo (Liu et al., 2024) with a vocabulary size of 32000 and use the transposed embedding matrix to produce the logits.
language modeling experiment
Optimization. We the AdamW optimizer (Loshchilov & Hutter 2019) with linear warmup starting from a learning rate of followed by cosine decay learning rate scheduler (Loshchilov Hutter 2017) that decays the learning rate 0.1 times the base learning rate. We exempt biases and LayerNore orm parameters from weight decay. the fuzzy logic task we use the mean- squared eтror on the logits of the query token whereas in SRAVEN use the softmax cross-entropy loss on the logits of all query tokens. The language modeling task uses an autoregressive softmax cross-entropy loss with causal masking applied to the attention mechanism well mixed- sexeedprrecsion training.
Hyperparameters. For all tasks and models we perform grid search over the learning rate, weight decay and warmup steps. We report the search grid as well all other hyperparameters Table A3
Table A3: Hyperparameters for all tasks and methods. Lists of values indicates that grid search over all points contained in the list have been conducted and for each method the optimal combination has been picked. The width factor multiplies the embedding dimension, key-query-vanl -aauu dimension and MLP dimension.
Parameter
Fuzzy logic
SRAVEN
Language modeling
batch_size
128
128
256
num_layers
2
4/ 8/8/ 1/8/10 16
_ictt_
1/2/4
emb_dim
128
128
512
kqv_dim
16
64
64
mlp_dim
256
256
2048
num_heam
16
learning_raan rate
00.001,00008
[0.0003, 0.001]
[0.001, 0.003]
weight_decay
[0.03,
0.1,
[0.1, 0.03]
warmup_steps
100
11000, 3000]
1000
26

---

Published as a conference paper at ICLR 2025
D ADDITIONAL DETAILS
D.1 COMPUTE RESOURCES
We used Linux workstation with Nvidia RTX 3090 GPUs with 24GB memory each for development and conducted hyperparameter searches and experiments using Linux server with Nvidia RTX 3090 GPUs as well as a Slurm cluster equipped with Nvidia RTX 4090 GPUs. With single GPU of these two GPU models, single run of the fuzzy logic task takes around 2 2-3 minutes, a single run on SRAVEN between
200 minutes. For the language modeling experiments, we used 16 Cloud TPU v5e with complete run taking 72-100 hours. In total, it takes around 24 GPU hours to reproduce all fuzzy logic results on our hardware, around 10 GPU days to reproduce all our SRAVEN results and approximately 163 TPU days to reproduce results on C4.
D.2 SOFTWARE AND LIBRARIES
For the results obtained in this paper we built on free and open-souree software. We implemented our experiments in Python using JAX (Bradbury et al. 2018 Apache License 2.0), Flax (Heek et al.)
2023 Apache License 2.0), NanoDo (Liu et al. 2024 Apache License 2.0) and the Deepmind Jax
Ecosystem (Babuschkin et al. 2020 Apache License 2.0). We utilized WandB (Biewald 2020) MIT license) to monitor the progress and results of experiments, and Plotly (Inc) 2015 MIT license) for generating the plots.
27


Reviews:
[{"review_id": "qC5BnYu53b", "reviewer": "Reviewer_ZLCD", "summary": "This work reformulates linear multi-head attention as a hypernetwork to study compositional generalisation and introduces a simple non-linear variant that substantially increases generalisation both on abstract symbolic tasks as well as language modeling.", "questions": "- In HYLA, you not only add the nonlinearity but also change the normalisation. Did you perform ablations w.r.t. the normalization, meaning how well does HYLA perform with the standard normalization (or vice versa, if you add RMSHead in linear attention)?\n\n- Please report the training performance of the models to shed light on the generalisation vs capacity question.\n\n- Please provide more details on the OOD splitting of the datasets.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and studies an important problem, namely compositional generalisation, from a new angle and with promising empirical results. The writing, structure and related work is outlined clearly, and (almost) all relevant details with regards to the theory and the empirical experiments are available in either the main text or the appendix. The perspective to see multi-head attention as a hypernetwork is novel and might open interesting future avenues for study.", "weaknesses": "One problem in the empirical evaluation is a certain lack of control over whether the suggested modification (HYLA) is actually increasing compositional generalization or whether it's just a way of increasing the capacity of the model. To this end, it would be interesting to look at the training loss of the models: does HYLA reach the same minimum as linear / softmax attention, or is it already performing much better on the training set (suggesting it is more the capacity increase of the model rather than a genuine way of increasing compositional generalisation abilities).\n\nSimilarly, it didn't get clear to me how the OOD sets were sampled? E.g., in the fuzzy logic dataset, when you say \"hold out a fraction of the combinations for OOD\" (line 206), does this refer to combinations of K terms, or are all K terms sampled and you just subsample each (meaning, in the OOD set all combinations of K are present)?", "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": "fGdZo7U6Hx", "reviewer": "Reviewer_59h3", "summary": "Transforms have been shown to compositionally generalize in some situations and not in others. This paper presents a novel perspective of multi-head self-attention, a core component of transformers, which better exposes the module's capacity for representing and reusing compositional structure. \n\nUnder this perspective:\n- The multi-head self-attention mechanism comprises of $\\\\text{num\\\\_queries}\\\\times\\\\text{num\\\\_keys}$ value networks that are generated by a hypernetwork (hypernet) parameterized by the value and output projection weights\n- The inputs of the hypernet are $\\\\text{num\\\\_heads}$-dimensional \"latent codes,\" each corresponding to a query-key pair. Relating to standard multi-head self-attention, these latent codes are obtained by indexing the attention matrix stacked for all heads at a specific key-value index.\n\nWith this perspective:\n- The authors show that the latent codes exhibit structure that reflects the learned compositionality for the task, which can explain why transformers compositionally generalize when they do. For instance, the latent codes are empirically predictive of the subtasks the network performs on unseen compositions.\n- The paper introduces the hypernetwork linear attention (HYLA), a drop-in replacement of standard multi-head self-attention, that modifies the hypernet-generated value network from linear to non-linear to make it more expressive. This modification also replaces the expensive global softmax normalizations per head with local key-query-specific normalizations across heads.\n\nThe authors introduce SRAVEN, a symbolic version of Raven Progressive Matrices human intelligence test, which gives them more precise control over the compositions during training and evaluation. Evaluation on SRAVEN shows clearer latent code structure and improved OOD accuracy of HYLA over softmax and linear attention.\n\nFinally, the authors evaluate HYLA on a language modeling task, but it fails to outperform softmax attention here.", "questions": "HYLA with a single head:\n- In the caption for Figure 5C, the authors state that the hypernet mechanism is absent for a single head. Why is this the case?\n\nScaling compositional difficulty:\n- In Section 4.2, the authors state that 25% of possible rule combinations are held out for evaluation. How does performance change when this fraction is changed? Specifically, as the compositional difficulty (% of held-out rules) varies from 0% to 100% on the X-axis, how do the performances of softmax attention, linear attention, and HYLA on the Y-axis vary? This important experiment will help us understand how these methods compare in a fully in-distribution setting to much harder OOD settings.\n\nDrop in performance of HYLA with softmax:\n- For HYLA, in Eq (6), both the value network weights $\\\\mathbf{W}'\\_{q,k}$ and $\\\\mathbf{W}\\_{q,k}$ depend on $a\\_{h,q,k}$, which makes sense from the hypernet perspective where $a$ is the latent code. In standard MHA, as shown in Eq (2), $\\\\mathbf{W}^\\\\text{out}$ operates after attention scaling has already been applied. However, from this perspective, Eq (6) indicates a re-application of attention scaling. Could this reapplication of attention scaling be responsible for the drop in performance when using softmax in Table A1, perhaps due to the reapplication of softmax being more prone to vanishing gradients?\n- Do the language modeling results improve if a potential issue with using softmax with HYLA is identified and fixed?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a novel perspective of multi-head attention as a hypernet. More specifically, the perspective that queries and keys dynamically create a function to apply to the values is a straightforward interpretation of the attention mechanism, but the proposed perspective further reveals where the learned compositional structure is potentially exhibited.\n- The core technical contributions of the \"multi-head self-attention as a hypernetwork\" perspective and HYLA are presented very clearly.\n- The proposed method HYLA is a novel drop-in replacement for standard multi-head self-attention.\n- The authors clearly show emergent compositionality in the \"latent codes\" for the intrinsic hypernets of multi-head self-attention and HYLA.", "weaknesses": "The paper seems to consider the special case of self-attention:\n- The proposed perspective of attention as hypernet is presented with the assumption of self-attention, which is a special case of attention (e.g., cross-attention is more general). While the assumption of self-attention can be seen in figures and equations, it is only explicitly mentioned in the text at line 114. Without clarifying that this perspective and the proposed modification (HYLA) consider the special case of self-attention, the paper can be seen as over-claiming its contribution. The paper's title, abstract, and introduction should clarify this constraint. Alternatively, the paper should clarify how its formulations work in more general cases of attention and if their experiments are limited to the particular case of self-attention.\n\nClarity and connection of multi-head to compositionality:\n- The necessity and significance of the \"multi-head\" aspect for compositional generalization is underplayed. A relevant discussion is found in the first paragraph of Section 6, which could be discussed earlier in the paper to set up a clear connection to compositionality.\n- The introduction of the paper is quite dense. After a clear flow of the first two paragraphs, it abruptly gets into too many technical specifics and jargon in the third paragraph that is difficult to parse without the context of the rest of the paper. Perhaps a less technical introduction of the paper's perspective and modification, with a motivation of how these relate to compositional generalization, can set up clearer expectations for what is to come and why it matters (maybe the Summary of this review can be helpful if accurate). Only after this, the details like how this paper's method differs from Schlag et al., 2021, or how the presented idea is evaluated become easier to place into context.\n\nCorrelation between compositionality and generalization:\n- The paper claims that the compositional structure seen in the intrinsic hypernet of MHA can explain the compositional generalization of transformers. However, there is a missing connection here. An important experiment is to evaluate with confidence estimates how the compositionality of attention correlates with the OOD generalization of the model.\n\nWeak performance on language modeling:\n- Despite the positive results on synthetic tasks, the paper's proposed method, HYLA, fails to outperform softmax attention for language modeling with a clear performance gap (Figure 6). The authors hypothesize that softmax is essential here, but find that HYLA with softmax results in a drop in performance (Table A1). Please look at the last question in the Questions section for a possible actionable avenue.\n\nRMSNorm:\n- It looks like RMSNorm at lines 179-180 is the operation from [1], but the paper is not cited.\n- The authors state that it ensures the parameters of the value network are appropriately initialized to prevent exploding/vanishing gradients. There needs to be an explanation of why this should be the case.\n\nReferences:  \n[1] Zhang, Biao and Sennrich, Rico. Root Mean Square Layer Normalization. NeurIPS 2019.", "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": "EKqjQ0HJbz", "reviewer": "Reviewer_47bR", "summary": "The paper aims at understanding the ability for compositional generalization for the transformer networks. The idea is to view the multihead attention module as a hypernetwork, where a composable, low-dimensional latent code specifies key-query specific operations. Empirically, the paper finds that using non-linear value networks for the hypernetwork (i.e., the multihead attention module) leads to improved generalization ability. The paper further introduces a symbolic RAVEN test for fine-grained evaluations of compositional generalization ability.", "questions": "Please see the weaknesses section.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea of viewing multi-head attention as a hypernetwork is interesting.\n- The experiments on a controlled domain such as symbolic RAVEN test are clear and thorough.\n- The research problem is of great interest to the ML community.\n- The paper is well-written and well-organized.", "weaknesses": "- I may be missing something. But the modification of the linear attention layer seems a little bit restricted. Will other non-linearity functions or $\\sigma(\\cdot)$ functions other than `RMSHead` also lead to improved compositional generalization ability? It would be great if the authors can comment on this and explain why this modification would work.\n- I am admittedly not an expert of this domain, and would like to hear my colleague reviewers' thoughts on this.", "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": "mI6OyAJ7I2", "reviewer": "Reviewer_Rtsz", "summary": "The paper investigates how Transformers generalize to new, unseen task compositions by reformulating multi-head attention as a hypernetwork—a model that reconfigures another model's computation based on data. This study reveals that a compact, low-dimensional latent code within the attention mechanism specifies key-query operations, which can generalize compositional tasks even if those specific combinations weren’t encountered during training.", "questions": "1. Could you provide a more formal explanation or theoretical basis for why framing multi-head attention as a hypernetwork enhances compositional generalization? Specifically, what is the role of the latent code in enabling the network to handle unseen task compositions?\n\n2. Given the computational demands of Transformers in large-scale applications, would HYLA introduce significant memory or efficiency constraints, and if so, how might these be addressed?\n\n3. Have you considered conducting deeper analyses, such as probing or interpretability methods, better to understand the latent codes’ alignment with specific subtasks?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper provides an innovative perspective by reformulating multi-head attention as a hypernetwork, a novel conceptual framework that opens up new avenues for understanding attention mechanisms in neural networks.\n\n2. The paper is well-structured, with clear explanations of its objectives, methods, and findings.", "weaknesses": "1. The paper makes a strong case for viewing multi-head attention as a hypernetwork, but it lacks a formal theoretical analysis showing why this approach specifically enhances compositional generalization over traditional attention mechanisms. Although the empirical results support the hypernetwork framing, a deeper theoretical exploration would add rigor and clarity.\n\n2. The paper primarily contrasts the proposed approach with standard softmax and linear attention but lacks comparisons with other recent methods that enhance compositional generalization in Transformers.\n\n3. While SRAVEN offers a modular benchmark for evaluating compositional generalization, it may not fully reflect the demands of compositional reasoning in language or perception tasks, which involve a broader range of dependencies and symbolic variations. \n\n4. Although the proposed Hypernetwork Linear Attention (HYLA) and hypernetwork framing show promising results, there is limited discussion on scalability and practical implementation challenges, especially as the model size increases. For instance, it remains unclear whether the approach would scale effectively in large language models (LLMs) or in high-dimensional data contexts.\n\n5. The authors posit that the structured latent code in multi-head attention serves as a reusable resource for different compositional tasks. While the paper provides clustering analyses, there is limited interpretability regarding how the latent codes correlate with specific task elements or subtasks.", "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": "jCzo5uJITQ", "reviewer": "Reviewer_yrk8", "summary": "This paper presents an interpretation of multi-head attention as a \"hypernetwork\", in the sense of Ha et al. (2017)---a network that generates weights that parameterize another network. The main insight comes from writing MHA as: $\\sum_{h=1}^{H} W_h^{out}\\sum_{k=1}^{K} a_{h,q,k} W_h^{val} x_k = \\sum_{k=1}^{T} \\left(\\sum_{h=1}^{H} a_{h,q,k} W_h^{out} W_h^{val}\\right) x_k$, and interpreting the inner summation in the RHS as the parameterization of the \"value network\". This comes from swapping the summation over heads $h$ with the summation over keys in the typical formulation of attention. The authors argue that this can give rise to some notion of \"compositionality\" and OOD generalization. While the value network here is parameterized as a linear combination of linear maps, the authors use this interpretation to propose a variant of linear attention where the value network is instead a two-layer MLP, with each linear map in the MLP being parameterized by the hypernetwork. The paper then explores two experimental settings: the first is based on ICL on \"fuzzy logic functions\" and the second is based on a symbolic version of Raven's progressive matrices. The authors show that networks with linear attention and softmax attention achieve some level of OOD generalization on these tasks, and that their proposed HYLA model achieves further improved performance.", "questions": "- How does this interpretation of MHA as a hypernetwork relate to the \"residual stream\" view, and related concepts on attentional circuits in previous mechanistic interpretability work? Does it have any implications on simple circuits like, say, induction heads?\n- How do MoE models compare in these compositional generalization experiments?\n- Typically, the value network in MHA is thought of as merely projecting out a certain component of the context token to be retrieved and incorporated into the query token. The value network is not intended to perform any complex processing, that is the job of the MLP that immediately follows. However, the interpretation of MHA as a hypernetwork seems to hinge on the hypothesis that the value network ought to be doing some more complex processing. Can you comment on this distinction? Why would we want the value network to perform complex processing rather than having that be performed by the MLP?\n- Under the hypernetwork view, the \"value network\" has a pretty specific parameterization as a query/key-dependent linear combination of linear maps. Can you comment on why this type of parameterization would make sense? For example, is there some simple construction of a network where there is an interpretable view of the class of functions this could capture?\n- In lines 370-374, the authors discuss the difference between standard RPM problems with visual inputs and their simplified version with symbolic input, arguing that the main source of difficulty in RPM tasks is the problem of finding correspondences of rules to feature dimensions. I found this unconvincing, however, and I hope the authors can expand. In particular, it seems to me that this argument relies on the assumption that the visual representation learning step of discovering the right symbolic abstractions is easy, citing work in cognitive science. But what might be easy for humans (who have in-built abilities through millions of years of evolution) may not be easy for machines. Learning the full visual RPM problem end-to-end is surely significantly more difficult for a machine learning system, no? Overall, I think it's unnecessary to try to make this argument. Yes, the symbolic RPM problem is easier, but it is still interesting; we don't need to argue that it captures all the interesting aspects of the full visual RPM problem.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The question of how compositional generalization arises in neural networks, and Transformers specifically, is of broad interest to the ML research community. I believe that the interpretation proposed in this paper is a meaningful contribution to this question.\n- The paper is well-written, well-structured, and clear.\n- To my knowledge, the interpretation of combinations of heads as a hypernetwork is novel.\n- The proposal of the HYLA attention mechanism is well-motivated and makes sense given the presented hypernetwork interpretation of MHA.\n- The experimental results provide some empirical support for the general claims made in the paper, with HYLA performing favorably in terms of compositional generalization on the tasks considered.\n- I found the figures showing the correlation between the latent codes and the target label (i.e., 2B and 5A) compelling, making the argument that some notion of compositionality is made possible by having multiple heads and", "weaknesses": "- The interpretation of MHA as a hypernetwork through re-ordering the summation over heads and keys makes sense in the case of linear attention, but seems less so in the case of softmax attention. In the case of softmax attention, the value network is a linear combination over heads, but with the normalization being over keys rather than heads. This doesn't fit into the claim that the value network is merely query-key dependent; rather it is a function of the entire sequence. Can the authors comment on this? You focus more on linear attention in the paper, and perhaps this is the reason. But, can you share more of your thinking about this?\n- The idea of a value network being input-dependent seems closely related to Mixtures-of-Experts. However this is not directly addressed in the paper. Can you comment on this? In particular, is there a type of function that HYLA can represent that standard MHA followed by a MoE-MLP could not? Experimentally, how would the latter perform compared to HYLA? Would it be better or worse at compositional generalization?\n- Although the experiments are clear, the scenarios where compositionality and the advantage of HYLA are demonstrated are somewhat restricted in scope. On complex real-world tasks such as language modeling, HYLA performs worse compared to Softmax attention (this is again related to HYLA being a variant of linear attention).\n\nSee also the questions below.\n\nOverall, I think this is a nice research paper, with meaningful contributions, and I look forward to discussion with the authors.", "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 (Oral)