Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

Reza Bayat; Ali Behrouz; Vahab Mirrokni; Aaron Courville

cs.LGAug 17, 2026

Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

Authors: Reza Bayat, Ali Behrouz, Vahab Mirrokni, Aaron Courville

Organizations: Mila · Google

Abstract

The quadratic cost of attention-based sequence models for long contexts has motivated a growing line of research on memory-based models that can compress context into a compact state. However, most existing memory models expose a static memory throughout the entire sequence. Because early tokens face no compression pressure, they occupy too many degrees of freedom and "pollute" the memory state, leaving little capacity for later context and increasing interference between what is stored and what arrives next. We study a new paradigm of incremental memory activation, where the effective capacity of memory is progressively expanded as the context grows. Imposing an early bottleneck forces the model to compress history more effectively, while unlocking fresh capacity over time reduces interference and improves retention of later context. We instantiate this paradigm in Proteus, a straightforward mechanism that can be incorporated into a broad class of neural memory architectures at no additional cost. We apply Proteus to state-of-the-art models, including SWLA, Comba, Titans, and Hope-Attention, and observe consistent improvements on standard language modeling and reasoning, as well as on long-context retrieval and understanding, with gains that grow at longer context lengths. Overall, our results show that static memory is suboptimal and that scheduling effective capacity is a simple and broadly applicable tool for sequence modeling.

Figures & tables

Figure 1: Incremental memory activation in Proteus . Proteus progressively expands the active subset of the memory state as the context advances. At each step, the model both reads from and writes to the currently activated blocks (together with previously activated “past” blocks), while the remaining blocks stay locked and do not participate in retrieval or updates. As the context grows, additional blocks are unlocked, providing fresh capacity for later tokens. This imposes an early capacity bottleneck that encourages compression of the initial context, while the fresh capacity introduced over time reduces overwriting and interference for later tokens.

Model	Wiki.	LMB.	LMB.	PIQA	Hella.	Wino.	ARC-e	ARC-c	SIQA	BoolQ	Avg.
	ppl $\downarrow$	ppl $\downarrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc_n $\uparrow$	acc $\uparrow$	acc $\uparrow$	$\uparrow$
760M params / 50B tokens
Transformer++	22.08	22.41	38.2	68.9	44.1	56.8	67.0	34.9	40.2	62.4	51.56
Hope-Attention	20.67	21.36	39.5	70.4	50.1	56.3	66.9	37.5	40.7	63.8	53.15
+ Proteus	19.87	19.72	40.5	70.2	51.9	58.2	67.3	38.5	41.4	63.9	53.99
RetNet	23.54	23.87	35.7	66.4	42.8	53.6	64.8	33.1	38.7	57.4	49.06

Table 1: Language modeling and commonsense reasoning results. We report perplexity on Wikitext and LAMBADA, and zero-shot accuracy on LAMBADA and seven commonsense benchmarks; Avg. denotes the mean accuracy across these eight accuracy columns. Results are shown for 760M (50B tokens) and 1.3B (100B tokens) models, comparing each backbone with and without Proteus (shaded rows). Best average per group in bold .

Figure 2: Long-context retrieval. Accuracy vs. context length on SWDE ( Lockard et al., 2019 ) , SQuAD ( Rajpurkar et al., 2016 ) , and FDA ( Arora et al., 2024 ) for Hope-Attention, Comba, and Titans. Proteus consistently improves robustness at longer contexts, mitigating degradation (notably for Comba and Titans) and improving Hope-Attention across lengths.

	S-NIAH-1			S-NIAH-2			S-NIAH-3
	(pass-key retrieval)			(number in haystack)			(uuid in haystack)
Model	4K	8K	16K	4K	8K	16K	4K	8K	16K
Transformer	97.8	86.2	83.8	100	99.2	96.4	82.2	72.0	48.4
Hope-Attention	100	100	100	100	99.4	96.8	83.4	72.8	50.0
+ Proteus	100	100	100	100	99.4	97.2	83.2	73.2	54.4
Comba	100	100	99.4	92.6	47.2	13.4	62.4	13.8	7.4

