Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers

Thanks for your helpful review!

Q1: The experiments in this paper are insufficient to demonstrate authors’ claim, they only evaluate with limited number of datasets and model structures, which might not be representative of the general behavior.

A1: Thanks for expressing this concern. We set out to understand how and why feed-forward and attention layers learn different parts of the data distribution in simple and controlled settings, which allows us to have theory and experiments that closely match, and thus provide robust evidence that may generalize to larger systems. More complex data models and architectures are likely to be much more complicated to analyze, which would make our arguments less transparent. We also note that despite their simplicity, our toy models in Section 3 already capture the complexities of the task at hand. Finally, our findings extend to large language models pretrained on real-world datasets, as shown in Section 4, which highlights the robustness of our insights.

Q2: The technique authors recommended, LASER, is not introduced by this paper, making this paper seems an extension or complement to LASER. Differentiation between this paper’s contributions and [1]’s should be claimed. Is there potential modifications or improvements for the authors to propose to LASER that could strengthen this paper’s contributions?

A2: We would like to emphasize that we use LASER truncation merely as a tool to probe whether certain associations are "stored" in different layers. The original paper's finding that LASER may improve reasoning performance was in fact a motivation for us to understand the phenomenon more deeply. This lead us to define the notions of distributional vs in-context associations, and to show how the training process may disentangle them into feed-forward and attention layers.

Q3: The methods and datasets this work employed is interesting but mostly appear a simple extension of that in [2]. This makes the research seem a supplementary work.

A3: We agree that [2] is a significant previous work for Section 3 of our work, but we would like to highlight the following key differences between these two:

1. Different motivations. We consider more realistic scenarios where there is competition between in-context reasoning and distributional associations (e.g. predicting "in the" and "in Spain" are both possible after "Madrid is located"), while [2] has a clear separation between in-context predictions (only on trigger tokens) and global bigrams (only on non-trigger tokens).
2. Different theoretical results. Our theoretical results differ significantly from [2]: our focus is on comparing feed-forward and attention layers in storing distributional associations, and we provide finite sample results, which is crucial for understanding how the model disentangles distributional associations from in-context reasoning. In contrast, [2, Theorem 3] only considers gradients with infinite samples, and only studies attention layers for learning the induction head mechanism.
3. Empirical verification of more tasks on pre-trained LLMs. We extend our study to more realistic reasoning tasks with pre-trained LLMs (see Section 4), which shows that our findings about disentanglement of distributional associations and in-context reasoning extends to such real-world scenarios. This includes by the impact of truncating feed-forward layers on complex reasoning paradigms such as GSM8k with CoT. In contrast, [2] only considers two-layer transformers on the synthetic bigram dataset.

References:

[1] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction. Sharma et al. ICLR 2024.

[2] Birth of a Transformer: A Memory Viewpoint. Bietti et al. NeurIPS 2023.

Thanks for your encouraging feedback!

Q1: The theorems and their consequences are relatively complicated to understand for a non-initiated reader of this literature, perhaps try to explain the different terms a little more.

A1: We agree with that and will keep thinking about how to improve the presentation. This is also why we provide Figure 1 and 2 to help understanding our motivations and results even without reading the theorems.

Q2: Experimental figures have no error bars.

A2: Thanks for the reminder! For the settings of IOI and factual recall in Figure 5, we add new plots of the prediction distributions in Figure 15 and 16 in the revision. Both are evaluated on the last checkpoints of the models.

1. Figure 15 uses the same 100 IOI sentences on Pythia-410M and 1B.
2. Figure 16 uses 20 more examples of factual recall (in Table 3) to measure the prediction distribution, instead of the single factual query ''Madrid is located in'' in Figure 5.

For other figures in Section 3, the curves are quite stable due to a synthetic setting, so we omit the error bars for better visualization.

Q3: Why did you use such a small model for the experiments? Would it be possible to perform the same kind of analysis with more layers?

