🤖 machine learning

Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

The paper introduces Proteus, a novel paradigm of incremental memory activation that progressively expands effective memory capacity to mitigate early token interference and improve retention, demonstrating consistent performance gains across various state-of-the-art long-context models.

Original authors: Reza Bayat, Ali Behrouz, Vahab Mirrokni, Aaron Courville

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Reza Bayat, Ali Behrouz, Vahab Mirrokni, Aaron Courville

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast landscape of artificial intelligence, a persistent challenge has long defined the limits of what machines can remember. For years, the most powerful systems relied on a mechanism that treated every piece of information it encountered with equal weight, storing a complete, unchanging record of everything it had ever read. While this approach allowed for incredible precision in recalling specific details, it came with a steep price: the more information the system processed, the more computationally expensive it became, eventually slowing to a crawl when faced with long stories or complex documents. To solve this, researchers turned to a different strategy, designing models that compress history into a single, fixed-size summary. This made them fast and efficient, but it introduced a new problem. Because the memory space was always fully open from the very first moment, the system tended to fill up with the earliest details it saw, leaving little room for the new information that arrived later. The result was a machine that remembered the beginning of a conversation perfectly but struggled to hold onto the end.

A team of researchers has now proposed a solution to this imbalance, introducing a method they call incremental memory activation. Instead of keeping the entire memory bank open and available from the start, their approach gradually unlocks the memory as the sequence of information grows. Imagine a library where the shelves are initially locked, forcing the librarian to summarize the first few books into a small, tight space. As more books arrive, new shelves are unlocked, providing fresh space for the later stories without erasing the summaries of the earlier ones. This simple shift in timing forces the system to compress the early context effectively, while ensuring that later information has room to be stored without interfering with what came before. The researchers tested this idea on several of the most advanced memory-based models currently in existence, finding that it consistently improved their ability to understand long texts and retrieve specific details from deep within them.

The core of this work addresses a specific failure mode in how these machines learn. When a model is allowed to use its full memory capacity immediately, the very first pieces of information it encounters face no competition for space. They can occupy a disproportionate amount of the system's attention, effectively "polluting" the memory state. By the time later information arrives, the memory is already saturated, and new details must squeeze in, often overwriting or distorting the earlier data. The researchers argue that this static approach is suboptimal. Their new paradigm, named Proteus, introduces a dynamic schedule where the effective capacity of the memory is not fixed but grows in step with the length of the input. At the beginning of a sequence, the system is forced to work with a restricted subset of its memory, acting as a bottleneck that encourages it to summarize and compress the initial context rather than memorizing it in detail. As the sequence continues, additional blocks of memory are progressively unlocked, offering fresh capacity for new information. This ensures that later tokens do not have to fight for space against the early ones, reducing interference and improving the retention of the most recent context.

To test this concept, the team applied the Proteus mechanism to a diverse set of state-of-the-art models, including architectures known as SWLA, Comba, Titans, and Hope-Attention. These models were trained on massive datasets containing billions of words, with the researchers comparing their performance with and without the new gating system. The results showed a clear and consistent pattern: adding Proteus improved the models' performance across the board. In standard language tasks, the models produced more accurate predictions and better reasoning scores. The improvements were particularly noticeable in long-context scenarios. When asked to retrieve a specific piece of information hidden within a very long document—a task often called "needle in a haystack"—the models using Proteus maintained their accuracy much better than their standard counterparts as the text length increased. For instance, on tasks involving documents up to 16,000 words long, the models with the incremental activation schedule showed significant gains in finding the correct information, whereas the baseline models saw their performance degrade sharply.

The study also explored how this principle could apply beyond just the memory state of the model, extending it to the model's internal parameters themselves. By treating the learning process of the model's internal layers as a form of memory, the researchers applied the same scheduling logic to how the model updates its own knowledge. This extension allowed the method to be used in architectures that do not rely on a traditional recurrent memory state, further demonstrating the flexibility of the approach. Across all experiments, the method required no additional memory storage or extra computational cost; it simply changed the timing of when different parts of the system were allowed to participate in learning and retrieval. The findings suggest that the way capacity is allocated over time is just as important as the architecture itself. By treating memory not as a static resource but as a schedule that evolves with the context, the researchers have provided a simple, broadly applicable tool that helps machines manage long sequences more effectively, proving that sometimes, the best way to remember everything is to open the doors slowly.

Technical Summary: Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

Problem Statement

The quadratic computational cost of attention-based Transformers has driven research toward memory-based models that compress context into a fixed-size state (e.g., modern RNNs). However, existing memory models typically expose a static memory capacity throughout the entire sequence. This design leads to two primary failure modes:

Early Under-Compression: Early tokens arrive when the memory state is nearly empty. Facing no competition for capacity, they occupy a disproportionate share of the memory's degrees of freedom without being forced to compress, effectively "polluting" the state.
Later Interference: As the sequence progresses, later tokens must be squeezed into the remaining capacity. This forces new information to overwrite or interfere with previously stored content, degrading the model's ability to retain later context.

