Attention as a Hypernetwork

We are grateful for your positive feedback on the quality of our paper's writing and its contribution towards understanding compositional generalization in neural networks. Based on your comments and suggestions, we have conducted additional experiments on MoEs and revised parts of the manuscript. In the following we answer your points one-by-one hoping to clarify any remaining concerns.

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. [...]

The interaction with the softmax is indeed an interesting point to consider. First of all we would like to clarify that we are not claiming that the value network is merely query-key dependent but rather query-key specific in the sense that from the hypernetwork perspective each query-key pair has a different latent code. We apologize if this was not clear enough in the manuscript - we have tried to improve the clarity of this point but please let us know if there are particular passages in the paper that need further improvement. As you correctly point out, the softmax introduces an interaction across keys that makes the key-query specific latents dependent on the whole sequence. Specifically it creates additional competition across individual entries of the latent code before this code is used to parameterize the value network. Our latent code analysis reveals that even with this competition among components a structured latent code can emerge that allows decoding operations performed by the network (see Figure 5B+F specifically).

The idea of a value network being input-dependent seems closely related to Mixtures-of-Experts. [...]

This is a very good question. One broader point to clarify in this context is that in principle multiple transformer blocks with standard MHA and standard MLPs can already emulate a single transformer block with HYLA. Intuitively, such a construction can be obtained by leveraging the nonlinearity of the MLP layer of the first block to function as the additional nonlinearity of HYLA. We primarily introduce HYLA as an experimental device that encourages the use of the hypernetwork circuit motif in order to evaluate whether it has the hypothesized effect on compositional generalization. Experimentally we observe in Figure 4 that at sufficient scale standard attention mechanisms can achieve comparable compositional generalization performance to HYLA. That being said, the idea that MoEs might similarly allow to compartmentalize and compose knowledge is interesting and we have added this analysis (see next response).

How do MoE models compare in these compositional generalization experiments?

We have conducted these additional experiments based on your suggestions and added them to Appendix A1. We replaced the MLP layers with MoEs with E=5 experts of the size of the original MLP layer and a router that sparsely activates the top k=2 experts, i.e. each MoE has ~5x more parameters than the original MLP block and ~2x the number of active parameters (discounting the parameters added by the router),

$f_{\mathrm{MoE}}({x})= \sum_{i=1}^{E}\mathrm{Softmax}(\mathrm{Top\_k}({W}^{\text{gate}}{x}))\cdot{W}^{\mathrm{up}}_i\phi({W}^{\mathrm{down}}_i{x}).$

This generally leads to a slight improvement in terms of OOD performance across attention variants which is in line with our observations increasing the capacity of the models by introducing more depth/width shown in Figure 4 but does not point to these types of MoEs specifically improving compositional generalization.

Thank you. I appreciate the continued discussion and your detailed responses to my questions and concerns.

I have raised my score in appreciation of the authors' efforts in improving the clarity of the presentation and the additional experimental results, which have addressed some of my concerns.

However, on a final note, I'm not entirely convinced of the argument that it is useful to view softmax attention as a hypernetwork. Ultimately, all this seems to build down to is a simple question---Is it more useful to understand MHA as $\sum_{h} \sum_{k} \alpha_{qk}^h \\, W_{ov}^h \\, x_k \quad \text{(Standard Interpretation)}$ or as $\sum_{k} \sum_{h} \alpha_{qk}^h \\, W_{ov}^h \\, x_k \quad \text{(Hypernetwork Interpretation)}$

This seems to be what you refer to as the "mathematical equivalence". But this isn't a deep mathematical derivation, it's just swapping the order of the summation. This might give rise to a useful interpretation (which I believe it does in the case of linear attention, as you demonstrate in the paper), but it is just an interpretation of the same underlying operation.

In the hypernetwork interpretation, you interpret $\sum_{h} \alpha_{qk}^h W_{ov}^h$ as a hypernetwork, because it is a linear combination of linear maps, with the linear weights determined by the input. This is a neat observation in the case of linear attention, where the two summations are intuitively exchangeable. But when you have (softmax) normalization over $k$, the hypernetwork interpretation obfuscates this crucial aspect and gives the wrong idea. You can stretch the definition of a "hypernetwork" to allow for this kind of scenario, but it is not clear that this continues to be a useful conceptual framework in that case.

we agree that a priori it is unclear if there is additional explanatory power in thinking of multiple heads as implementing a circuit together vs. each head simply forming their own circuit, that can be understood independently of other heads (but potentially still combined).

