Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

G-MemLLM: Gated Latent Memory Augmentation for Long-Context Reasoning in Large Language Models

Xun Xu

TL;DR

G-MemLLM addresses the bottleneck of fixed context windows in large language models by introducing a trainable Latent Memory Bank managed through a GRU-style gate. The architecture maintains a frozen LLM backbone while selectively updating persistent memory slots via a Memory Loop that includes Extraction, Retrieval, Injection, and Consolidation, with a loss that jointly optimizes task performance and memory efficiency. Empirical results across GPT-2 and Llama 3.1-8B show significant gains on HotpotQA and ZsRE, including an absolute 13.3 percentage-point improvement on ZsRE for the 8B model and notable gains in Answer and Sup Fact F1 on HotpotQA; 1024 memory slots emerge as a sweet spot. These findings suggest that explicit, gated latent memory becomes increasingly beneficial as model scale grows, enabling more robust multi-hop reasoning with minimal parameter overhead.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in natural language understanding, yet they remain constrained by the finite capacity of their context windows and the inherent difficulty of maintaining long-term factual consistency during multi-hop reasoning. While existing methods utilize context compression or recurrent tokens, they often suffer from ``context rot'' or the dilution of information over long horizons. In this paper, we propose \textbf{G-MemLLM}, a memory-augmented architecture that integrates a frozen LLM backbone with a trainable \textbf{Latent Memory Bank}. Our key innovation is a GRU-style gated update logic that allows the model to selectively update, preserve, or overwrite latent memory slots, preventing the vanishing gradients of knowledge common in recurrent systems. We evaluate G-MemLLM across scales, from GPT-2 (124M) to Llama 3.1 (8B), on the HotpotQA and Zero-Shot Relation Extraction (ZsRE) benchmarks. Our results demonstrate that G-MemLLM significantly enhances multi-hop reasoning and relational precision, achieving a 13.3\% accuracy boost on ZsRE for Llama 3.1-8B, and it also yields improvements across model scales, boosting Answer F1 by 8.56 points for GPT-2 and increasing Supporting Fact F1 by 6.89 points for Llama 3.1-8B on HotpotQA.

G-MemLLM: Gated Latent Memory Augmentation for Long-Context Reasoning in Large Language Models

TL;DR

G-MemLLM addresses the bottleneck of fixed context windows in large language models by introducing a trainable Latent Memory Bank managed through a GRU-style gate. The architecture maintains a frozen LLM backbone while selectively updating persistent memory slots via a Memory Loop that includes Extraction, Retrieval, Injection, and Consolidation, with a loss that jointly optimizes task performance and memory efficiency. Empirical results across GPT-2 and Llama 3.1-8B show significant gains on HotpotQA and ZsRE, including an absolute 13.3 percentage-point improvement on ZsRE for the 8B model and notable gains in Answer and Sup Fact F1 on HotpotQA; 1024 memory slots emerge as a sweet spot. These findings suggest that explicit, gated latent memory becomes increasingly beneficial as model scale grows, enabling more robust multi-hop reasoning with minimal parameter overhead.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in natural language understanding, yet they remain constrained by the finite capacity of their context windows and the inherent difficulty of maintaining long-term factual consistency during multi-hop reasoning. While existing methods utilize context compression or recurrent tokens, they often suffer from ``context rot'' or the dilution of information over long horizons. In this paper, we propose \textbf{G-MemLLM}, a memory-augmented architecture that integrates a frozen LLM backbone with a trainable \textbf{Latent Memory Bank}. Our key innovation is a GRU-style gated update logic that allows the model to selectively update, preserve, or overwrite latent memory slots, preventing the vanishing gradients of knowledge common in recurrent systems. We evaluate G-MemLLM across scales, from GPT-2 (124M) to Llama 3.1 (8B), on the HotpotQA and Zero-Shot Relation Extraction (ZsRE) benchmarks. Our results demonstrate that G-MemLLM significantly enhances multi-hop reasoning and relational precision, achieving a 13.3\% accuracy boost on ZsRE for Llama 3.1-8B, and it also yields improvements across model scales, boosting Answer F1 by 8.56 points for GPT-2 and increasing Supporting Fact F1 by 6.89 points for Llama 3.1-8B on HotpotQA.
Paper Structure (16 sections, 2 equations, 3 figures, 3 tables)

This paper contains 16 sections, 2 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Context Compression and Information Bottlenecks
  4. Memory Tokens and Recurrent State-Passing
  5. Dynamic Latent Memory and Gated Architectures
  6. Method
  7. Latent Memory Bank
  8. System Dynamics: The Memory Loop
  9. Training Objective
  10. Experiments
  11. Datasets
  12. Baselines
  13. Results
  14. Scaling Analysis
  15. Ablation Study: Memory Density
  16. ...and 1 more sections

Figures (3)

  • Figure 1: Performance enhancement brought by the G-MemLLM memory module. (a) illustrates the absolute improvement across models, while (b) shows the direct comparison between Vanilla and G-MemLLM on specific metrics.
  • Figure 2: Overview of G-MemLLM architecture.
  • Figure 3: Trend of performance gains and computational overhead.