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

• URL: http://arxiv.org/abs/2503.07903v1
• Date: Mon, 10 Mar 2025 22:48:53 GMT
• Title: Can Memory-Augmented Language Models Generalize on Reasoning-in-a-Haystack Tasks?
• Abstract summary: We propose MemReasoner, a new and simple memory-augmented LLM architecture.<n>MemReasoner is trained end-to-end, with optional supporting fact supervision.<n>We show strong generalization of MemReasoner on both single- and two-hop tasks.
• Score: 31.192647141427233
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Dual-Latent Memory Routing for Vision-Language Reasoning [39.90357280328134]
  Multimodal large language models (MLLMs) have recently made strong progress in vision-language reasoning, yet their performance often degrades as generations grow longer.<n>Inspired by how humans separately recall what they see and what they infer when solving complex tasks, we propose DLMR.<n>It equips MLLMs with Dual Latent Memories: a visual memory that compresses image evidence and a reasoning memory that tracks intermediate conclusions and constraints.
  arXiv  Detail & Related papers  (2026-09-02T10:45:25Z)
• DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning [91.94042219655866]
  Large language models achieve strong performance on many reasoning tasks when allowed to externalize intermediate steps as Chain-of-Thought (CoT)<n>We study this challenge through two-hop reasoning, a representative task where the model must compose multiple pieces of knowledge within a single forward pass.<n>We propose DiscoLoop, a looping architecture whose recurrence carries both a discrete embedding channel and a continuous hidden-state channel.
  arXiv  Detail & Related papers  (2026-07-01T02:32:02Z)
• WorldMemArena: Evaluating Multimodal Agent Memory Through Action-World Interaction [72.1620416874118]
  Multimodal large language models are increasingly deployed as long-horizon agents.<n>Existing benchmarks measure recall over static dialogue, collapse memory into a single end-of-task accuracy, and reduce visual observations to captions.<n>We formulate multimodal agent memory as an Action-World Interaction Loop with an observable four-stage lifecycle.
  arXiv  Detail & Related papers  (2026-05-28T04:27:20Z)
• Structured Reasoning for Large Language Models [59.215789462977206]
  We propose Structured Reasoning (SCR), a framework that decouples reasoning trajectories into explicit, evaluable, and trainable components.<n>SCR substantially improves reasoning efficiency and self-verification.<n>Compared with existing reasoning paradigms, it reduces output token length by up to 50%.
  arXiv  Detail & Related papers  (2026-01-12T04:04:01Z)
• CogMem: A Cognitive Memory Architecture for Sustained Multi-Turn Reasoning in Large Language Models [21.427373172124167]
  Large language models (LLMs) excel at single-turn reasoning but often lose accuracy and coherence over extended, multi-turn interactions.<n>We introduce CogMem, a memory-augmented LLM architecture that supports sustained iterative reasoning through structured, persistent memory.<n> Experiments on TurnBench show that this layered design mitigates reasoning failures, controls context growth, and improves consistency across extended reasoning chains.
  arXiv  Detail & Related papers  (2025-12-16T06:01:08Z)
• Agentic Learner with Grow-and-Refine Multimodal Semantic Memory [50.81667005063605]
  ViLoMem is a dual-stream memory framework that constructs compact, schema-based memory.<n>It encodes visual distraction patterns and logical reasoning errors, enabling MLLMs to learn from their successful and failed experiences.
  arXiv  Detail & Related papers  (2025-11-26T18:55:08Z)
• Evaluating Long-Term Memory for Long-Context Question Answering [100.1267054069757]
  We present a systematic evaluation of memory-augmented methods using LoCoMo, a benchmark of synthetic long-context dialogues annotated for question-answering tasks.<n>Our findings show that memory-augmented approaches reduce token usage by over 90% while maintaining competitive accuracy.
  arXiv  Detail & Related papers  (2025-10-27T18:03:50Z)
• Beyond Context Limits: Subconscious Threads for Long-Horizon Reasoning [11.810780820457927]
  We propose the Thread Inference Model (TIM), a family of large language models, and TIMRUN, an inference runtime.<n>Together, TIM hosted on TIMRUN supports virtually unlimited working memory and multi-hop tool calls within a single language model inference.<n>We show that our system sustains high inference throughput, even when manipulating up to 90% of the KV cache in GPU memory.
  arXiv  Detail & Related papers  (2025-07-22T17:30:04Z)
• MMATH: A Multilingual Benchmark for Mathematical Reasoning [94.05289799605957]
  We introduce MMATH, a benchmark for multilingual complex reasoning spanning 374 high-quality math problems across 10 typologically diverse languages.<n>We observe that even advanced models like DeepSeek R1 exhibit substantial performance disparities across languages and suffer from a critical off-target issue-generating responses in unintended languages.<n>Our findings offer new insights and practical strategies for advancing the multilingual reasoning capabilities of large language models.
  arXiv  Detail & Related papers  (2025-05-25T12:47:39Z)
• The Unreasonable Effectiveness of Model Merging for Cross-Lingual Transfer in LLMs [54.59207567677249]
  Large language models (LLMs) still struggle across tasks outside of high-resource languages.<n>In this work, we investigate cross-lingual transfer to lower-resource languages where task-specific post-training data is scarce.
  arXiv  Detail & Related papers  (2025-05-23T20:28:31Z)
