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

Authors: Lei Chen, Joan Bruna, Alberto Bietti
Primary area: foundation or frontier models, including LLMs
Keywords: reasoning, in-context learning, associative memory, transformers, distribution shift

Abstract:
Large language models have been successful at tasks involving basic forms of in-context reasoning, such as generating coherent language, as well as storing vast amounts of knowledge. At the core of the Transformer architecture behind such models are feed-forward and attention layers, which are often associated to knowledge and reasoning, respectively. In this paper, we study this distinction empirically and theoretically in a controlled synthetic setting where certain next-token predictions involve both distributional and in-context information. We find that feed-forward layers tend to learn simple distributional associations such as bigrams, while attention layers focus on in-context reasoning. Our theoretical analysis identifies the noise in the gradients as a key factor behind this discrepancy. Finally, we illustrate how similar disparities emerge in pre-trained models through ablations on the Pythia model family on simple reasoning tasks.

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

Reviews and discussion:
<thread>
creation date:
1730384060639

thread_type:
official_review

<thread_object>
creator: reviewer_p7ta

<content>
summary:
The authors studied how attention mechanisms and feed-forward layers handle distributional associations and in-context inputs differently. They considered simple two-layer transformers on an in-context recall task and observed the distinct focuses of feed-forward layers and attention mechanisms through analyzing gradient.
They further employed LLMs like GPT-2 and Pythia, concentrating on IOI dataset . Additionally, the authors validate the effectiveness of the LASER technique.
Overall, the authors underscore the differing roles of attention layers and feed-forward layers in language models.

strengths:
1. The authors conducts an in-depth study on how attention and feed-forward layers differently process distributional associations and contextual inputs, yielding clear and significant conclusions.
2. The flow of the paper is easy to follow. 
3. The authors employ clear theoretical proofs and a logical structure, supporting the authors' arguments.

weaknesses:
1. 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.

2. The technique authors recommended, LASER[1], 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.

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

**Reference**

[1] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction

[2] Birth of a Transformer: A Memory Viewpoint

questions:
1. Is there potential modifications or improvements for the authors to propose to LASER that could strengthen this paper’s contributions?
2. In Section 4.2, why LASER struggles with more shots? 
3. Is it possible to make other modifications beyond those in [2], such as designing an entirely new dataset to better validate your claims?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730392209211

thread_type:
official_review

<thread_object>
creator: reviewer_d4xh

<content>
summary:
In this paper, the authors investigate the role of transformers' self-attention and feed-forward networks (FFNs) both experimentally, in a controlled environment with a simplified model, and theoretically, using an even simpler model that allows them to derive calculations. The authors thus highlight an important distinction: FFNs are biased towards learning distributional associations (i.e. co-occurrence cues), while self-attention focuses on in-context reasoning. 

Concretely, as it is common to have the word “the” after the word “in”, a model using distributional cues only would be tempted to complete the sentence “Madrid is located in” with “the”, whereas a model using in-context reasoning would complete it with “Spain”.

So the authors set up an experiment in which the model is trained on a text containing with equal probability the “a + b” and “a + noise” patterns (where “a” varies from sentence to sentence, while the “noise” token remains the same) and then evaluated on a text containing only the “a + b” pattern. The idea is then to measure in-context recall in such a scenario: indeed, a purely distributional model will perform at 50% accuracy where a perfectly in-context reasoning model will achieve 100%. Using LASER to reduce the rank of the second-layer MLP matrix, until it is completely suppressed (rho = 0), they note that distributional behavior disappears with MLP rank, suggesting that FFNs are biased in this direction. They also note that self-attention manages to focus only on “a + b” patterns, ignoring “a + noise” patterns. 

The authors then propose a theoretical explanation for such behavior, analyzing in particular the first gradient descent step. Finally, the authors propose to test their findings in a more realistic setting, i.e. on real models and real-world reasoning benchmarks.

strengths:
- It is very well written and clear.
- Their experimental environment is very clearly established, from the creation of their dataset to the architecture of the model.
- I think that experimenting with LLMs in a controlled setting is an excellent methodology.
- The paper proposes a theoretical justification for their discovery that is more than welcome and very rare in the field.
- The findings are very interesting and help to better understand the behavior of LLMs, in particular FFNs, which are often less studied than self-attention.

weaknesses:
- The main experiments and theoretical results are based on highly simplified models. However, the observed behavior may not scale with larger models.
- 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.
- Experimental figures have no error bars.

