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?
• Authors: Payel Das, Ching-Yun Ko, Sihui Dai, Georgios Kollias, Subhajit Chaudhury, Aurelie Lozano,
• 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

• 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)
• 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.