Attention as a Hypernetwork

Authors: Simon Schug, Seijin Kobayashi, Yassir Akram, Joao Sacramento, Razvan Pascanu
Primary area: interpretability and explainable AI
Keywords: attention, compositional generalization, abstract reasoning, in-context learning, transformer, mechanistic interpretability

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.

Decision: accept (oral)
Reviewer scores: 8, 8, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730503408658

thread_type:
official_review

<thread_object>
creator: reviewer_yrk8

<content>
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.

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.
- The paper is well-written, well-structured, and clear.
- To my knowledge, the interpretation of combinations of heads as a hypernetwork is novel.
- The proposal of the HYLA attention mechanism is well-motivated and makes sense given the presented hypernetwork interpretation of MHA.
- 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.
- 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?
- 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?
- 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).

See also the questions below.

Overall, I think this is a nice research paper, with meaningful contributions, and I look forward to discussion with the authors.

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?
- How do MoE models compare in these compositional generalization experiments?
- 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?
- 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?
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671255459

thread_type:
official_review

<thread_object>
creator: reviewer_rtsz

<content>
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.

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.

2. 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.

2. 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.

3. 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. 

4. 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.

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

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?

2. 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?

3. Have you considered conducting deeper analyses, such as probing or interpretability methods, better to understand the latent codes’ alignment with specific subtasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673151712

thread_type:
official_review

<thread_object>
creator: reviewer_47br

<content>
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.

strengths:
- The idea of viewing multi-head attention as a hypernetwork is interesting.
- The experiments on a controlled domain such as symbolic RAVEN test are clear and thorough.
- The research problem is of great interest to the ML community.
- 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.
- I am admittedly not an expert of this domain, and would like to hear my colleague reviewers' thoughts on this.

questions:
Please see the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730674900141

thread_type:
official_review

<thread_object>
creator: reviewer_59h3

<content>
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. 

Under this perspective:
- 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
- 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.

With this perspective:
- 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.
- 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.

The 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.

Finally, the authors evaluate HYLA on a language modeling task, but it fails to outperform softmax attention here.

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.
- The core technical contributions of the "multi-head self-attention as a hypernetwork" perspective and HYLA are presented very clearly.
- The proposed method HYLA is a novel drop-in replacement for standard multi-head self-attention.
- 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:
- 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.

Clarity and connection of multi-head to compositionality:
- 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.
- 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.

Correlation between compositionality and generalization:
- 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.

Weak performance on language modeling:
- 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.

RMSNorm:
- It looks like RMSNorm at lines 179-180 is the operation from [1], but the paper is not cited.
- 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.

References:  
[1] Zhang, Biao and Sennrich, Rico. Root Mean Square Layer Normalization. NeurIPS 2019.

questions:
HYLA with a single head:
- In the caption for Figure 5C, the authors state that the hypernet mechanism is absent for a single head. Why is this the case?

Scaling compositional difficulty:
- 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.

Drop in performance of HYLA with softmax:
- 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?
- Do the language modeling results improve if a potential issue with using softmax with HYLA is identified and fixed?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730974737587

thread_type:
official_review

<thread_object>
creator: reviewer_zlcd

<content>
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.

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).

Similarly, 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)?

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)?

- Please report the training performance of the models to shed light on the generalisation vs capacity question.

- Please provide more details on the OOD splitting of the datasets.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732153867804

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all reviewers for their positive reviews and insightful comments. We have revised the manuscript based on your feedback to improve the clarity of the exposition (see individual responses) and have added new results and figures based on several experiments suggested by the reviewers. 

We believe these additional results provide further evidence for the central hypothesis of the paper that the hypernetwork mechanism present in multi-head attention supports compositional generalization.

Notably, we added the following to the revised version of the manuscript:

* As suggested by Reviewer ZLCD, we added appendix figures A3 and A4 on page 18-19 that show the training loss on all fuzzy logic and SRAVEN experiments conducted in the main text. They provide further evidence that encouraging the hypernetwork mechanism identified in multi-head attention does not simply result in additional capacity to fit the training data but specifically leads to improvements in compositional generalization.
* As suggested by Reviewer 59h3, we conduct a new experiment that compares the compositional generalization performance as the fraction of OOD tasks varies for SRAVEN in Figure A5 on page 19 similar to the existing experiment on the fuzzy logic task in Figure 2C. The experiment confirms that encouraging the hypernetwork mechanism improves compositional generalization even for large fractions of held-out tasks before performance eventually degrades across models.
* As suggested by Reviewer yrk8, we have conducted new experiments to evaluate to what extent replacing the MLP layer of the transformer block with Mixture of Experts (MoE) improves compositional generalization in Appendix A1 and Table A1. We find only moderate improvements in line with our findings on how increasing model capacity in general benefits compositional generalization in Figure 4.
* As suggested by Reviewer 47bR, we perform additional model ablations varying the nonlinearity of the value network introduced in HYLA (see individual response below) showing that encouraging the hypernetwork mechanism with other nonlinearities leads to similar improvements in compositional generalization.

We hope to have resolved your concerns (please see the individual responses) and are looking forward to engaging in further discussions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper demonstrates a mathematical equivalence between multihead attention and linear hypernetworks, which is then used to provide an explanation for why Transformers are able to compositionally generalize to some extent. The attention scores per head can be interpreted as a low-dimensional latent code that predicts the resulting linear value network that processes the input. Then, this insight is used to construct a more expressive approach to compositional processing by modifying MHA to make the value network nonlinear, and also introduce a normalization scheme that leads to linear value networks that are more consistent with desirable network initialization properties. Through experiments on two toy reasoning datasets (a fuzzy boolean logic task and a symbolic Raven matrices task), the paper shows evidence of compositional generalization that is enhanced for the nonlinear version, lending support for the hypernetwork view of attention. Qualitative visualizations show that the heads are consistent with hypothesized compositional functions. Finally, experiments on language modeling show that the nonlinear hypernetwork approach improver over standard linear attention baselines, adding further support for the hypernetwork view of attention.

Strengths:
Well-written paper posing an interesting hypothesis with provocative results on compositional generalization. May open new design space choices for improving attention schemes in general and towards better generalization in reasoning tasks. The nonlinear modification and normalization scheme of HYLA may be valuable in its own right as a drop-in replacement for attention. The paper may also spur new ways of approaching interpretability in attention models. Experiments are nice in that they allow controlling degree of compositionality.

Weaknesses:
Not really a core weakness, but it would be nice to see controlled experiments with cross-attention, since it was claimed that the same treatment should hold. There are also some unexplained aspects of standard softmax self-attention that are not captured by the current treatment of linear MHA as a hypernetwork (and the HYLA variant), and one reviewer is not fully convinced that the hypernetwork equivalence is a useful construct to explain compositional generalization in softmax attention. 

Decision reasoning:
Reviewers unanimously felt the paper is a strong contribution, with 4/5 reviewers giving scores of 8. I agree with the reviewers’ assessments and find the paper quite interesting, particularly towards better understand why Transformers are able to reason compositionally in at least some cases. The authors also put a lot of effort into addressing reviewer concerns, which made the final paper significantly stronger. Papers like this are useful for understanding methodology, but may also lead to new developments — while the authors caution that the HYLA variant of MHA is more of an ‘experimental device’ rather than a legitimate enhancement, it may inspire new advances that enhance the expressiveness of attention.