# Can Memory-Augmented Language Models Generalize on Reasoning-in-a-Haystack Tasks?

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge- and tool-augmented reasoning, one of 17 directions in LLM reasoning.
- Builds on 28 of its 41 references that are in the map.
- First posted on arXiv 10 Mar 2025.

## Placed in

- LLM reasoning / [Knowledge- and tool-augmented reasoning](https://mlsanity.com/lr/knowledge-and-tool-augmented-reasoning/)

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

- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Beyond the Imitation Game: Quantifying and extrapolating the capabilities of language models](https://mlsanity.com/paper/2206.04615/) (2022)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [Towards AI-Complete Question Answering: A Set of Prerequisite Toy Tasks](https://mlsanity.com/paper/1502.05698/) (2015)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Iterative Reasoning Preference Optimization](https://mlsanity.com/paper/2404.19733/) (2024)
- [BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack](https://mlsanity.com/paper/2406.10149/) (2024)
- [An Empirical Study of Mamba-based Language Models](https://mlsanity.com/paper/2406.07887/) (2024)
- [Same Task, More Tokens: the Impact of Input Length on the Reasoning Performance of Large Language Models](https://mlsanity.com/paper/2402.14848/) (2024)
- [Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks](https://mlsanity.com/paper/2402.04248/) (2024)
- [Memory Transformer](https://mlsanity.com/paper/2006.11527/) (2020)
- [Exposing Attention Glitches with Flip-Flop Language Modeling](https://mlsanity.com/paper/2306.00946/) (2023)
- [Larimar: Large Language Models with Episodic Memory Control](https://mlsanity.com/paper/2403.11901/) (2024)
- [GRACE: Discriminator-Guided Chain-of-Thought Reasoning](https://mlsanity.com/paper/2305.14934/) (2023)
- [NeedleBench: Evaluating LLM Retrieval and Reasoning Across Varying Information Densities](https://mlsanity.com/paper/2407.11963/) (2024)
- [Not All LLM Reasoners Are Created Equal](https://mlsanity.com/paper/2410.01748/) (2024)
- [Cofca: A Step-Wise Counterfactual Multi-hop QA benchmark](https://mlsanity.com/paper/2402.11924/) (2024)
- [LogicAsker: Evaluating and Improving the Logical Reasoning Ability of Large Language Models](https://mlsanity.com/paper/2401.00757/) (2024)
- [Investigating Multi-Hop Factual Shortcuts in Knowledge Editing of Large Language Models](https://mlsanity.com/paper/2402.11900/) (2024)
- [End-To-End Memory Networks](https://mlsanity.com/paper/1503.08895/) (2015)
- [Generation Constraint Scaling Can Mitigate Hallucination](https://mlsanity.com/paper/2407.16908/) (2024)

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

- [Cognitive Foundations for Reasoning and Their Manifestation in LLMs](https://mlsanity.com/paper/2511.16660/) (2025)
- [Retrieval--Reasoning Processes for Multi-hop Question Answering: A Four-Axis Design Framework and Empirical Trends](https://mlsanity.com/paper/2601.00536/) (2026)

## Abstract (arXiv)

Large language models often expose their brittleness in reasoning tasks, especially while executing long chains of reasoning over context. We propose MemReasoner, a new and simple memory-augmented LLM architecture, in which the memory learns the relative order of facts in context, and enables hopping over them, while the decoder selectively attends to the memory. MemReasoner is trained end-to-end, with optional supporting fact supervision of varying degrees. We train MemReasoner, along with existing memory-augmented transformer models and a state-space model, on two distinct synthetic multi-hop reasoning tasks. Experiments performed under a variety of challenging scenarios, including the presence of long distractor text or target answer changes in test set, show strong generalization of MemReasoner on both single- and two-hop tasks. This generalization of MemReasoner is achieved using none-to-weak supporting fact supervision (using none and 1% of supporting facts for one- and two-hop tasks, respectively). In contrast, baseline models overall struggle to generalize and benefit far less from using full supporting fact supervision. The results highlight the importance of explicit memory mechanisms, combined with additional weak supervision, for improving large language model's context processing ability toward reasoning tasks.