Associative Recurrent Memory Transformer

• ARMT is a neural network architecture that integrates local self-attention, segment-level recurrence, and a quasi-linear associative memory mechanism to process multi-million-token contexts.
• It employs a delta-rule update for fast-weight associative memory, maintaining constant memory usage and enabling efficient retrieval of task-specific information.
• The design uses Diagonal Batching to parallelize segment processing, resulting in significant speedups and state-of-the-art performance on long-context tasks.

The Associative Recurrent Memory Transformer (ARMT) is a neural network architecture designed to efficiently process and reason over extremely long sequences, providing constant-time updates per token and constant memory usage with respect to total context length. ARMT integrates local self-attention, segment-level recurrence, and a lightweight, quasi-linear associative memory mechanism at every layer, enabling the model to encode and retrieve vast quantities of task-specific information over multi-million-token contexts while remaining practical for parallel and hardware-optimized deployment (Rodkin et al., 2024, Sivtsov et al., 5 Jun 2025).

1. ARMT Architecture and Computational Principles

At its core, ARMT extends the segment-level recurrent Transformer (RMT) architecture by introducing, at each Transformer layer $l$, both a compact memory $M^l_s$ and a fast-weight associative memory $A^l_s$ per segment $s$ (Rodkin et al., 2024). ARMT processes sequences in fixed-size segments of length $S$ (commonly $S=512$ tokens), applying standard multi-head self-attention within each segment for local dependency modeling.

Between segments, instead of relying solely on a small set of recurrent memory tokens (as in RMT), each layer maintains a key–value association matrix $A^l_s \in \mathbb{R}^{d_v \times d_k}$ and a normalization vector $z^l_s \in \mathbb{R}^{d_k}$, updated via a delta-rule procedure motivated by fast-weight associative memory research. These per-layer, per-segment fast weights enable constant-time memory storage and associative recall of arbitrary facts from the entire sequence history, without growing memory or read cost.

Segment-level recurrence is implemented by concatenating $M^l_{s-1}$ to the input of the next segment and updating $M^l_s$ during each Transformer block:

$M^l_s$0

The novelty of ARMT is in folding $M^l_s$1 into $M^l_s$2 each segment, and allowing subsequent recall to all tokens via efficient, normalization-corrected associative addressing.

2. Associative Memory Mechanisms

The associative component at each layer operates through a sequence of update and read steps. After normal Transformer processing, each new memory token $M^l_s$3 is projected into key $M^l_s$4, value $M^l_s$5, and scaling factor $M^l_s$6. The feature mapping $M^l_s$7 uses the DPFP-3 nonlinearity. For each $M^l_s$8: $M^l_s$9

$A^l_s$0

At the end of the segment, $A^l_s$1, $A^l_s$2 (for $A^l_s$3 memory updates). To recall for each token $A^l_s$4: $A^l_s$5 with $A^l_s$6 added residually to the token representation. This structure enables finely-graded associative retrieval over the global context at each layer, in constant time per segment.

3. End-to-End Processing Flow

ARMT operates in the following steps for an input sequence of length $A^l_s$7:

1. Split the sequence into segments $A^l_s$8 of length $A^l_s$9, initialize all $s$0.
2. For each segment $s$1 and each layer $s$2:
   • Concatenate token embeddings $s$3 with recurrent memory $s$4.
   • Apply standard Transformer operations to produce $s$5.
   • Perform the associative-memory delta-rule update as described above.
   • Recall from $s$6 for each token, adding $s$7 via a residual connection.
3. After all layers, pass $s$8 to the output or next segment, and $s$9 as state.

This design ensures that every token can, at each layer, retrieve salient information across the entire history, with no buffer growth or computational bottleneck.

4. Theoretical Complexity and Diagonal Batching

Time and memory complexity for ARMT are governed by the segment operations:

• Local self-attention per segment: $S$0
• Associative memory updates and recall: $S$1 and $S$2 per segment With $S$3 and all updates constant-sized, the per-token cost is $S$4 with respect to overall context length.