questions:
1) 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?
2) What is your interpretation of the poorer performance of truncated MLPs after 8-shot on GSM8K (Table 2)?
3) 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730649436156

thread_type:
official_review

<thread_object>
creator: reviewer_gq5y

<content>
summary:
Large language models excel in in-context reasoning tasks like coherent language generation and knowledge storage. Their Transformer architecture relies on feed-forward and attention layers linked to knowledge and reasoning. In controlled experiments, the authors found feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning. Gradient noise explains this difference. Pre-trained models, like Pythia, show similar disparities in simple reasoning tasks.

strengths:
1. This paper uses a large number of experiments and theoretical proof to draw the following conclusions:
For the Transformer architecture of large language model LLM, feed-forward layers learn simple distributional associations, while attention layers focus on context reasoning.

2. The experimental results and conclusions are relatively solid, and the theoretical proof is quite convincing.

3. This work may be helpful for future researchers studying LLM.

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

questions:
1. 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?

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

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

4. In terms of format:
    Line 1935, Line 1938, Line 1952, Line 2001, Line 2069 exceed the margins and have incorrect formatting.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695869823

thread_type:
official_review

<thread_object>
creator: reviewer_u1qi

<content>
summary:
This work theoretically examines a phenomenon regarding the purpose of the feedforward and attention layers of a Transformer model, relative to the reasoning and knowledge learning tasks. It attempts to provide an explanation for the difference through a generic noise token, demonstrating that the feedforward layers store noise while the induction heads in the attention layers are insensitive to the noise tokens.

strengths:
* The paper sets out to address a clear objective, and designs a 2-layer model that enables a theoretical analysis of the phenomenon of interest, then reduces to a 1-layer model for theoretical analysis, revealing important, mathematically-grounded supporting evidence for said phenomenon. 
* I appreciate that the tasks of interest, and the distinguishment between the two abilities of interest (distributional associations and in-context reasoning) is made clear upfront. 
* Figure 2 is especially helpful to follow the 2-layer setup for the in-context recall task with noise tokens — the finding that the induction head filters out the noisy token is an interesting insight, and reinforced by Theorem 2 and its proof in Appendix D. 
* While the paper is dense, it is also organized quite well and in a manner that emphasizes the key takeaways.

weaknesses:
* The novelty presented the results section (Section 4) is unclear — are there any architectural insights informed by the theoretical analyses that could be applied that have been considered, beyond an application of LASER? Aside from this, this section appears to be a more in-depth, investigative analysis on the behavior of the chosen pre-trained models for the selected tasks; it is also unclear how / whether these insights scale as the models considered are small. 
* I have no major concerns about the theoretical results in the work, which reveal the means through which noise is stored both in the presence and absence of feedforward layers, under a simple 1-layer model. I must acknowledge that while I have largely studied and followed the proofs of Theorems 1 and 2 and the theory in Appendix A, I have not thoroughly verified the correctness of the theory in the rest of the appendix.

questions:
* 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. 
* 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. 
* The use of the term 'sequel' on line 125 is unclear; this should be clarified or rephrased accordingly. 
* As noted in the weaknesses section, are there any architectural changes that may be proposed, conditioned on the results of the exploration? 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732549451834

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updates in the Revision

comment:
Dear reviwers and AC,

Thanks for your helpful feedback! We posted individual reponses to each reviewer, along with a revised submission.

The revision has the following major updates:

1. **A new synthetic IOI task** in Appendix J, following the question about additional tasks by Reviewer P7Ta. This task requires counting the *correct* token $\bar{y}$ and the *distractor* token $y^D$ to predict the less frequent $\bar{y}$. It turns out truncating the last-layer MLP weights consistently improve test performance for layers $L=3,4,5,6,7$.
2. **More examples of factual recall** in Table 5 with results presented in Appendix C.3, to show the robustness of our arguments. 
3. Re-orgnized sections in Appendix, as suggested by Reviewer U1qi.
4. A discussion on our contributions and implications in Appendix A.

We are very glad to answer any further questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The papers provide an analysis about feed-forward layers and attention layers. The main observation is that FFN is generally responsible for distributional association, whereas attention layers are responsible for in-context reasoning. The authors back up their claim by both empirical and theoretical analysis, although the setups are a little bit toy.

Reviewers generally recognize contributions of this paper. Concerns are raised regarding the simplified setups, but it is understandable that theoretical analysis on full-sized models could be extremely difficult. On the other hand, the implication of the finding could be further discussed.