Table 2: Needle-In-A-Haystack (NIAH) results. Single-needle tasks at three levels of difficulty—S-NIAH-1 (pass-key retrieval), S-NIAH-2 (numerical needle), and S-NIAH-3 (UUID-based needle)—together with multi-key (MK), multi-query (MQ), and multi-value (MV) settings, at context lengths 4K, 8K, and 16K. Each backbone is shown with and without Proteus (shaded rows).

Model	Narrative	Qasper	MultiField	Hotpot	2WikiMulti	Musique	Average
Hope-Attention	12.1	9.4	19.2	20.6	26.7	6.3	15.72
+ Proteus	12.6	10.2	20.9	21.5	27.3	7.4	16.65
Comba	7.6	10.4	15.9	14.6	23.1	6.7	13.05
+ Proteus	7.8	10.8	16.2	14.5	23.2	6.9	13.23
Titans	8.0	10.3	17.4	15.6	24.8	6.7	13.8
+ Proteus	8.4	10.4	17.9	16.4	25.1	6.7	14.15

Table 3: LongBench evaluation. Performance on LongBench ( Bai et al., 2023 ) across six long-context tasks (Narrative, Qasper, MultiField, Hotpot, 2WikiMulti, Musique), reported for each backbone with and without Proteus (shaded rows); Average denotes the mean across the six tasks. Proteus improves the average score for all three backbones.

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 189 papers

Sessa: Selective State Space Attention
Apr 20, 2026Liubomyr HorbatkoAttention Mechanisms Self-Attention
Raven: High-Recall Sequence Modeling with Sparse Memory Routing
Jul 28, 2026Arshia Afzal, Aviv Bick, Eric P. Xing +2Memory-Augmented Language Models Long-Context Language Modeling
EPFL · Carnegie Mellon University · MBZUAI +1
MMLA: Memory-Mediated Learning Architecture for Predictive Dual-State Adaptation
Jun 27, 2026Junyi Zou, Avrova DonzContinual Learning for LLMs Test-Time Adaptation
1MMLA-org · 2Communication University of China (CUC), No. 1 Dingfuzhuang East Street, Chaoyang District, Beijing 100024, China

Apr 20, 2026J/K move · Enter open · S save

Sessa: Selective State Space Attention

Liubomyr Horbatko

Attention Mechanisms Self-Attention Long-Context Language Modeling Persistent Memory for Language Models State Space Models Recurrent Transformers+3

Open paper AskPDF arXiv

Modern sequence modeling is dominated by two families: Transformers, whose self-attention can access arbitrary elements of the visible sequence, and structured state-space models, which propagate information through an explicit recurrent state. These mechanisms face different limitations on long contexts: when attention is diffuse, the influence of individual tokens is diluted across the effective support, while recurrent state propagation can lose long-range sensitivity unless information is actively preserved. As a result, both mechanisms face challenges in preserving and selectively retrieving information over long contexts. We propose Sessa, a decoder that places attention inside a recurrent feedback path. This creates many attention-based paths through which past tokens can influence future states, rather than relying on a single attention read or a single recurrent chain. We prove that, under explicit assumptions and matched regimes, Sessa admits power-law memory tails

O(\ell^{-β})

for

0 < β< 1

, with slower decay than in the corresponding Transformer and Mamba-style baselines. We further give an explicit construction that achieves this power-law rate. Under the same assumptions, Sessa is the only model class among those considered that realizes flexible selective retrieval, including profiles whose influence does not decay with distance. Consistent with this theoretical advantage, across matched experiments, Sessa achieves the strongest performance on long-context benchmarks while remaining competitive with Transformer and Mamba-style baselines on short-context language modeling.

Abstract

Long-Context Language Modeling

Explore similar work

Sessa: Selective State Space Attention

Raven: High-Recall Sequence Modeling with Sparse Memory Routing

MMLA: Memory-Mediated Learning Architecture for Predictive Dual-State Adaptation