Gated Differentiable Working Memory for Long-Context Language Modeling
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2026
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866914263418798080 |
|---|---|
| author | Mei, Lingrui Liu, Shenghua Wang, Yiwei Ge, Yuyao Bi, Baolong Yao, Jiayu Wan, Jun Yin, Ziling Guo, Jiafeng Cheng, Xueqi |
| author_facet | Mei, Lingrui Liu, Shenghua Wang, Yiwei Ge, Yuyao Bi, Baolong Yao, Jiayu Wan, Jun Yin, Ziling Guo, Jiafeng Cheng, Xueqi |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_12906 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Gated Differentiable Working Memory for Long-Context Language Modeling Mei, Lingrui Liu, Shenghua Wang, Yiwei Ge, Yuyao Bi, Baolong Yao, Jiayu Wan, Jun Yin, Ziling Guo, Jiafeng Cheng, Xueqi Computation and Language 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. |
| title | Gated Differentiable Working Memory for Long-Context Language Modeling |
| topic | Computation and Language |
| url | https://arxiv.org/abs/2601.12906 |