Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Gated Differentiable Working Memory for Long-Context Language Modeling

Lingrui Mei, Shenghua Liu, Yiwei Wang, Yuyao Ge, Baolong Bi, Jiayu Yao, Jun Wan, Ziling Yin, Jiafeng Guo, Xueqi Cheng

TL;DR

Gated Differentiable Working Memory (Gdwm) reframes test-time adaptation for long-context LLMs as budget-constrained memory consolidation, introducing a Write Controller that gates updates based on Contextual Utility. Contextual Utility, grounded in CPMI, identifies high-value long-range dependencies, enabling a budget-aware allocation that guarantees global coverage across context chunks. Through chunk-restricted consolidation of LoRA adapters with a prefilled KV cache, Gdwm reduces gradient steps by up to $4\times$ while maintaining or improving performance, achieving a notable speedup of approximately $39\%$ in wall-clock time on ZeroSCROLLS and LongBench v2. The approach yields the strongest gains on sparse-information tasks, provides ablations validating each component, and offers theoretical variance-reduction guarantees via chunked updates, representing a meaningful advance in efficient long-context adaptation for LLMs.

Abstract

Long contexts challenge transformers: attention scores dilute across thousands of tokens, critical information is often lost in the middle, and models struggle to adapt to novel patterns at inference time. Recent work on test-time adaptation addresses this by maintaining a form of working memory -- transient parameters updated on the current context -- but existing approaches rely on uniform write policies that waste computation on low-utility regions and suffer from high gradient variance across semantically heterogeneous contexts. In this work, we reframe test-time adaptation as a budget-constrained memory consolidation problem, focusing on which parts of the context should be consolidated into working memory under limited computation. We propose Gdwm (Gated Differentiable Working Memory), a framework that introduces a write controller to gate the consolidation process. The controller estimates Contextual Utility, an information-theoretic measure of long-range contextual dependence, and allocates gradient steps accordingly while maintaining global coverage. Experiments on ZeroSCROLLS and LongBench v2 demonstrate that Gdwm achieves comparable or superior performance with 4$\times$ fewer gradient steps than uniform baselines, establishing a new efficiency-performance Pareto frontier for test-time adaptation.

Gated Differentiable Working Memory for Long-Context Language Modeling

TL;DR

Gated Differentiable Working Memory (Gdwm) reframes test-time adaptation for long-context LLMs as budget-constrained memory consolidation, introducing a Write Controller that gates updates based on Contextual Utility. Contextual Utility, grounded in CPMI, identifies high-value long-range dependencies, enabling a budget-aware allocation that guarantees global coverage across context chunks. Through chunk-restricted consolidation of LoRA adapters with a prefilled KV cache, Gdwm reduces gradient steps by up to while maintaining or improving performance, achieving a notable speedup of approximately in wall-clock time on ZeroSCROLLS and LongBench v2. The approach yields the strongest gains on sparse-information tasks, provides ablations validating each component, and offers theoretical variance-reduction guarantees via chunked updates, representing a meaningful advance in efficient long-context adaptation for LLMs.

Abstract

Long contexts challenge transformers: attention scores dilute across thousands of tokens, critical information is often lost in the middle, and models struggle to adapt to novel patterns at inference time. Recent work on test-time adaptation addresses this by maintaining a form of working memory -- transient parameters updated on the current context -- but existing approaches rely on uniform write policies that waste computation on low-utility regions and suffer from high gradient variance across semantically heterogeneous contexts. In this work, we reframe test-time adaptation as a budget-constrained memory consolidation problem, focusing on which parts of the context should be consolidated into working memory under limited computation. We propose Gdwm (Gated Differentiable Working Memory), a framework that introduces a write controller to gate the consolidation process. The controller estimates Contextual Utility, an information-theoretic measure of long-range contextual dependence, and allocates gradient steps accordingly while maintaining global coverage. Experiments on ZeroSCROLLS and LongBench v2 demonstrate that Gdwm achieves comparable or superior performance with 4 fewer gradient steps than uniform baselines, establishing a new efficiency-performance Pareto frontier for test-time adaptation.
Paper Structure (85 sections, 2 theorems, 23 equations, 6 figures, 6 tables, 1 algorithm)

This paper contains 85 sections, 2 theorems, 23 equations, 6 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Context Engineering
  4. LLM Memory
  5. Test-Time Adaptation
  6. Problem Formulation
  7. Fast Weights.
  8. Budget-Constrained Optimization.
  9. Method
  10. Memory Interface
  11. Contextual Utility Estimation
  12. Interpretation.
  13. Information-Theoretic Grounding.
  14. Budget-Aware Allocation Policy
  15. Structured Memory Consolidation
  16. ...and 70 more sections

Key Result

Theorem 1

Let $g$ be the gradient estimator for a single update step. Assume the document consists of $M$ chunks, where chunk $c$ has gradient mean $\mu_c$ and variance $\sigma_c^2$. Under global uniform sampling: Under chunk-restricted sampling (conditioned on chunk $c$), we have $\mathrm{Var}(g \mid c) = \sigma_c^2$, eliminating the inter-chunk variance term. Equality holds only when $\mu_c = \bar{\mu}$

Figures (6)

  • Figure 1: Efficiency vs Performance on ZeroSCROLLS (Qwen3-4B).Gdwm achieves comparable performance to qTTT-32 with only 8 gradient steps (4$\times$ fewer), establishing a new Pareto frontier. Context-aware budget allocation enables faster convergence than uniform or sampling-based alternatives.
  • Figure 2: High-level overview of Gdwm. The framework proceeds in four stages: Chunk the input into fixed-size units (approximating semantic segments), Gate each chunk via Contextual Utility (CPMI-based divergence), Allocate gradient budget proportionally subject to coverage constraints, and Consolidate into LoRA adapters.
  • Figure 3: Technical details of Gdwm. Left: Prefill-and-freeze KV cache enables efficient chunk-wise processing. Middle: Contextual Utility is computed as CPMI between global and local predictions, then converted to budget weights via softmax allocation. Right: LoRA adapters on $W_Q$/$W_O$ projections are updated through chunk-wise next-token prediction loss.
  • Figure 4: Task-wise Performance on LongBench v2. Radar charts comparing Gdwm-32 (light blue) against baselines on 4B (left) and 8B (right) models. Gdwm achieves consistent improvements on Code Repositories and Multi-Doc QA where information is sparse and localized, while showing competitive performance on Long Dialogue where global coverage is required.
  • Figure 5: Time Breakdown per Sample (Qwen3-4B). Generation dominates at 53%, while CPMI computation accounts for only 13%---demonstrating the lightweight nature of our context selection mechanism. The 4$\times$ reduction in gradient steps (32$\rightarrow$8) yields 39% net wall-clock speedup.
  • ...and 1 more figures

Theorems & Definitions (6)

  • Definition 1: Contextual Utility
  • Theorem 1: Variance Reduction
  • Definition 2: Evidence Span
  • Definition 3: Contextual Utility as Mutual Information
  • Proposition 1: Utility Underestimation Under Fragmentation
  • proof