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

arXiv: https://arxiv.org/abs/2406.03068 · HTML: https://mlsanity.com/paper/2406.03068/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 17 of its 53 references that are in the map.
- First posted on arXiv 5 Jun 2024.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)

## Builds on (most cited first, 17 of 17)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models](https://mlsanity.com/paper/2206.04615/) (2022)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction](https://mlsanity.com/paper/2312.13558/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [Augmenting Self-attention with Persistent Memory](https://mlsanity.com/paper/1907.01470/) (2019)
- [In-Context Language Learning: Architectures and Algorithms](https://mlsanity.com/paper/2401.12973/) (2024)
- [Transformers, parallel computation, and logarithmic depth](https://mlsanity.com/paper/2402.09268/) (2024)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [Approximating How Single Head Attention Learns](https://mlsanity.com/paper/2103.07601/) (2021)
- [Learning Associative Memories with Gradient Descent](https://mlsanity.com/paper/2402.18724/) (2024)
- [Training Dynamics of Contextual N-Grams in Language Models](https://mlsanity.com/paper/2311.00863/) (2023)

## Built on by (most cited first, 12 of 12)

- [Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformer](https://mlsanity.com/paper/2506.01115/) (2025)
- [HTMuon: Improving Muon via Heavy-Tailed Spectral Correction](https://mlsanity.com/paper/2603.10067/) (2026)
- [On the Emergence of Induction Heads for In-Context Learning](https://mlsanity.com/paper/2511.01033/) (2025)
- [Transmuting prompts into weights](https://mlsanity.com/paper/2510.08734/) (2025)
- [Facts in Stats: Impacts of Pretraining Diversity on Language Model Generalization](https://mlsanity.com/paper/2510.16096/) (2025)
- [Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledge](https://mlsanity.com/paper/2603.20969/) (2026)
- [Sparse Growing Transformer: Training-Time Sparse Depth Allocation via Progressive Attention Looping](https://mlsanity.com/paper/2603.23998/) (2026)
- [Hessian-Enhanced Token Attribution (HETA): Interpreting Autoregressive LLMs](https://mlsanity.com/paper/2604.13258/) (2026)
- [Transformers Provably Learn to Internalize Chain-of-Thought](https://mlsanity.com/paper/2605.28600/) (2026)
- [How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalization](https://mlsanity.com/paper/2505.15009/) (2025)
- [Rethinking Associative Memory Mechanism in Induction Head](https://mlsanity.com/paper/2412.11459/) (2024)
- [GSM-Agent: Understanding Agentic Reasoning Using Controllable Environments](https://mlsanity.com/paper/2509.21998/) (2025)

## Abstract (arXiv)

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.