PaperlayerView exact revision on arXiv

Source-linked AI summary

Proteus: Incremental Memory Activation for Long-Context Sequence Modeling

Reza Bayat, Ali Behrouz, Vahab Mirrokni, Aaron Courville

arXiv:2608.16844v1cs.LGcs.AIcs.CL

TL;DR

Long-context sequence modeling must balance attention’s quadratic cost against recurrent memories’ tendency to overfit early context under static capacity. The paper introduces incremental memory activation and instantiates it as Proteus, which progressively expands active capacity. Across four memory architectures and multiple long-context evaluations, Proteus consistently improves robustness and length extrapolation, especially at the longest contexts, while its activation schedule remains hand-designed.

  • Problem

    Static memory can let early tokens use disproportionate capacity, polluting the state and increasing interference with later context.

  • Method

    Proteus progressively activates memory capacity, restricting reads and writes to an active subset before unlocking fresh capacity over time.

  • Results

    Proteus consistently improves long-context robustness and length extrapolation across SWLA, Comba, Titans, and Hope-Attention, with gains at the longest context lengths.

  • Takeaways & Limitations

    Static capacity is suboptimal, while scheduling capacity is a practical and broadly applicable tool for long-context computation.

  • Takeaways & Limitations

    The activation schedule is fixed and hand-designed, and the paper does not characterize the optimal schedule or its dependence on data or update rule.

Abstract

from arXiv · show

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.

1 Introduction

Transformers provide lossless retrieval but incur quadratic sequence costs, while fixed-size recurrent memories are efficient yet can over-allocate capacity to early tokens. Proteus addresses this imbalance through progressively activated memory capacity and improves long-context performance across models and tasks.

  • Motivation: Softmax attention stores every token, making memory grow linearly and computational cost quadratic with sequence length.This supports retrieval but limits long-context handling.
  • Motivation: Fixed-size recurrent memories reduce decoding and overall computational costs by compressing context into a constant-size hidden state.Their efficiency motivates memory-based alternatives to attention.
  • Static-memory problem: Early tokens can occupy disproportionate memory capacity, causing later tokens to overwrite or interfere with stored information.This biases memory toward initial tokens and weakens incorporation of later context.
  • Approach: Incremental memory activation imposes an early capacity bottleneck and unlocks fresh capacity as context grows.The bottleneck encourages compression, while newly available capacity reduces later interference.
  • Results: Proteus integrates incremental activation into modern recurrent models at no additional cost and improves standard and long-context benchmarks, with larger gains at longer contexts.Evaluations span language modeling, commonsense reasoning, long-context understanding, and needle-in-a-haystack tasks.

2 Preliminaries

Associative-memory models learn key–value mappings online through architecture-specific objectives, optimizers, and parameterizations. Static capacity can favor memorization of early inputs and interfere with later updates, motivating capacity scheduling as an orthogonal intervention.

  • Associative memory: Associative memory maps keys to values through parameterized operators learned by minimizing an internal objective.In sequence models, keys, values, and queries are typically linear projections of inputs.
  • Architecture choices: Different architectures arise from choices of internal objective, optimizer, retention or momentum, and memory parameterization.Examples include linear attention, delta-rule memories, Titans, and TTT.
  • Online optimization: Memory-based sequence models update and read their state online as each token is revealed.The update writes the new input into memory, while the read produces a memory-conditioned output.
  • Effective capacity: Effective capacity is the number of active memory parameters participating in both updates and reads, while total capacity remains fixed.Scheduling the active subset changes which parameters are exposed without changing the objective, optimizer, or architecture.
  • Capacity scheduling: Incremental activation restricts writes and reads to an active subset, then unlocks components that provide fresh degrees of freedom for later inputs.The active subset creates an early compression bottleneck while newly unlocked components reduce interference.
  • Static capacity: Static capacity lets early inputs use all degrees of freedom, favoring memorization and increasing interference when later tokens update the heavily written state.The paper frames this as a failure mode of static capacity for long-context modeling.

3 Incremental Memory Activation

Incremental memory activation schedules a growing active subset of a fixed memory over input position. It combines an early capacity bottleneck, which forces compression, with fresh later capacity that reduces interference.

  • Static capacity creates two failure modes: early inputs are memorized rather than compressed, while later inputs interfere with a saturated state.
  • A capacity-scheduled memory restricts both reads and writes to the active subset while holding inactive components fixed.
  • The activation operator selects memory components at each step, with inactive components locked and unchanged until they are activated.
  • The active subset remains smaller than the full memory for most of the context, forcing compression before progressively unlocked components absorb later inputs.
  • The total memory remains fixed, activation is monotone, and all components become available by the end rather than being discarded.
  • The paradigm applies to capacity-bearing parameters as well as recurrent memory, including MLP-block parameters under the Nested Learning view.

4 Proteus

