# Associative memory inspires improvements for in-context learning using a novel attention residual stream architecture

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 165 references that are in the map.
- First posted on arXiv 19 Dec 2024.

## Builds on (most cited first, 30 of 48)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets](https://mlsanity.com/paper/2201.02177/) (2022)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [TinyStories: How Small Can Language Models Be and Still Speak Coherent English?](https://mlsanity.com/paper/2305.07759/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Dense Associative Memory for Pattern Recognition](https://mlsanity.com/paper/1606.01164/) (2016)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Language Conditioned Imitation Learning over Unstructured Data](https://mlsanity.com/paper/2005.07648/) (2020)
- [Linearity of Relation Decoding in Transformer Language Models](https://mlsanity.com/paper/2308.09124/) (2023)
- [Large Language Models as General Pattern Machines](https://mlsanity.com/paper/2307.04721/) (2023)
- [ResViT: Residual vision transformers for multi-modal medical image synthesis](https://mlsanity.com/paper/2106.16031/) (2021)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Nemotron-CC: Transforming Common Crawl into a Refined Long-Horizon Pretraining Dataset](https://mlsanity.com/paper/2412.02595/) (2024)
- [The mechanistic basis of data dependence and abrupt learning in an in-context classification task](https://mlsanity.com/paper/2312.03002/) (2023)
- [Transformers represent belief state geometry in their residual stream](https://mlsanity.com/paper/2405.15943/) (2024)
- [Attractor and integrator networks in the brain](https://mlsanity.com/paper/2112.03978/) (2021)
- [Universal Hopfield Networks: A General Framework for Single-Shot Associative Memory Models](https://mlsanity.com/paper/2202.04557/) (2022)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)

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

- [Swallowing the Poison Pills: Insights from Vulnerability Disparity Among LLMs](https://mlsanity.com/paper/2502.18518/) (2025)
- [Large Language Models as Model Organisms for Human Associative Learning](https://mlsanity.com/paper/2510.21408/) (2025)
- [Context-Gated Associative Retrieval: From Theory to Transformers](https://mlsanity.com/paper/2605.10970/) (2026)

## Abstract (arXiv)

Large language models (LLMs) demonstrate an impressive ability to utilise information within the context of their input sequences to appropriately respond to data unseen by the LLM during its training procedure. This ability is known as in-context learning (ICL). Humans and non-human animals demonstrate similar abilities, however their neural architectures differ substantially from LLMs. Despite this, a critical component within LLMs, the attention mechanism, resembles modern associative memory models, widely used in and influenced by the computational neuroscience community to model biological memory systems. Using this connection, we introduce an associative memory model capable of performing ICL. We use this as inspiration for a novel residual stream architecture which allows information to directly flow between attention heads. We test this architecture during training within a two-layer Transformer and show its ICL abilities manifest more quickly than without this modification. We then apply our architecture in small language models with 8 million and 1 billion parameters, focusing on attention head values, with results also indicating improved performance at these larger and more naturalistic scales.