• The Reasoning-Memorization Interplay in Language Models Is Mediated by a Single Direction [34.86855316803838]
  We identify a set of linear features in the model's residual stream that govern the balance between genuine reasoning and memory recall. We show that intervening in these reasoning features helps the model more accurately activate the most relevant problem-solving capabilities during answer generation.
  arXiv  Detail & Related papers  (2025-03-29T14:00:44Z)
• Sketch-of-Thought: Efficient LLM Reasoning with Adaptive Cognitive-Inspired Sketching [60.04718679054704]
  We introduce Sketch-of-Thought (SoT), a novel prompting framework.<n>It combines cognitive-inspired reasoning paradigms with linguistic constraints to minimize token usage.<n>SoT achieves token reductions of 76% with negligible accuracy impact.
  arXiv  Detail & Related papers  (2025-03-07T06:57:17Z)
• Memorize or Generalize? Evaluating LLM Code Generation with Code Rewriting [54.48306552577881]
  We argue that large language models (LLMs) are mostly doing memorization (i.e., replicating or reusing large parts of their training data) versus generalization.<n>Existing evaluations largely proxy neglecting surface/structural similarity, thereby conflating benign reuse of repeated code with harmful recall and memorization task correctness.<n>We propose Memorization Risk Index (MRI), a normalized score that combines two signals: (i) how similar the model's answer for the rewritten task is to the original ground-truth solution, and (ii) how much performance drops from the original task to its rewritten counterpart.
  arXiv  Detail & Related papers  (2025-03-04T05:39:24Z)
• Unveiling Reasoning Thresholds in Language Models: Scaling, Fine-Tuning, and Interpretability through Attention Maps [3.8936716676293917]
  This study investigates the in-context learning capabilities of various decoder-only transformer-based language models with different model sizes and training data. We identify a critical parameter threshold (1.6 billion), beyond which reasoning performance improves significantly in tasks such as commonsense reasoning in multiple-choice question answering and deductive reasoning.
  arXiv  Detail & Related papers  (2025-02-21T00:48:32Z)
• Mixture of Parrots: Experts improve memorization more than reasoning [72.445819694797]
  We show that as we increase the number of experts, the memorization performance consistently increases while the reasoning capabilities saturate. We find that increasing the number of experts helps solve knowledge-intensive tasks, but fails to yield the same benefits for reasoning tasks.
  arXiv  Detail & Related papers  (2024-10-24T17:54:41Z)
• Generalization v.s. Memorization: Tracing Language Models' Capabilities Back to Pretraining Data [76.90128359866462]
  We introduce an extended concept of memorization, distributional memorization, which measures the correlation between the output probabilities and the pretraining data frequency.<n>This study demonstrates that memorization plays a larger role in simpler, knowledge-intensive tasks, while generalization is the key for harder, reasoning-based tasks.
  arXiv  Detail & Related papers  (2024-07-20T21:24:40Z)
• Uncertainty Guided Global Memory Improves Multi-Hop Question Answering [3.7013865226473848]
  We propose a two-stage method that first collects relevant information over the entire document to the memory and then combines it with local context to solve the task. Our experimental results show that fine-tuning a pre-trained model with memory-augmented input, including the most certain global elements, improves the model's performance.
  arXiv  Detail & Related papers  (2023-11-29T23:45:57Z)
• Memory Injections: Correcting Multi-Hop Reasoning Failures during Inference in Transformer-Based Language Models [4.343604069244352]
  We propose an approach to pinpoint and rectify multi-hop reasoning failures through targeted memory injections on attention heads. We show that a simple, efficient, and targeted memory injection into a key attention layer can often increase the probability of the desired next token in multi-hop tasks, by up to 424%.
  arXiv  Detail & Related papers  (2023-09-11T16:39:30Z)
• Think Before You Act: Decision Transformers with Working Memory [44.18926449252084]
  Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. We propose a working memory module to store, blend, and retrieve information for different downstream tasks.
  arXiv  Detail & Related papers  (2023-05-24T01:20:22Z)
• Uni-Perceiver v2: A Generalist Model for Large-Scale Vision and Vision-Language Tasks [86.66733026149892]
  We propose Uni-Perceiver v2, which is the first generalist model capable of handling major large-scale vision and vision-gnostic tasks. Specifically, images are encoded as general region proposals, while texts are encoded via a Transformer-based language model. Uni-Perceiver v2 achieves competitive performance on a broad range of vision and vision-language tasks.
  arXiv  Detail & Related papers  (2022-11-17T18:59:52Z)
• Memory-Guided Semantic Learning Network for Temporal Sentence Grounding [55.31041933103645]
  We propose a memory-augmented network that learns and memorizes the rarely appeared content in TSG tasks. MGSL-Net consists of three main parts: a cross-modal inter-action module, a memory augmentation module, and a heterogeneous attention module.
  arXiv  Detail & Related papers  (2022-01-03T02:32:06Z)
• Reparameterizing Convolutions for Incremental Multi-Task Learning without Task Interference [75.95287293847697]
  Two common challenges in developing multi-task models are often overlooked in literature. First, enabling the model to be inherently incremental, continuously incorporating information from new tasks without forgetting the previously learned ones (incremental learning) Second, eliminating adverse interactions amongst tasks, which has been shown to significantly degrade the single-task performance in a multi-task setup (task interference)
  arXiv  Detail & Related papers  (2020-07-24T14:44:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.