Proteus realizes incremental memory activation with block-wise gating for memory-based recurrent models and parameter-based architectures. It progressively exposes capacity without additional cost and improves performance across language, reasoning, and long-context evaluations.

  • Proteus partitions memory into blocks and uses a deterministic activation schedule to expose progressively more blocks over the context.
  • Inactive blocks are excluded from retrieval and learning, while multiplicative gating controls effective capacity at no additional cost.
  • Proteus applies updates only to active components, preserves locked blocks, and gates retrieval to the active subspace.
  • Tokens successively access the first 1/E, 2/E, and later fractions of memory until the full memory is available after the final expansion step.
  • Proteus extends the same scheduling principle from recurrent memory states to model parameters, including MLP blocks in Hope-Attention.
  • Across all retrieval datasets and baselines, Proteus improves robustness and length extrapolation at longer contexts, especially for Comba and Titans.

5 Experiments

Experiments show that Proteus consistently improves recurrent backbones across language modeling, reasoning, retrieval, and long-context understanding, with especially strong benefits at longer contexts.

  • Experimental setup and overall results: Proteus improves average downstream accuracy across all four model families and both tested scales, while lowering perplexity in nearly every setting.The evaluated families are Hope-Attention, SWLA, Comba, and Titans.
  • Language modeling and commonsense reasoning: At 760M, Proteus raises every backbone’s average commonsense score, including Hope-Attention from 53.15 to 53.99 and Titans from 52.65 to 53.36.Perplexity also decreases for each backbone at this scale.
  • Language modeling and commonsense reasoning: At 1.3B, Titans+Proteus achieves the table’s best overall performance, with Wiki perplexity 14.94, LMB perplexity 13.03, and average accuracy 58.00.The gains persist at the larger model scale and are strongest for Titans.
  • Needle-in-a-Haystack: On NIAH, Proteus is nearly neutral on saturated short-context tasks but substantially improves harder retrieval variants where the base model degrades.At 16K, Titans improves from 21.4 to 29.8 on S-NIAH-3 and from 69.4 to 74.2 on S-NIAH-2.
  • Long-context understanding: Proteus improves long-context retrieval robustness, length extrapolation, and LongBench average scores across the evaluated backbones without adding parameters or memory.The benefits are reported across retrieval tasks and the six-task LongBench evaluation.

6 Related Work

Related work addresses efficient recurrent computation, adaptive capacity allocation, and associative memory, while Proteus differs by incrementally activating a fixed-size memory rather than growing the memory itself.

  • Modern linear recurrent neural networks: Modern recurrent models reduce Transformer cost by using fixed-size hidden states, yielding constant per-token decoding cost and linear or sub-quadratic overall computation.This line includes linear attention, RetNet, RWKV, S5, and nonlinear recurrent architectures.
  • Modern linear recurrent neural networks: Other recurrent approaches improve memory initialization, update rules, architectures, or effective state growth to address limitations in sequential reasoning and memory use.Growing-memory methods expand the memory itself, unlike Proteus.
  • Adaptive computation: Adaptive computation allocates processing non-uniformly through conditional computation, early exiting, and token routing such as Mixture-of-Depths.These methods reduce cost by avoiding full computation for every input.
  • Non-uniform capacity allocation: Capacity-allocation research links effective-capacity control with generalization and identifies memorization of atypical examples as a high-capacity failure mode.This motivation parallels Proteus’s focus on compressing early context.
  • Associative memory and fast-weight programs: Associative-memory research evolved from classical key-value storage toward dense memories and exponential kernels, providing foundations for modern memory-based architectures.Proteus is grounded in associative-memory and fast-weight principles.

7 Conclusion

The paper presents incremental memory activation as a broadly applicable strategy for long-context modeling and finds that its benefits are largest at the longest contexts, while identifying schedule design as an open limitation.

  • Conclusion: Proteus progressively expands effective capacity over context, restricting reads and writes to active memory blocks before unlocking fresh capacity.The mechanism is applied across SWLA, Comba, Titans, and Hope-Attention.
  • Conclusion: Across four architectures and multiple language-modeling, reasoning, retrieval, and understanding benchmarks, Proteus consistently improves long-context robustness and length extrapolation.The largest gains appear at the longest context lengths.
  • Conclusion: The results suggest that performance benefits arise from allocating capacity over context rather than from any single architecture, making scheduling a practical long-context tool.The conclusion characterizes static capacity as suboptimal within the evaluated scope.
  • Limitations: The activation schedule is fixed and hand-designed, and the study does not characterize the optimal schedule or its dependence on data and memory update rules.The extension to MLP blocks is also demonstrated on only one architecture.
  • Future directions: Future work includes learned data-dependent schedules and combining incremental activation with growing-memory architectures to cover complementary context regimes.The paper also proposes extending activation beyond the training window at inference time.

A Ablations and Analysis

Proteus improves perplexity across token positions, with its advantage widening as capacity is progressively unlocked and remaining positive beyond the training context. An ablation shows that performance depends on the number of partition blocks rather than improving without bound.

  • Effect of the number of blocks: Perplexity improves sharply as the number of partition blocks increases from E = 1 to E = 8.E = 1 exactly recovers the base model, isolating the effect of capacity scheduling.
  • Perplexity by token position: Proteus achieves lower perplexity than the base model at every token position, with the gap largest around the 8K training context length.The gap narrows slowly beyond 8K but remains positive through 32K.
  • Perplexity by token position: The token-position analysis indicates that scheduled capacity helps throughout the context rather than trading early-context quality for late-context gains.The advantage accumulates while fresh capacity is unlocked and persists after the memory becomes fully active.
Loading 2608.16844v1…