A3: In Section 3, we use a two-layer transformer because it is the simplest model capable of solving the induction head task (the recent work [1] proves that one layer is not enough). This simplicity provides results and architectures that are more readable and interpretable. For deeper models with more than two layers, since the attention scores are almost uniform at initialization regardless of the layer, all layers should behave similarly at the first gradient step, and we expect a result similar to Theorem 1 to hold for each layer. However, subsequent steps would likely become more complex and tedious to analyze.

In Section 4, small models are used in experiments for the following reasons:

• For IOI, as a quick demo following the setting in [2], we study ignored top predictions about generic words from GPT-2 small, while [2] focuses on only subjects ([S]) and indirect objects ([IO]). Then we extend the study to larger pythia models.
• As we are interested in the training dynamics from scratch, training checkpoints of pythia models are one of the very few open-source options that are accessible to us.

[1] One-layer transformers fail to solve the induction heads task. Sanford et al. arXiv:2408.14332

[2] Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. Wang et al. ICLR 2023.

Q4: What is your interpretation of the poorer performance of truncated MLPs after 8-shot on GSM8K (Table 2)?

A4: Thanks for this question. One possibility is that truncating MLPs may remove too many distributional associations, including ones that are useful in the reasoning chain, while few-shot prompting with many few-shot examples may succeed in only removing spurious distributional associations. Indeed, many-shot prompting has been found to promote in-context reasoning (to the point where it can break safety guardrails), suggesting that spurious distributional associations will be automatically dropped.

Q5: Your results seem to indicate that reducing FFNs greatly improves in-context reasoning, which is often desirable. However, FFNs also play an important role. It would therefore be interesting to study what is lost when they are reduced/removed?

A5: Thanks for this question, this is indeed an important point. One example where this might arise as a limitation is in GSM8K, as described in our answer above, where intermediate reasoning steps may be negatively affected by truncation, which leads to poor performance compared to using many few-shot examples. Understanding how to truncate layers more selectively, e.g., by erasing only certain well-chosen associations rather than entire subspaces, would be an important next step to overcome these limitations.

We'd be happy to clarify any further questions.

Thanks for your helpful review!

Q1: This work is highly theoretical and lacks certain practical significance and applicability.

A1: We agree with the difficulty of making practical benefits in LLMs with theoretical understanding. Nevertheless, we position this work as a careful study of ignored empirical phenomenon to inspire better post-training. For instance, in [1], they studied a curcuit of attention heads in a trained LLM on IOI task, totally focusing on probabilities of predicting subjects (S) and indirect objects (IO) at the end of training. However, we carefully study the probabilities of generic ''the'' word along the training in Figure 5 (left), as well as LASER. Following many previous works on alignment, this work is another evidence of improving a pre-trained model during post-training.

[1] Interpretability in the wild: a circuit for indirect object identification in gpt-2 small. Wang et al. ICLR 2023.

Q2: In addition to feed-forward layers and attention layers, there are also wording embedding layers, residual connections, and LayNorm layers in the Transformer architecture. What special role do these layers play in helping the model understand the semantics of the natural language? Did the authors explore?

A2: Thanks for raising this point! Our toy models do include residual connections and word embeddings, but they do not include LayerNorm, and the word embeddings are not trained. This is mostly to simplify our analysis, and we note that this simple model already suffices to capture the different behavior of FF and attention layers. While learned embeddings are crucial in real-world natural language problems, our toy task is mostly algorithmic and does not involve rich semantics. We expect that learning embeddings would be important in tasks with richer semantic structure (as in [2]). Layernorm is also crucial in practice for efficient and stable training, but we found it to have minor effects in our simple setup where data is well-balanced.

[2] How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding. Li et al. ICML 2023.

Q3: Because the attention module is the interaction between all tokens in the input sequence, it intuitively plays a role in understanding context. It seems obvious that the attention layer mainly focuses on in-context reasoning, right?

