# Gated Differentiable Working Memory for Long-Context Language Modeling

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 146 of its 198 references that are in the map.
- First posted on arXiv 19 Jan 2026.

## Builds on (most cited first, 30 of 146)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [A Survey on Large Language Model based Autonomous Agents](https://mlsanity.com/paper/2308.11432/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [A-MEM: Agentic Memory for LLM Agents](https://mlsanity.com/paper/2502.12110/) (2025)
- [ExpeL: LLM Agents Are Experiential Learners](https://mlsanity.com/paper/2308.10144/) (2023)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [Passage Re-ranking with BERT](https://mlsanity.com/paper/1901.04085/) (2019)
- [A Survey on the Memory Mechanism of Large Language Model based Agents](https://mlsanity.com/paper/2404.13501/) (2024)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [Efficient Test-Time Model Adaptation without Forgetting](https://mlsanity.com/paper/2204.02610/) (2022)
- [Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time](https://mlsanity.com/paper/2305.17118/) (2023)
- [Entity-Based Knowledge Conflicts in Question Answering](https://mlsanity.com/paper/2109.05052/) (2021)
- [Adapting Language Models to Compress Contexts](https://mlsanity.com/paper/2305.14788/) (2023)
- [LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression](https://mlsanity.com/paper/2403.12968/) (2024)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Trusting Your Evidence: Hallucinate Less with Context-aware Decoding](https://mlsanity.com/paper/2305.14739/) (2023)
- [Memory in the Age of AI Agents](https://mlsanity.com/paper/2512.13564/) (2025)
- [RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation](https://mlsanity.com/paper/2310.04408/) (2023)
- [Evaluating Memory in LLM Agents via Incremental Multi-Turn Interactions](https://mlsanity.com/paper/2507.05257/) (2025)
- [AutoPrompt: Eliciting Knowledge from Language Models with Automatically Generated Prompts](https://mlsanity.com/paper/2010.15980/) (2020)
- [MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent](https://mlsanity.com/paper/2507.02259/) (2025)

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

- [Semantic-Aware Adaptive Visual Memory for Streaming Video Understanding](https://mlsanity.com/paper/2605.07897/) (2026)
- [VideoStir: Understanding Long Videos via Spatio-Temporally Structured and Intent-Aware RAG](https://mlsanity.com/paper/2604.05418/) (2026)
- [Memory for Large Language Models](https://mlsanity.com/paper/2607.25380/) (2026)
- [Self-Manager: Parallel Agent Loop for Long-form Deep Research](https://mlsanity.com/paper/2601.17879/) (2026)
- [SR-TTT Does Not Learn Retrieval: A Correction and Mechanistic Post-Mortem of Surprisal-Aware Residual Test-Time Training](https://mlsanity.com/paper/2603.06642/) (2026)
- [V-DEAL: Diagnosing Video Safety De-Calibration as an Understanding-Refusal Coupling Failure](https://mlsanity.com/paper/2607.21151/) (2026)

## Abstract (arXiv)

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.