The authors argue that static capacity is suboptimal because the "right amount" of capacity differs between the beginning and end of a sequence. Early inputs require a bottleneck to enforce compression, while later inputs require fresh capacity to avoid overwriting.

Methodology: Incremental Memory Activation

The paper proposes a new paradigm called Incremental Memory Activation, where the effective capacity of the memory is scheduled as a function of the context length rather than being fixed. This paradigm addresses the failure modes through two complementary mechanisms:

Early Bottleneck: By restricting the active memory subset at the start of the sequence, the model is forced to summarize (compress) initial context rather than memorize it.
Fresh Capacity: As the context grows, previously locked memory blocks are progressively unlocked. This provides fresh degrees of freedom for new information, reducing interference with earlier stored content.

The Proteus Mechanism

The authors instantiate this paradigm in Proteus, a lightweight, block-wise gating scheme applicable to a broad class of associative-memory architectures (including linear RNNs and models with MLP-based memory) at no additional parameter cost.

Core Components:

Block Partitioning: The memory parameters (e.g., a matrix $M \in \mathbb{R}^{d_k \times d_v}$ ) are partitioned into $E$ equal-sized contiguous blocks.
Deterministic Schedule: A simple, uniform schedule determines the number of active blocks $k(t)$ $k (t)$ at step $t$ $t$ . The active fraction grows in uniform increments across the context length $N$ $N$ .
- Tokens in the early segment $[1, \Delta]$ interact only with the first $1/E$ of the memory.
- As $t$ increases, additional blocks are unlocked until the full memory is available.
Gated Operations:
- Write: The update rule (e.g., gradient descent) is applied only to the currently active blocks. Locked blocks remain fixed and do not receive updates.
- Read: Retrieval operations read only from the active subset of the memory.
- Mathematical Formulation: For a memory state $M_{t-1}$ and activation mask $g_t$ , the update is:
  $M_t = M_{t-1} - \theta_t (g_t \odot \nabla \tilde{L}(M^{(g)}_{t-1}; k_t, v_t))$
  where $M^{(g)}_{t-1}$ is the active subset and $\odot$ denotes element-wise multiplication.

Extension to Model Parameters (Nested Learning):
Leveraging the "Nested Learning" perspective (Behrouz et al., 2025d), which views MLP blocks as associative memory, Proteus extends this scheduling principle to the model's parameters themselves. In architectures like Hope-Attention, the update of MLP block parameters is gated by the same schedule, ensuring that early training samples do not overfit the full parameter space before later samples arrive.

Key Contributions

Incremental Memory Activation Paradigm: A novel approach to controlling effective model capacity over the context flow, encouraging early compression and reducing later overwrite/interference.
Proteus Mechanism: A general, drop-in block-wise gating scheme that applies incremental activation to both memory reads/writes and, via Nested Learning, to MLP parameters. It requires no additional parameters or memory overhead.
Empirical Validation: Demonstration that Proteus consistently improves state-of-the-art models (SWLA, Comba, Titans, Hope-Attention) across multiple benchmarks.

Experimental Results

The authors evaluated Proteus on four model families (Hope-Attention, SWLA, Comba, Titans) at 760M and 1.3B parameter scales, trained on 50B and 100B tokens respectively.

Language Modeling & Commonsense Reasoning: Proteus consistently lowered perplexity on Wikitext and LAMBADA and improved zero-shot accuracy on seven commonsense benchmarks (PIQA, HellaSwag, etc.). The gains were observed across all backbones, with Titans+Proteus achieving the best overall performance (e.g., 58.00 average accuracy at 1.3B).
Needle-in-a-Haystack (NIAH): Improvements were most pronounced at longer context lengths (16K) and on harder variants (e.g., UUID retrieval). For instance, Proteus improved Titans' performance on S-NIAH-3 at 16K from 21.4% to 29.8%. At short contexts where base models already saturated, Proteus was neutral.
Long-Context Retrieval & Understanding: On tasks like SWDE, SQuAD, and FDA, Proteus improved robustness and length extrapolation, particularly for models that previously suffered sharp accuracy degradation (Comba, Titans). On the LongBench benchmark, Proteus improved average scores across all three evaluated backbones.

Key Observation: The performance gains generally scale with context length, confirming that the mechanism specifically aids long-context retention and extrapolation.

Significance and Claims

The paper claims that static memory is suboptimal for long-context sequence modeling. By treating memory capacity as a static resource, models fail to balance the need for early compression with the need for later retention.

The significance of Proteus lies in its simplicity and broad applicability:

It is a simple scheduling tool that does not alter the underlying memory architecture, update rules, or internal objectives.
It is architecture-agnostic, improving diverse models from linear RNNs to complex memory-augmented networks.
It demonstrates that scheduling effective capacity is a fundamental lever for improving long-context performance, offering a practical alternative to increasing model size or memory complexity.

The authors remain modest, noting that their schedule is currently fixed and hand-designed, and that the extension to MLP blocks is a proof of concept. They suggest future work could explore learned, data-dependent activation policies and combining incremental activation with growing-memory architectures.