📄 Gated Differentiable Working Memory for Long-Context Language Modeling

Summary

Long inputs remain difficult for transformer LLMs. Attention spreads thin over thousands of tokens, and evidence placed mid-document is often overlooked. Test-time training methods such as qTTT address this by briefly fine-tuning "fast weights" on the given context. However, they sample positions uniformly, so they spend compute on uninformative text and mix conflicting gradients from unrelated parts of the document.

The authors treat this adaptation as a budget allocation problem and propose GDWM (Gated Differentiable Working Memory). The method:

• splits the input into fixed chunks;
• scores each chunk by Contextual Utility, which is how much full-context predictions differ from short-window predictions (an absolute CPMI quantity);
• gives every chunk a minimum number of updates and distributes the remaining gradient steps with a softmax over the scores;
• writes the updates into LoRA adapters on the query and output projections, while the prefilled KV cache stays frozen.

A variance argument based on the Law of Total Variance explains why per-chunk sampling stabilizes training.

On ZeroSCROLLS and LongBench v2 with 1.7B, 4B and 8B models, GDWM equals or beats uniform baselines while using a quarter of the gradient steps. Long-context adaptation therefore becomes notably cheaper.

Key points

• Efficiency: On Qwen3-4B, GDWM with 8 gradient steps roughly matches qTTT with 32 steps. The authors also report a 39% wall-clock speedup.
• Equal-budget comparison (Qwen3-4B, 8× compute): GDWM averages 45.7. The alternatives score lower:
  • qTTT: 40.9
  • Self-Consistency (SC-8): 38.4
  • Beam Search (k=8): 35.3
• Where it helps most: The biggest improvements come on tasks where the relevant information is sparse, reaching +12.7% on Qasper and +11.2% on GovReport, both with the 1.7B model. Smaller models benefit more overall.
• Chunk size: S = 1024 is the most robust setting. At S = 512, MuSiQue falls from 25.8 to 10.0 because multi-hop evidence gets split across chunks.
• Limitations:
  • Gains are small on tasks that need dense coverage of the whole document, e.g., QMSum, NarrativeQA, and QuALITY (0.0% on the 8B model).
  • Fixed-size chunking can cut through semantic units in irregular documents.
  • The method adds hyperparameters; the defaults are τ = 1.0 and k_min = 1.