Sequential segment processing previously limited practical speedup, but the introduction of Diagonal Batching (Sivtsov et al., 5 Jun 2025) enables all operations on the $S$5-th diagonal set $S$6 in the $S$7 grid to execute in parallel. For $S$8 layers and $S$9 segments, this reduces kernel launches to $S=512$0—optimal for the dependency DAG—eliminating the sequential constraint of RMT. This yields large, architecture-neutral speedups on modern hardware without retraining, while memory footprint remains constant (e.g., 0.43 GB at 128K tokens, compared to 70–80 GB for full-attention Transformer).

Method	Time @131k Tokens	Relative Speedup
LLaMA-1B full-attention	8.16 s	baseline
ARMT (sequential)	4.48 s	×1.82
ARMT + Diagonal Batching	2.48 s	×3.29 (vs. baseline)

5. Empirical Results on Long-Context Tasks

ARMT achieves state-of-the-art results on associative retrieval and ultra-long context NLP evaluations:

• In key–value retrieval tasks ("Remember" and "Rewrite"), ARMT outperforms memory RNNs (Mamba) and RMT, sustaining high exact-match for up to $S=512$1 key–value pairs. Gamma-corrected delta rule ($S=512$2) is critical for this property.
• In the BABILong benchmark (multi-task Q&A over contexts up to 50 million tokens): ARMT achieves 79.9% exact-match in QA1 (single supporting fact) at 50M tokens, the first model to reach high accuracy at this scale. On QA1–QA5, ARMT surpasses RMT and Mamba at all tested sequence lengths up to 10M, with generalization far beyond the training regime (train @16K, test up to 50M tokens).
• On language modeling (Wikitext-103), ARMT matches RMT in perplexity but does not yet exploit the full associative memory capacity.

Empirically, Diagonal Batching improves ARMT throughput by 1.8–3.3× on 131K-token sequences and matches batch-8–16 micro-batch throughput without increased memory. Numerical differences vs. sequential ARMT are within 0–2% (logits, task accuracy).

6. Training Strategies, Hyperparameters, and Ablation Studies

The reference implementation extends GPT-2 (137M parameters) to an ARMT with ≈145M parameters. Default hyperparameters include $S=512$3, $S=512$4, $S=512$5 memory tokens per layer, associative dimensions $S=512$6, and nonlinearity $S=512$7 DPFP-3. Training on BABILong uses AdamW, curriculum learning over segment count ($S=512$8), and batch sizes tuned to $S=512$9B tokens (BABILong) or $A^l_s \in \mathbb{R}^{d_v \times d_k}$0K segments (associative retrieval).

Ablation studies demonstrate:

• Gamma-correction ($A^l_s \in \mathbb{R}^{d_v \times d_k}$1) in delta updates prevents catastrophic forgetting in $A^l_s \in \mathbb{R}^{d_v \times d_k}$2.
• Removal of associative memory (Parallel-memory RMT, PRMT) reduces capacity to RMT levels.
• ARMT matches RMT on language modeling; however, nontrivial associative retrieval benefits are specific to tasks exploiting key–value memory over long horizons.

7. Deployment Implications and Practical Integration

ARMT, combined with Diagonal Batching, maintains full task accuracy while providing high GPU utilization, constant memory demand, and architectural determinism for long-sequence inference. No special pipelining or sharding is required; runtime adaptations (grouped forwards, execution on compute diagonals, batch-size $A^l_s \in \mathbb{R}^{d_v \times d_k}$3 number of layers) suffice for efficient hardware scaling. Segment size $A^l_s \in \mathbb{R}^{d_v \times d_k}$4 can be tuned for accuracy, independent of memory/speed concerns due to the decoupled scheduler.

Any ARMT or PRMT codebase can be retrofitted at runtime with Diagonal Batching; the grouping of linear and attention layers with CUTLASS GroupedGEMM and FlashAttention supports models with up to $A^l_s \in \mathbb{R}^{d_v \times d_k}$5 layers achieving $A^l_s \in \mathbb{R}^{d_v \times d_k}$6-fold reduction in kernel-launch overhead. For deep models, this produces strongly superlinear practical speedups for very long-context deployments (Sivtsov et al., 5 Jun 2025).