I don't think that the hypernetwork interpretation is necessary to understand multiple heads as implementing a circuit together. There is clearly an interaction between heads since they add up in their contribution to the residual stream, and the attention patterns of the different heads are a function of the same underlying input. I don't think that one needs to adopt the hypernetwork interpretation to understand how heads "collaborate". If you'd like to argue for the hypernetwork interpretation, I think you would need to demonstrate that this very particular way of parameterizing the linear weights that define the hypernetwork (including the softmax normalization over keys) gives rise to some interpretable operation. In the paper, you show that the latent codes have structure, but you don't interpret the exact "hypernetwork" that was learned.

I do believe that the exploration of the "latent code" $[\alpha_{qk}^1, ..., \alpha_{qk}^H]$, and interpreting its correlation to the target is interesting. I had not seen it before. But again, even under the standard "$\sum_{h} \sum_{k} \alpha_{qk}^h \\, W_{ov}^h \\, x_k$" interpretation, you would still expect some structure in the latent code (since different heads will activate in different places in an input-dependent way). You don't necessarily need the hypernetwork interpretation (which obfuscates the normalization over keys, which is often crucial).

My final suggestion to the authors, in case it's useful, is to consider revising the framing of the paper to more explicitly explain the purpose of proposing HYLA (i.e., as an "experimental device", and their argument about where the hypernetwork interpretation is or is not useful. I hope some of my feedback was helpful.

We are happy to hear that you believe our paper offers an innovative perspective that opens up new avenues for understanding attention mechanisms in neural networks. It is in this light that we introduce HYLA, namely as an experimental device to improve our understanding of attention. We believe that some of your concerns might be a result of considering HYLA as a suggested replacement for the attention mechanism in transformers to improve their performance but this is not what we propose. We apologize for this potential misunderstanding. In the following we answer your points more specifically, hoping that you reconsider your overall rating of our submission in light of this clarification.

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.

In section 2.1 we provide a mathematical derivation that links multi-head attention to hypernetworks. We would like to clarify that this derivation applies to traditional attention mechanisms (e.g. multi-head softmax and linear attention) and is not by itself a separate approach to improve compositional generalization. Prior work has theoretically and empirically shown that hypernetworks support compositional generalization [1]. We apologize that this point has not been made more explicit in the original submission and have revised section 2.1 to clarify this. With the theoretical link between multi-head attention, hypernetworks and compositional generalization in place, we design our experiments to verify the predictions made by this theoretical framework, for instance showing that the latent code of the hypernetwork is predictive of the operations implemented by multi-head attention.

[1] Schug et al. “Discovering modular solutions that generalize compositionally.” (2024)

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.

Given that our goal is to understand to what extent the hypernetwork mechanism explains network computations of learned transformers, the introduction of HYLA serves primarily as an experimental device to examine the hypothesis that the intrinsic hypernetwork of multi-head attention is what supports compositional generalization on abstract reasoning tasks that require recombining previously acquired operations. We attempt to clarify this part in the abstract (line 16-19). Furthermore as we show in Figure 4, at sufficient scale, linear and softmax attention can similarly compositionally generalize. Importantly, their internal workings can be understood through the hypernetwork mechanism as for instance the analyses in Figure 5 aim to show.

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.

We consider abstract reasoning problems such as SRAVEN and the fuzzy logic task since this gives us precise control over the problem compositions encountered during training and evaluation and allows us to verify specific hypotheses pertaining to the mathematical equivalence of multi-head attention to a hypernetwork - for example the ability to decode the operations performed by the network based on the latent code in Figure 5F. We furthermore conduct language modeling experiments with 50M parameter models trained on 150B tokens. We find that also in this setting the experimental intervention of reinforcing the hypernetwork mechanism of standard multi-head linear attention through HYLA leads to a noticeable improvement in performance performing closely to softmax attention. This is noteworthy since it is in a setting where the softmax is known to be important due to its role in binding and associative recall problems on which linear attention methods typically struggle [2,3,4].

[2] Arora et al., Zoology: Measuring and Improving Recall in Efficient Language Models (2023)

[3] Schlag et al., Linear Transformers Are Secretly Fast Weight Programmers (2021)

[4] Olsson et al., In-Context Learning and Induction Heads (2022)

Thank you for your encouraging review - we are glad that you found our experiments to be clear and the paper well written.

I may be missing something. But the modification of the linear attention layer seems a little bit restricted. Will other non-linearity functions or 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.

Thank you for this suggestion, we have run additional experiments to test the influence of different nonlinearities on the performance of HYLA on SRAVEN (same task and model setup as in Table A1).

These results show that the ReLU nonlinearity is generally a well performing choice but not the only one that can lead to improvements in compositional generalization. In addition to varying the nonlinearity of HYLA, we perform many model ablations in Appendix A1 to verify how the different modifications of HYLA affect compositional generalization. This includes testing a model that applies the RMSHead to linear attention and a variant that removes RMSHead from HYLA. Overall we find that all the modifications we introduce in HYLA over linear attention together benefit OOD performance on both SRAVEN and the fuzzy logic task.