A3: While this may seem obvious, the fact that distributional and in-context associations are disentangled into different blocks of a transformer is a nontrivial result! Indeed, in terms of expressivity alone, attention could easily capture distributional associations, and training an MLP on top of an untrained/frozen attention could capture in-context associations. The fact that different layers automatically specialize to disentangle between these two tasks is nontrivial, and follows from a careful study of training dynamics.

Q4: Are the given conclusions applicable to all models based on the Transformer architecture? Is this applicable to all LLM models?

A4: We do believe that our conclusions would apply to any Transformer LLMs. We have verified this empirically on GPT-2, Pythia models up to 2.8B, and Llama models up to 8B, and there is no reason to believe these behaviors would change at larger scale.

We did not investigate other architectures beyond Transformers, but that would be an interesting avenue for future work.

Q5: In terms of format: Line 1935, Line 1938, Line 1952, Line 2001, Line 2069 exceed the margins and have incorrect formatting.

A5: Thanks for pointing out! We have fixed them in the revision.

We'd be happy to clarify any further questions.

Appreciate the reply. After reading the comments from other reviewers and the authors' responses, the authors answered most of my questions, and fixed the format in the revision version. I will keep my score.

Thanks for your helpful review!

Q1: The novelty presented the results section (Section 4) is unclear [...]

A1: The novelty of our work is not any specific architectural change or the application of LASER. Rather, our contribution is an understanding of how FF and attention layers are able to disentangle different parts of the data distribution during the training process, namely by preferentially storing distributional associations in FF layers, while attention layers tend to pick up in-context reasoning mechanisms. Our use of weight truncation / LASER is simply as a tool for probing for this behavior in both our toy models and on pretrained LLMs, and Section 4 is meant to investigate the validity of our claims on actual LLMs beyond the synthetic settings of Section 3.

Q2: Is the key takeaway that one should perform low-rank truncation as in LASER on the feed-forward layers since they store noise? Including discussion on this would improve the actionability of the findings of this work and its impact.

A2: As mentioned above, our contribution focuses on understanding the different roles of attention and FF weights in disentangling distributional vs in-context associations, both empirically and theoretically. The application of low-rank truncation is simply a way to verify our claims, and is consistent with the findings in the LASER paper that truncating some FF layers may improve performance on some reasoning tasks.

Nevertheless, our perspective based on distributional associations versus in-context reasoning may be helpful in thinking about how to allocate parameters to feed-forward versus attention layers: for instance, on our synthetic task, we found that for a fixed total parameter budget, models with fewer MLP parameters achieve higher loss on distributional predictions (e.g., non-contextual bigrams) compared to models with more MLP parameters (and fewer attention parameters). These notions may also provide a different way to reason about circuit discovery in mechanistic interpretability from the perspective of training dynamics and properties of the training data. Finally, this disentanglement may inform more effective ways to fine-tune models, e.g., by selectively choosing which layers to fine-tune.

We've included these points and the results on varying MLP vs attention parameters in the revision, see Appendix A.

Q3: I would suggest moving the proofs of Lemmas C.1 and C.2 up, before they are invoked in the proof of Theorem 1 / 4. This would make it easier to follow the progression of claims made.

A3: We have reorganized the appendix in the revision.

Q4: In the abstract, the term “gradient noise” is used in line 21, but not referred to in the rest of the paper — please clarify (or remove, if not relevant) this term.

A4: Thanks for pointing this out. "Gradient noise" is referring to the fact that the attention gradients at initialization are much more noisy than the feed-forward gradients in storing the distributional association (see Theorem 1 and section 3.1 more generally). We have changed this to "noise in the gradients" in the abstract to avoid confusion.

Q5: The use of the term 'sequel' on line 125 is unclear; this should be clarified or rephrased accordingly.

A5: Thanks for pointing this out, we replaced "in the sequel" by "below": the in-context recall task is studied in Section 3 while the other two are explored empirically in Section 4.

We'd be happy to clarify any further questions.