Thank you for having taken the time for an elaborate and insightful review. We appreciate the many suggestions for improvements you have made and believe that they have helped to improve the paper. We found your remark that you would like to rate the paper higher encouraging and hope that we were able to address your remaining concerns.

The proposed perspective of attention as hypernet is presented with the assumption of self-attention [...]

Thank you for raising this concern, we are indeed focusing on self-attention in our experiments. The main derivation we present equally applies to cross-attention and we are clarifying this point in our updated version of the manuscript: In the case of cross-attention the keys and values come from the output of the encoder whereas the queries come from the output of the decoder. This means that based on the latent codes obtained from combining queries from the decoder and keys from the encoder, the hypernetwork mechanism generates a value network that processes inputs from the encoder stream and adds them to the decoder stream.

Clarity and connection of multi-head to compositionality [...]

We agree that the mentioned paragraph was overly technical in the original submission and the important role of multiple heads was not sufficiently well communicated. We have rewritten it based on your feedback in our updated version of the manuscript hoping to have improved its clarity.

[...] An important experiment is to evaluate with confidence estimates how the compositionality of attention correlates with the OOD generalization of the model.

It is difficult to operationalize the compositionality of attention directly and we would be curious if you had something else in mind: We can intervene on the compositionality of attention by removing all heads but one. With a single head (also see our response to your question in this regard), attention no longer supports the type of compositionality we hypothesize to help compositional generalization while maintaining the same number of parameters (we set $D^{\text{head}} = \frac{D}{H}$). The result of this intervention is shown in Figure 5C: In line with our hypothesis, for a single head the compositional generalization performance noticeably drops for all attention variants.

Despite the positive results on synthetic tasks, the paper's proposed method, HYLA, fails to outperform softmax attention for language modeling [...]

[...] 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?

Following your suggestion, we have checked whether the ablation that combines the nonlinear value network of HYLA with the softmax normalization over keys of standard attention leads to vanishing gradients but have not observed this pathology. This is in line with other works that introduce additional multiplicative interactions inside of neural networks [e.g. 1,2] have similarly not reported instabilities in learning.

One point we would like to stress is that we are introducing HYLA as an experimental device in an attempt to examine the hypothesis that the intrinsic hypernetwork of multi-head attention is what supports compositional generalization on abstract reasoning tasks that require recombining previously acquired operations. We take the fact that emphasizing this ability through HYLA improves language modeling performance compared to linear attention as evidence that the hypernetwork mechanism is also relevant in the context of language modeling. However we also acknowledge that softmax attention can be very flexible supporting other modes of operation such as for instance induction heads. HYLA was designed with compositionality in mind, but by emphasizing this aspect the layer might struggle with other behaviors that are expected from the attention layer in large complex tasks like language, as for example retrieval (which is not compositional).

[1] Jayakumar et al. Multiplicative Interactions and Where to Find Them. (2019)

[2] Chrysos et al. Deep Polynomial Neural Networks (2020)

Thank you for all your responses! My primary concerns have been addressed, and I have increased my score.

Thank you for your positive review highlighting that you found the paper to provide a novel perspective on an important problem. Your review helped us clarify an important point on the possibility of HYLA mereling increasing the capacity of the model rather than specifically improving compositional generalization in the updated version of the manuscript (and in our response below of course).

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

Thank you for raising this important point. We have added the training loss on the fuzzy logic tasks corresponding to Figure 2C in appendix figure A3D to complement the training R2 score we already reported in the appendix of the original submission. Similarly the new appendix figure A4 shows the training loss on the SRAVEN scaling experiments of Figure 4. We will point more prominently to these in the main text.

We generally find no consistent differences between methods in terms of in-distribution training loss. Together with the property that HYLA has exactly the same number of parameters as linear and softmax attention, we take this as an indication that it specifically improves compositional generalization and is not simply a way of improving overall model capacity.

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

Apologies for the confusion, we will try to clarify this point in the manuscript: In the fuzzy logic benchmark, each task is defined by a combination of K terms. So for instance if we have 3 possible terms A, B, C, and K=2 then the possible tasks are AB, AC, BC. During training we only train on a subset of tasks, e.g. we train on AB, BC and during OOD evaluation we evaluate on the held-out task, e.g. AC.

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

Yes, we perform a number of ablations in Appendix A1 including removing the RMSHead normalization from HYLA and adding it to linear attention. We find that while simply applying RMSHead to linear attention (Linear Attention + RMSHead) and also the attention-weighted sum operation without nonlinearity (HYLA − nonlinearity) both by themselves can slightly boost performance, the full HYLA model generally performs best in terms of compositional generalization. We will point more prominently to the ablations in the main text.