Extending LLM Context via Associative Recurrent Memory

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Extending LLM Context via Associative Recurrent Memory

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

artificialintelligence knowledgerepresentation reasoning

artificial intelligence, reasoning, knowledge representation

Extending LLM Context via Associative Recurrent Memory

arXiv:2607.11614v1 [cs.CL] 13 Jul 2026

Gleb Kuzmin1,4,8 Ivan Rodkin2,6 Aydar Bulatov3,6 Yuri Kuratov3,6 Lyudmila Rvanova1 Mikhail Katkov9,10 Ilia Sochenkov7 Misha Tsodyks9,10 Timothy Baldwin2 Mikhail Burtsev5 Artem Shelmanov2 1 FusionBrain Lab 2 MBZUAI 3 Cognitive AI Systems Lab 4 RUDN 5 London Institute for Mathematical Sciences 6 MIRAI 7 Lomonosov Moscow State University 8 Laboratory for Analysis and Controllable Text Generation Technologies RAS 9 School of Natural Sciences, Institute for Advanced Study, Princeton 10 Department of Brain Sciences, Weizmann Institute of Science [email protected] [email protected] 2024). Therefore, since the introduction of the Abstract transformer architecture, long-context processing has emerged as a central and rapidly-evolving research direction (Beltagy et al., 2020; Katharopoulos et al., 2020; Bulatov et al., 2022). Traditionally, efficient long-context approaches have been built using recurrent architectures (Gu and Dao, 2024; Peng et al., 2023); however, such models must typically be trained from scratch, limiting the ability to leverage existing pre-trained LLMs. Moreover, fully-recurrent LMs update the memory at each time step, which complicates high-level information processing in tasks such as structured copying (Jelassi et al., 2024) and instruction following (Park et al., 2024). Recent studies (Bulatov et al., 2024; Rodkin et al., 2024) have explored enhancing transformers with segment-wise context processing and recurrent memory mechanisms. Using human memory as an analogy (Cowan, 2008), full attention within a segment models short-term/working memory, while the module that recurrently propagates crucial information from segment to segment can be viewed as long-term memory. These approaches preserve strong intra-segment modeling performance while enabling linear scaling with respect to context length. In this work, we focus on the Associative Recurrent Memory Transformer (ARMT: Rodkin et al. (2024)), which introduces a capacious segmentlevel associative memory and features strong scaling to extremely long input sizes. Prior work on ARMT-based models has been limited to scales below 200M parameters and evaluated on a narrow set of tasks (Rodkin et al., 2024), leaving their behavior at larger model sizes mostly unexplored. Models at this scale typically struggle to handle complex real-world workloads. In this work, we extend ARMT to small- and medium-sized LMs (1B parameters), which are substantially more capable in practical settings. These models provide a practical

Extending the context length of large language models (LLMs) is critical for many real-world applications, yet standard transformers remain constrained by quadratic compute and linear memory scaling. In this work, we investigate the Associative Recurrent Memory Transformer (ARMT) as a practical approach for enabling long-context processing in LLMs, constant memory scaling, and better efficiency. We make three main contributions. First, we construct two domain-specific long-context datasets designed to evaluate realistic workloads, focusing on narrow-domain fine-tuning scenarios. Second, we propose a comprehensive training recipe for ARMT-based context extension, combining continued pre-training, synthetic long-context data generation, curriculum learning, and selective integration of associative memory into chosen model layers. Third, we present an extensive experimental study demonstrating that ARMT-augmented models: (i) process inputs well beyond their original context limits without degrading performance relative to in-limit baselines; (ii) generalize more effectively to out-of-distribution context lengths; and (iii) need 30% less FLOPs while preserving baseline performance within the original context window.

Introduction

Long-context understanding is crucial for many tasks, such as processing and understanding technical and financial reports, software development, and multi-document reasoning in scientific and legal domains. These scenarios often require models to integrate information distributed across hundreds of thousands or even millions of tokens. However, standard transformer architectures (Vaswani et al., 2017) struggle to scale to such contexts, as the computational and memory costs of self-attention grow quadratically with sequence length. Moreover, transformer performance degrades as the context length increases (Liu et al., 2024; Kuratov et al., 1

middle ground, enabling linear-compute, constantmemory long-context processing while maintaining strong performance on real-world tasks. Our contributions are as follows: • We construct two new datasets over real-world tasks to train and evaluate the long-context performance of LLMs, focusing on narrow domain fine-tuning. • We propose a training recipe for extending the effective context of LLMs using ARMT. This approach incorporates continued pre-training, synthetic generation of training data, curriculum learning, and the strategic integration of the associative memory mechanism into selected LLM layers. • We present an experimental study on extending the context windows of state-of-the-art LLMs using ARMT. Our findings demonstrate that: (1) ARMT-augmented models process inputs well beyond their original context limits without degrading performance relative to in-limit baselines; (2) they exhibit superior generalization to out-of-distribution context lengths compared to base models; and (3) they need 30% less FLOPs while maintaining baseline performance within the original context limits.

some long-context tasks; however, they still face significant limitations compared to transformers in complex algorithmic tasks (Jelassi et al., 2024; Merrill et al., 2024), instruction following (Park et al., 2024), and long-context reasoning (Kuratov et al., 2024). The ability to efficiently address the context is limited by the bottleneck in the fixed-size recurrent state. Hybrid models with interleaving attention and recurrent layers (Lenz et al., 2025; Fu et al., 2025; Qwen Team, 2026) partially address these performance issues, but they still face expensive quadratic compute scaling with context size and require training from scratch. Models with segment-level recurrence present a promising middle ground: they retain linear scaling on longer sequences and maintain the strong short-context performance of transformers with full attention (Dai et al., 2019; Rae et al., 2020). ARMT (Rodkin et al., 2024) has been shown not only to scale linearly with context size, but also to efficiently process tens of millions of tokens on selected tasks. Scaling ARMT opens up possibilities to develop LLMs that are compute-efficient, potentially scalable to extremely long contexts, and maintain all the strengths of full attention within each segment. An overview of the ARMT architecture is presented in Figure 1. The backbone LLM is passed to a wrapper that enables segment-wise processing of input segments, splitting long-context inputs into non-overlapping segments of fixed length. Each segment is also augmented with trainable memory embeddings (“Mem”), which are processed throughout the model. At the core of ARMT lies the associative block, a layerwise memory module that retrieves and updates key-value associations from previous segments and injects the recalled information into the current segment representation. The associative memory mechanism consists of three parts: • Memory extraction: each transformer layer compresses an input segment into memory embeddings. • Memory consolidation: memory embeddings are then consolidated in a per-layer associative matrix as key-value pairs. • Association: every embedding of the following segment is transformed into a query vector and multiplied by the associative matrix. Appendix A presents the formal definition of the associative block. ARMT combines the best of both worlds: the

Related Work and Background

In recent years, the demand for long-context task handling has grown as the problem-solving capabilities of LLMs have evolved (OpenAI, 2023; Reid et al., 2024; Anthropic, 2024). However, even for narrow tasks, both small and large models experience sharp performance degradation when the context size increases. Additionally, scaling the transformer’s context is computationally expensive due to the quadratic computational complexity of attention. This problem has been addressed by sparse attention (Child et al., 2019; Zaheer et al., 2020), attention linearization (Wang et al., 2020; Katharopoulos et al., 2020), and recurrent approaches. Modern RNNs are increasing in popularity due to their linear scaling with length: xLSTM (Beck et al., 2024) builds on traditional recurrent architectures to achieve efficient scaling, RWKV (Peng et al., 2023) and state space models such as Mamba (Gu and Dao, 2024) and Gated DeltaNet (Yang et al., 2024a) utilize simple linear recurrence to achieve faster parallel training. Fullyrecurrent models exhibit strong performance on 2

Base LLM Output Transformer Block

Associative Recurrent Memory Transformer (ARMT) Mem

...

Transformer Block

Output Transformer Block

Associative Block

Intermediate Blocks

Intermediate Blocks Mem

...

Transformer Block

Associative Block

Input Embeddings

Segment 1

Mem

Segment 2

Mem

Figure 1: Base LLM architecture (left) and ARMT architecture (right). ARMT divides the input text into segments and processes them sequentially, allowing the model to handle long contexts.

ability of recurrent architectures to propagate information across arbitrarily long contexts (in principle) and the strong performance of full selfattention within a limited context window. In contrast to Mamba (Gu and Dao, 2024) and RWKV (Peng et al., 2023), ARMT’s computational depth scales with sequence length, enabling more effective multi-step reasoning by combining precise local attention with deep, high-capacity long-term memory (Rodkin et al., 2024; Kuratov et al., 2024).

forms near-randomly on challenging tasks that require propagating salient information across multiple segments, thereby yielding little useful learning signals from such instances. We argue that, much like human memory, machine memory should be trained through a gradual increase in task complexity, i.e. curriculum learning (Bengio et al., 2009). Accordingly, we apply curriculum learning during fine-tuning by progressively increasing task difficulty. We control task difficulty via the maximum context length, defined by the number of ARMT segments. Specifically, the model first learns how to propagate important information across 2 segments, then 4, and finally 8 segments. In addition, we anneal the learning rate across the curriculum stages, gradually reducing it in the later stages. Generating synthetic long context training data. Long context instances might be scarce in the original training dataset; therefore, fine-tuning ARMT on such data might be challenging. Moreover, in order to use curriculum learning, we need relatively large bins of instances that have contexts of specific lengths. To mitigate this problem, we suggest generating synthetic long-context training instances (see Section 4 for further details). We first sample multiple short passages (paragraphs) from a long document and generate a question-answer pair for each passage. These passages are then concatenated into a single long context, which is used as input during training while preserving the original QA supervision. To increase diversity and reduce model-specific biases in the generated data, we produce QA pairs using multiple LLMs from different model families. Associative memory layers pruning. We further investigate whether associative memory is needed in every layer for effective context processing. Similar to modern LLMs that apply sliding-window at-

LLM Context Extension via ARMT

To enable LLMs to handle long contexts, we developed a set of techniques essential for effective ARMT training: continued pre-training, synthetic long context supervision, curriculum learning, and pruning of ARMT layers. Continued pre-training. When associative memory is introduced into a pre-trained LLM, the newly added parameters remain uninitialized. As a result, although the backbone LLM is already pre-trained, ARMT still requires additional adaptation to learn how to use the memory mechanism effectively. We hypothesize that continued unsupervised languagemodel pre-training can properly initialize the associative memory parameters and simplify subsequent task-specific fine-tuning. This stage should be performed on sufficiently long contexts so that the memory mechanism is actively used to propagate salient information across multiple segments. In our experiments, we used 8 segments. To make such continued pre-training feasible, we developed an optimized ARMT implementation compatible with DeepSpeed ZeRO Stage 3 (Rasley et al., 2020; Rajbhandari et al., 2020). Curriculum learning. When associative memory layers are trained from scratch, ARMT initially per3

tention only in a subset of layers (Team et al., 2025), we hypothesize that associative memory can likewise be retained only in selected layers. Therefore, as an optional efficiency step, we suggest removing redundant associative layers after training. We also propose a recipe for the pre-selection of layers with associative memory, showing that one can train the ARMT model with only a few associative layers and achieve performance comparable to that of the full ARMT model.

keeping the validation and test sets unchanged. For both datasets, we place the question before the context to reduce reliance on parametric knowledge. Examples of MT and GR datasets are presented in Tables 6 and 7 in Appendix B.1. Dataset statistics are provided in Table 5 in Appendix B.1.

Experiments

5.1

Experimental Setup

The main objective of our experiments is to demonstrate that ARMT, together with the proposed training techniques, can: (1) effectively extend the context capacity of the base LLM; and (2) replace full quadratic self-attention with a piece-wise attention mechanism that remains quadratic only within a small local window, while preserving the performance of the original LLM through the recurrent memory mechanism. To this end, we conducted an experiment in which we compared the base LLM and an ARMT-augmented LLM in which the native context window is drastically reduced relative to the base model (to 1,024 tokens), and the primary burden of long-context processing is placed on the recurrent memory mechanism instead of full selfattention. This means, e.g., that if we process a 32k size context, we need to split it into 32 segments, in each of which the full-attention only spans 1k tokens, and important information is propagated through memory. Base LLMs. We experimented with Gemma-31B-IT (Team et al., 2025) and SmolLM-2-360MIT (Allal et al., 2025) as base LLMs. Gemma-3-1B-IT was selected because it is a strong state-of-the-art model in the 1B-parameter category, while still having a relatively limited maximum context length of 32k tokens. Moreover, it exhibits substantial performance degradation in long-context settings, even within this nominal context window (Team et al., 2025). We hypothesize that this arises from the use of interleaving sliding-window attention in its architecture instead of monolithic full attention. Therefore, this model is a good candidate for context extension. SmolLM2-360M-IT is also well suited for our experiments: it is small enough to train under a limited GPU memory budget, yet remains highly capable for its size. Its limited 8k-token context window makes it a natural candidate for context extension. Evaluation and metrics. We fine-tuned the models on the training sets and evaluated their performance on the test sets of MT and GR. We used

Long-Context Datasets

Approaches to LLM context extension rely either on continued LM pre-training (Gao et al., 2025) or on long-context supervised fine-tuning (SFT) (Xu et al., 2024). Continued LM training is prohibitively computationally expensive, requiring hundreds of GPU hours and is limited in effectiveness. Existing long-context SFT dataset selection is scarce and mostly targets selected domains. In our experiments, we aim to extend the context in a compute-efficient way by carefully crafting longcontext datasets with sufficient diversity, complexity, and size, as well as controllable sample length suitable for curriculum learning. We construct training and evaluation datasets using the ManyTypes4Py (Mir et al., 2021) and GovReport-QS (Cao and Wang, 2022) datasets. We focus on these datasets to demonstrate the possibility of context extension with the ARMT model on domain-specific data, even with limited compute resources, while avoiding more compute-demanding instruction following datasets. ManyTypes-long (MT) targets variable type prediction in code with long context. We split the original dataset into non-overlapping repositories for training, validation, and test splits, and stacked the code scripts from each repository to obtain a long text with the desired length of up to 64k tokens. We reused the original labels for the variable types from ManyTypes4Py (Mir et al., 2021). GovReport-long (GR) focuses on long-document question answering. The original GovReport-QS consists of triplets {report, question, answer} with ground-truth paragraphs in the report for each question. We used ground-truth paragraphs from the report and mixed them with non-relevant paragraphs, keeping the paragraph order to build datasets of up to 64k in length. Due to the limited size of the original dataset, we augment the training split with synthetic examples to create GR-100+, while 4

GPU RAM Usage, MB

50000

GPU memory usage on MT dataset Base, BS=8 ARMT (ours), BS=8 ARMT (ours), BS=32

Max. train length

Inference time on MT dataset Max. base model context

40000 30000 20000 10000

0.200

Inference time per token, seconds

60000

0.175

Base, BS=8 ARMT (ours), BS=8 ARMT (ours), BS=32

0.150

Max. train length

0.125

Max. base model context

0.100 0.075 0.050 0.025 0.000

0 103

104

103

Context length, log scale

(a) GPU memory usage.

104

Context length, log scale

(b) Inference time.

Figure 2: ARMT performance in terms of inference time and memory usage. (1) Gemma-3-1B-IT (Base) with batch size of 8 uses approximately 40 GB of GPU RAM on 32k length, while ARMT model uses only 12 GB GPU RAM (left) while having approximately the same inference time (right). (2) ARMT shows constant memory usage with respect to context length, improving base model performance and allowing for long-context data processing.

the exact match (EM) metric for the MT dataset, as the target answers are short and require precise variable type prediction. For the GR and GR-100+ datasets, we use the ROUGE-L (Lin, 2004) metric, as the ground-truth answers in GR are free-form and relatively long. Training details. We fine-tune the LLMs using LoRA (Hu et al., 2021) for GPU memory efficiency. We used the training setup from Section 3 for both the GR and MT datasets. Continued pre-training of the Gemma-3 ARMT model was conducted on 19B tokens from FineWeb-Edu (Penedo et al., 2024). We also experimented with additional ARMT pretraining on a synthetic long-context QA task. Training hyperparameters are provided in Appendix E. 5.2

ken throughput, since GPU cores may already be saturated. However, because self-attention in ARMT is quadratic only within short, fixed-size segments, ARMT also requires substantially fewer self-attention FLOPs than full-attention-based models. For a sequence of length T and segment size S, ARMT reduces global-attention FLOPs by approximately T /S; for T =32,768 and S=1,024, this corresponds to a 32× reduction in global-attention FLOPs. When feed-forward and projection layers are taken into account, ARMT theoretically requires roughly one-third fewer FLOPs in total than Gemma under our approximation; see Section A.1 for the full derivation. Figure 2(b) empirically validates this reduction by illustrating the average token inference time. For a batch size of 8, ARMT is slower than the base model due to the sequential processing of segments and the under-utilization of GPU cores. However, when we increase the batch size to 32, ARMT substantially outperforms the base Gemma model. Optimized ARMT implementations could further improve efficiency (Sivtsov et al., 2025).

Analysis of Computational Efficiency

GPU memory consumption is the main advantage of ARMT, as it remains constant while the sequence length increases, whereas even the most efficient implementations of full self-attention scale at least linearly. This is illustrated in Figure 2(a): for Gemma, memory consumption grows as the context length increases and can eventually exceed memory limits. In contrast, ARMT does not require additional memory to process longer sequences, so it can use larger batch sizes. For example, under the same memory budget required by Gemma to process a 65k-token sequence with a batch size of 8, ARMT can process sequences of the same length using a batch size of 32. FLOPs and token throughput. Scaling the batch size does not automatically increase to-

5.3

Analysis of Model Performance

Extending the context window of base LLM. The main results with Gemma-3 on the MT dataset are shown in Figure 3a and Table 1 (see also Table 8 in Appendix C). The base model without fine-tuning achieves near-zero performance across all context lengths, highlighting the necessity of task-specific adaptation. After fine-tuning on 8k contexts (Base, MT, 8k), the model performs well 5

Main results on MT dataset

Main results on GR-100+ dataset 0.40

0.9

Max. train length

0.35

0.8

Max. base model context

Base Mamba ARMT (ours) ARMT, pre-selected layers, trained (ours) ARMT, pretrained (ours) Base, mean Mamba, mean ARMT (ours), mean ARMT, pre-selected layers, trained (ours), mean ARMT, pretrained (ours), mean

0.6 0.5 0.4 0.3

0.25

ROUGE-L

Exact match (EM)

0.30 0.7

0.05 0.00

104

103

0.15 0.10

Max. base model context

Max. train length

0.20

104

103

Context length

(a) MT results.

(b) GR-100+ results.

Figure 3: Main results. (1) Gemma-3-1B-IT (Base) with full attention and Mamba-2 trained with a context length of 8,192 show a sharp performance drop beyond 32k tokens. (2) ARMT model with Gemma-3-1B-IT backbone (ARMT) maintains stable performance across all evaluated context lengths and outperforms other models in the long out-of-distribution (Long-OOD) regimen (>32k). (3) ARMT with continued pre-training followed by fine-tuning consistently outperforms the variant trained with fine-tuning only. (4) ARMT fine-tuned with only five pre-selected associative layers out of 26 achieves comparable or better performance than the full ARMT model.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, MT, 8k

ARMT, MT, 8k

ARMT, MT, 8k pre-selected layers, trained

ARMT, MT, 8k, pretrain

0.000 0.000 0.000 0.000 0.000 0.000 0.011 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.001 0.000 0.001

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.705 0.717 0.712 0.837 0.840 0.857 0.777 0.685 0.812 0.762 0.800 0.802 0.744 0.793 0.801 0.776

0.756 0.743 0.736 0.740 0.859 0.827 0.868 0.766 0.697 0.757 0.812 0.791 0.762 0.767 0.783 0.777 0.777

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Table 1: Best results on the MT dataset for Gemma-31B-IT model, metric = EM. ARMT performs comparably with the base model overall, and outperforms it on OOD and Long-OOD.

Base, No FineTuning

Base, GR-100+, 8k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k pre-selected layers, trained

ARMT, GR-100+, 8k, pretrain

0.188 0.161 0.140 0.114 0.113 0.106 0.101 0.091 0.088 0.083 0.078 0.078 0.103

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.363 0.317 0.304 0.287 0.250 0.231 0.231 0.221 0.243 0.259 0.204 0.205 0.246

0.370 0.316 0.293 0.303 0.260 0.245 0.260 0.252 0.262 0.260 0.221 0.207 0.300

0.143

0.351

0.299

0.304

0.308

0.092

0.266

0.238

0.234

0.252

0.084

0.175

0.215

0.214

0.228

0.116

0.306

0.266

0.267

0.278

Table 2: Best results on the GR-100+ dataset for Gemma-3-1B-IT model, metric = ROUGE-L. ARMT performs comparably to the base model overall, and outperforms it on OOD and Long-OOD.

up to its maximum supported length of 32k tokens, but degrades sharply beyond this limit due to context truncation. In contrast, ARMT shows no degradation beyond 32k tokens. Notably, after curriculum training with only 4k contexts, ARMT already generalizes to inputs as long as 65k, demonstrating the effectiveness of the recurrent memory mechanism. The final ARMT model (ARMT, MT, 8k) matches the fine-tuned base model on in-domain inputs (less than 8k) while substantially improving performance on out-of-distribution (OOD) context lengths (larger than 8k). The results with Gemma-3 on GovReport-long-

100+ are presented in Figure 3b, Table 2 and in Table 9 in Appendix C. They exhibit a similar trend. The base model achieves low zero-shot in-domain performance (ROUGE-L 0.116) and improves to 0.306 after fine-tuning on GR-100+. Its performance drops substantially on Long-OOD inputs (0.175). ARMT partially mitigates this degradation, improving the Long-OOD performance to 0.215, with only a minor trade-off on shorter contexts. Experiments with SmolLM show a similar pattern (Tables 23 and 24 in Appendix D.6). While the fine-tuned base model degrades beyond its native context length (8k), ARMT maintains competitive 6

performance up to 32k tokens. Overall, these results highlight two main findings: (1) ARMT, together with the proposed training techniques, effectively extends the ability of the base model to handle longer inputs; and (2) for contexts that fall within the original LLM’s native context window, our piecewise attention mechanism with recurrent memory largely preserves the performance of the original full-attention model. Combined with constant memory consumption and the estimated 30% reduction in FLOPs derived in Section 5.2, these results demonstrate the strong practical applicability of ARMT. Comparison to other baselines. We fine-tuned a 1.3B-parameter Mamba-2 model (Dao and Gu, 2024) on the MT and GR-100+ datasets (Figure 3). Mamba-2 is a variant of a state space model that shows comparable performance with transformers and linear scaling in sequence length, making it feasible for long-context tasks. Although Mamba-2 handles long contexts (32–64k) better than the base full-attention LLM (on MT), it still substantially underperforms ARMT in this setting. However, in shorter contexts (on MT), Mamba-2 performs better than ARMT and the base full-attention model, which can be attributed to the larger model size. Overall, ARMT shows better length generalization than other baselines. Additional results for Mamba (Gu and Dao, 2024), Mamba-2 (Dao and Gu, 2024), DeltaNet (Yang et al., 2024b), and xLSTM (Beck et al., 2024) are reported in Table 40 (Appendix J). We also compared the ARMT model with the simpler RMT model (Bulatov et al., 2022), which does not use associative memory layers (Table 36 in the Appendix G). Results show that the associative memory mechanism introduced in the advanced ARMT architecture substantially improves performance and generalization compared to RMT. Other long-context benchmarks. We further trained and evaluated ARMT on BABILong (Kuratov et al., 2024) and ContractNLI from SCROLLS (Shaham et al., 2022) to validate the training recipe and model. The results are reported in Tables 38 and 39 (Appendix I) and Table 37 (Appendix H). On both benchmarks, ARMT with the proposed training procedure outperforms the base full-attention model. Ablation: synthetic data. To study the impact of synthetic data, we train ARMT on GR variants with different proportions of synthetic samples (Appendix D.4, Figure 5). Adding synthetic longcontext data consistently improves performance up

Model/ Lengths

Base, MT, 8k

ARMT, MT, 8k

ARMT, MT, 8k only top-4 layers

ARMT, MT, 8k only top-4 layers, trained

ARMT, MT, 8k pre-selected layers, trained

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713

0.756 0.724 0.708 0.712 0.793 0.765 0.758 0.755 0.663 0.729 0.693 0.765 0.713

0.756 0.714 0.736 0.699 0.815 0.827 0.802 0.777 0.708 0.757 0.762 0.765 0.752

0.756 0.705 0.717 0.712 0.837 0.840 0.857 0.777 0.685 0.812 0.762 0.800 0.802

0.797

0.749

0.738

0.744

0.709

0.741

0.730

0.767

0.793

0.463

0.741

0.759

0.801

0.741

0.744

0.733

0.759

0.776

Table 3: Associative layers ablation on the MT dataset for Gemma-3-1B-IT model, metric = EM. ARMT with only top-4 associative blocks retains almost the same performance as the full ARMT without training. The model with 5 associative blocks (approximately 20%) achieves the same performance as the full ARMT.

Num. of Segments 2 4

Model ARMT ARMT, pre-selected layers ARMT ARMT, pre-selected layers

Num. of Trainable Parameters

GPU RAM Usage ↓

Steps per Second ↑

90.5M 59.6M 90.5M 59.6M

45.4GB 44.5GB 71.8GB 69.8GB

2.86 3.64 1.62 2.14

Table 4: ARMT with pre-selected layers saves 30% of trainable parameters and speed-ups training up to 30%, while reducing GPU memory usage. Results for ARMT with Gemma-3-1B-IT backbone; only five associative layers out of 26 are used in the pre-selected setup.

to a moderate scale, after which the gains saturate. The best results are achieved with GR-100+, corresponding to a synthetic-to-real ratio of 5.5, which we use in subsequent experiments. Ablation: continued pre-training. We analyze the effect of continued pre-training on ARMT prior to task-specific fine-tuning in Figure 3 and in Tables 10 and 11 in Appendix D.2. Starting from Gemma-3-1B-IT, we add an uninitialized associative memory with 1024-token segments and continue pre-training on 19B FineWeb-Edu tokens (Penedo et al., 2024) using 8-segment sequences. The loss curves of the pre-trained ARMT model are presented in Figure 4 in Appendix D.1. After pre-training, the model is fine-tuned on the GR-100+ dataset using curriculum learning. Continued pre-training consistently improves performance for both in-domain and out-of-domain context lengths. It effectively initializes the associative memory, enabling the model to propagate 7

information across segments and making this capability easier to exploit during fine-tuning. Ablation: curriculum learning. To further examine the role of curriculum learning, we also fine-tune the pre-trained ARMT model without curriculum scheduling (ARMT w/o CL). As shown in Table 10 in Appendix D.2, this simpler training procedure achieves performance comparable to the curriculum-based variant, indicating that continued pre-training alone is sufficient to effectively initialize associative memory. However, according to our results, for models without continued pre-training, curriculum learning is necessary. Associative layer pruning and selection. Adding associative memory layers to a pre-trained LLM increases the number of trainable parameters and training cost, and adds the need for re-adaptation as inserted associative memory layers alter the internal representations of LLM. This motivates the following RQ: what is the minimal architectural transformation needed for effective associative memory integration, and which layers of the LLM are most important for introducing associative memory? To investigate this, we conduct two sets of experiments: (1) after fine-tuning, we replace a subset of ARMT layers with an identity function; (2) we introduce associative memory only into a subset of LLM layers before fine-tuning. The results of the first (pruning) experiment are provided in Tables 12 and 13 in Appendix D.3 for the MT and GR-100+ datasets. For both datasets, the most important layers are concentrated in the third quarter of the network (layers 13–18), while the early layers (0–6) contribute the least and can be removed after training without performance loss. Fine-grained ablations further show that only a few layers are critical, with layer 14 consistently being the most important. Retaining only the top-1 or top-4 layers preserves performance, reducing the number of associative layers by up to six times. Detailed layer-wise ablation results are provided in Tables 15 and 16 in Appendix D.3. In the second experiment, we investigate whether we can add the associative memory only to a subset of LLM layers before fine-tuning. We first train ARMT models on GR-100+ and MT using only the 4 most important layers according to the previous analysis (see Table 14 in Appendix D.3 and Table 3). On GR-100+, the model with the most important ARMT layers achieves slightly lower overall performance than the full ARMT model but performs better on the Long-OOD subset. On

the MT dataset, such a model outperforms the full ARMT model. The drawback is that identifying important layers requires training a full ARMT model first, which is not convenient. To address this limitation, we propose a universal pre-selection strategy: two middle layers, the final layer, and one layer in each of the first and last quarters. For ARMT with Gemma-3-1B-IT (26 layers), this corresponds to layers 7, 13, 14, 19, and 25. This configuration outperforms the full ARMT model on both MT and GR-100+ (Figure 3, Table 14 in Appendix D.3 and Table 3), while eliminating the need for prior layer importance analysis. It also reduces the number of trainable parameters and speeds up fine-tuning by 30% (see Table 4). Other ablations. We also evaluated the supervised pre-training of ARMT on synthetic QA data prior to training on GR-100+ (Table 22 in Appendix D.5). Although this improves in-domain performance, it degrades long-context generalization, likely due to limited sequence lengths in the synthetic data. In contrast, continued language model pre-training is more effective for initializing associative memory.

Conclusion

We present a practical recipe for extending the context length of LLMs on domain-specific tasks using ARMT. The proposed approach combines several techniques: (i) continued pre-training, (ii) synthetic long-context data generation, (iii) curriculum learning, and (iv) optional pruning of associative layers. Together, these components enable small local models to achieve strong long-context understanding, helping to close an important gap for privacypreserving applications that cannot rely on remote API-based LLMs. ARMT also enables constant GPU memory utilization for arbitrarily long contexts and a 30% reduction of TFLOPs. We introduced two datasets for the training and testing of LLMs in long-context tasks: ManyTypes-long and GovReport-long. Finally, we show that only a few of the most important layers need to be augmented with ARMT. Based on this observation, we propose a universal strategy for pre-selecting layers for augmentation prior to fine-tuning. Our experiments demonstrate that fine-tuning ARMT with only 20% of the associative layers is sufficient to match the performance of the full ARMT model while reducing the training time by approximately 30%. 8

Limitations

References

We experimented only with relatively small LLMs (up to 1B parameters) due to the chosen scope of the paper and computational constraints. However, we argue that this setting is still important, especially for practical domain-specific deployments. Although we trained and evaluated ARMT on multiple long-context benchmarks that covered code, question answering, and document understanding, our study still covers only a subset of possible long-context tasks and modalities. Future work should evaluate the proposed approach across a broader range of tasks and domains. Finally, our analysis of associative memory focuses primarily on layer-level ablations within the ARMT architecture. Although we identify a subset of associative layers that contribute most to long-context performance, the underlying mechanisms of how associative memory interacts with transformer representations remain only partially understood. We used AI assistants in a limited and controlled manner for two purposes: (i) generating synthetic training data for fine-tuning of the models and (ii) improving the grammar and clarity of the manuscript. The AI tools were used according to their intended purpose and with careful consideration of responsible and ethical research practices.

Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydlícek, Agustín Piqueres Lajarín, Vaibhav Srivastav, Joshua Lochner, Caleb Fahlgren, Xuan-Son Nguyen, Clémentine Fourrier, Ben Burtenshaw, Hugo Larcher, Haojun Zhao, Cyril Zakka, Mathieu Morlon, and 3 others. 2025. SmolLM2: When Smol goes big - datacentric training of a small language model. arXiv preprint arXiv:2502.02737. Anthropic. 2024. Introducing the next generation of Claude. urlhttps://www.anthropic.com/news/claude3-family. Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, and 1 others. 2024. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd annual meeting of the association for computational linguistics (volume 1: Long papers), pages 3119– 3137. Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. 2024. xLSTM: extended long shortterm memory. In Proceedings of the 38th International Conference on Neural Information Processing Systems, pages 107547–107603. Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. Preprint, arXiv:2004.05150.

Ethical Considerations

Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. 2009. Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning - ICML ’09, pages 1–8, Montreal, Quebec, Canada. ACM Press.

This work focuses on improving long-context capabilities of small LLMs intended for local deployment in privacy-sensitive environments. By enabling efficient processing of long documents without reliance on large external services, the proposed approach can help reduce the risks of data leakage when handling confidential or proprietary information. The datasets used in this work are derived from previously published sources. We do not introduce new personal or sensitive data in presented datasets. As with other LLM technologies, improved longcontext processing can enable more effective analysis of large collections of text, which could potentially be misused for large-scale information extraction or surveillance. We therefore encourage careful evaluation of downstream deployments, particularly in settings involving sensitive data, and recommend transparency about dataset construction and model limitations when releasing derived resources.

Aydar Bulatov, Yuri Kuratov, Yermek Kapushev, and Mikhail Burtsev. 2024. Beyond attention: Breaking the limits of transformer context length with recurrent memory. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17700– 17708. Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. 2022. Recurrent memory transformer. In Advances in Neural Information Processing Systems, volume 35, pages 11079–11091. Curran Associates, Inc. Shuyang Cao and Lu Wang. 2022. HIBRIDS: Attention with hierarchical biases for structure-aware long document summarization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 786–807, Dublin, Ireland. Association for Computational Linguistics.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509.

Yuta Koreeda and Christopher Manning. 2021. ContractNLI: A dataset for document-level natural language inference for contracts. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 1907–1919, Punta Cana, Dominican Republic. Association for Computational Linguistics.

Nelson Cowan. 2008. What are the differences between long-term, short-term, and working memory? Progress in Brain Research, 169:323–338.

Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. 2024. BABILong: Testing the limits of LLMs with long context reasoning-in-a-haystack. Advances in Neural Information Processing Systems, 37:106519–106554.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. 2019. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978–2988.

Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, and 1 others. 2025. Jamba: Hybrid Transformer-Mamba language models. In International Conference on Learning Representations, volume 2025, pages 67959–67984.

Tri Dao and Albert Gu. 2024. Transformers are SSMs: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning, pages 10041–10071.

Chin-Yew Lin. 2004. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pages 74–81, Barcelona, Spain. Association for Computational Linguistics.

Yonggan Fu, Xin Dong, Shizhe Diao, Hanrong Ye, Wonmin Byeon, Yashaswi Karnati, Lucas Liebenwein, Maksim Khadkevich, Alexander Keller, Jan Kautz, and 1 others. 2025. Nemotron-Flash: Towards latency-optimal hybrid small language models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems.

Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157–173.

Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. 2025. How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7376– 7399.

William Merrill, Jackson Petty, and Ashish Sabharwal. 2024. The illusion of state in state-space models. In International Conference on Machine Learning, pages 35492–35506. PMLR.

Albert Gu and Tri Dao. 2024. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling.

A. M. Mir, E. Latoskinas, and G. Gousios. 2021. ManyTypes4Py: A benchmark Python dataset for machine learning-based type inference. In IEEE/ACM 18th International Conference on Mining Software Repositories (MSR), pages 585–589. IEEE Computer Society.

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. 2024. RULER: What’s the real context size of your long-context language models? In First Conference on Language Modeling.

OpenAI. 2023. New models and developer products announced at DevDay. https://openai.com/index/ new-models-and-developer-products-announced-at-devday/. Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. 2024. Can Mamba learn how to learn? a comparative study on in-context learning tasks. In Forty-first International Conference on Machine Learning.

Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, and Weizhu Chen. 2021. Lora: Low-rank adaptation of large language models. CoRR, abs/2106.09685. Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. 2024. Repeat after me: transformers are better than state space models at copying. In Proceedings of the 41st International Conference on Machine Learning, pages 21502–21521.

Guilherme Penedo, Hynek Kydlíček, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. 2024. The FineWeb datasets: decanting the web for the finest text data at scale. In Proceedings of the 38th International Conference on Neural Information Processing Systems, NIPS ’24, Red Hook, NY, USA. Curran Associates Inc.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156–5165.

Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi

Cao, Xin Cheng, Michael Chung, Leon Derczynski, Xingjian Du, Matteo Grella, Kranthi Gv, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartłomiej Koptyra, and 13 others. 2023. RWKV: Reinventing RNNs for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 14048– 14077, Singapore. Association for Computational Linguistics.

Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, and 1 others. 2025. Gemma 3 technical report. arXiv preprint arXiv:2503.19786. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in Neural Information Processing Systems, 30.

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. 2024. Yarn: Efficient context window extension of large language models. In International Conference on Learning Representations, volume 2024, pages 31932–31951.

Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. Preprint, arXiv:2006.04768. Peng Xu, Wei Ping, Xianchao Wu, Chejian Xu, Zihan Liu, Mohammad Shoeybi, and Bryan Catanzaro. 2024. ChatQA 2: Bridging the gap to proprietary LLMs in long context and RAG capabilities. In The Thirteenth International Conference on Learning Representations.

Qwen Team. 2026. Qwen3.5: Accelerating productivity with native multimodal agents. https://qwen.ai/ blog?id=qwen3.5. Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. 2020. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations.

Songlin Yang, Jan Kautz, and Ali Hatamizadeh. 2024a. Gated delta networks: Improving Mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1– 16. IEEE.

Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. 2024b. Parallelizing linear transformers with the delta rule over sequence length. Advances in Neural Information Processing Systems, 37:115491–115522.

Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. 2020. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 3505–3506.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. 2020. Big Bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems, volume 33, pages 17283–17297. Curran Associates, Inc.

Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, and 1 others. 2024. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530. Ivan Rodkin, Yuri Kuratov, Aydar Bulatov, and Mikhail Burtsev. 2024. Associative recurrent memory transformer. arXiv preprint arXiv:2407.04841. Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. 2022. SCROLLS: Standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 12007–12021, Abu Dhabi, United Arab Emirates. Association for Computational Linguistics. Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, and Ivan Oseledets. 2025. Diagonal batching unlocks parallelism in recurrent memory transformers for long contexts. arXiv preprint arXiv:2506.05229.

ARMT. ARMT splits the sequence into TS segments. Each segment is processed with attention context length S + M , where M is the number of memory tokens. Therefore, the global-attention FLOP cost becomes

Formal Description of the ARMT Architecture

For the hidden states of segment s in layer l Hsl , memory tokens from the previous layer Msl−1 , the associative matrix Als , and normalization zsl are updated as follows: Msl = {mi },

FARMT_Attn ∼

Msl = TransLayer(Hsl−1 , Msl−1 ) (1)

When M ≪ S, this simplifies to

ki , vi = WK mi , WV mi ; βi = σ(Wβ mi ); (2) Al0 = ⃗0; z0l = ⃗0; (3) vi =

Als−1 ϕ(ki ) ; (zs−1 )T ϕ(ki )

Als = Als−1 +

γi = 1 −

FARMT_Attn ≈ 4Ng Hdh ST.

(zs−1 )T ϕ(ki ) ; ∥ϕ(ki )∥2 (4)

βi (vi − v i ) ⊗ ϕ(ki );

(5)

(6)

Thus, for global attention alone, the reduction factor is approximately FGemma_Attn T ≈ . FARMT_Attn S

i l zsl = zs−1 +

γi ϕ(ki ).

T 4Ng Hdh (S + M )2 . S

For T =32,768 and S=1,024, this gives

T = 32, S

Reading from memory in the following segments l−1 for embedding xj from Hs+1 : qj = WQ xj ;

yj =

Als ϕ(qj ) , (zsl )T ϕ(qj )

i.e., ARMT achieves a 32× reduction in globalattention FLOPs when memory-token overhead is neglected.

(7)

where yj is an association for xj . A.1

Overall FLOP ratio. To estimate the total FLOP ratio, we combine the global-attention term with the dense-layer term. Ignoring local attention and memory-token overhead, we have

FLOP Analysis

We provide the FLOP calculation used in Section 5.2. We compare Gemma-3 and ARMT on a sequence of length T , with segment size S, M memory tokens, N layers, Ng global-attention layers, Nl local-attention layers, hidden dimension d, head dimension dh , and H attention heads.

FGemma ≈ 4Ng Hdh T 2 + 30N T d2 , and

Gemma-3. For a full sequence of length T , the attention cost per layer is 4HT 2 dh for global attention and 4HT W dh for local attention with window size W . For long sequences, the sequence-lengthdependent attention cost is dominated by global attention, giving

FARMT ≈ 4Ng Hdh ST + 30N T d2 . Dividing both numerator and denominator by 2T gives 2Ng Hdh S + 15N d2 FARMT ≈ . FGemma 2Ng Hdh T + 15N d2

FGemma_Attn ∼ 4Ng Hdh T 2 . Feed-forward and projection layers scale linearly with the sequence length and quadratically with the hidden dimension. Using the same coarse accounting as in the main text, the dense-layer cost is

For Gemma-3-1B-IT, substituting the model constants into this expression gives an overall FLOP ratio of approximately 0.67. Thus, although ARMT reduces global-attention FLOPs by 32× for T =32,768 and S=1,024, the total FLOP reduction is more modest because feed-forward and projection layers remain unchanged. Under this approximation, ARMT uses roughly one third fewer total FLOPs than Gemma.

FFFN = (2 × 2 × 6 + 2 × 3)N T d2 = 30N T d2 . This term is independent of the segmentation and is therefore approximately equal for Gemma and ARMT when the memory-token overhead is neglected. 12

Dataset/ Split

Train

Validation

Test

MT, 2k MT, 4k MT, 8k GR, 2k GR, 4k GR, 8k GR-100+, 2k GR-100+, 4k GR-100+, 8k ContractNLI (SCROLLS)

229.4k 289.2k 328.3k 9.9k 17.4k 19.7k 90.5k 121.7k 128.4k 7.2k

3.1k 4.5k 5.1k 0.5k 0.9k 1.0k 0.5k 0.9k 1.0k 1.0k

1.3k 1.3k 1.3k 1.0k 1.0k 1.0k 1.0k 1.0k 1.0k 2.1k

Table 5: Dataset statistics - number of samples for each dataset and for each split.

Datasets Statistics and Ablations

B.1

Dataset Statistics and Examples

Figure 4: ARMT-augmented Gemma-3-1B-IT pretraining on 19B tokens from the FineWeb-Edu dataset. Pre-training was performed with 8,192-token sequences, divided into 8 segments of 1024 tokens each. We additionally report the final checkpoint loss on the train subset, averaged over the first 1, 2, 4, and 8 segments.

In this section, we present statistics for all dataset versions and splits used in our experiments. The dataset overview is summarized in Table 5. We also provide short illustrative examples from the MT and GR datasets in Tables 6 and 7.

D.2

After ARMT pre-training, we also conducted finetuning using the same setup as for the ARMT model without continued pre-training on GR-100+ and MT datasets. The results are shown in Tables 10 and 11. For both datasets, pre-training substantially improves performance on downstream tasks, especially for long-context samples from 8k to 65k in length.

ARMT Training Dynamics

Tables 8 and 9 presents results obtained using curriculum learning for the ARMT model on the MT and GR-100+ datasets. The performance of the ARMT models gradually increases during curriculum learning.

Additional Experimental Results

D.1

ARMT language modeling pre-training

ARMT Fine-Tuning After Pre-Training

D.3

Ablation Study for Associative Memory Layers

We conducted a layer-wise ablation for associative layers in a trained ARMT model for the MT and GR-100+ datasets. The results are presented in Tables 15 to 18. For the GR-100+ dataset, the most important representations are from the 14th, 25th, 9th, and 13th layers, while other layers contribute substantially less. For the MT dataset, the most important representations are from the 14th, 18th, 25th, and 19th layers. Layers are numbered from the starting embeddings layer (i. e., the 0th layer goes after the embeddings layer, while the 25th layer goes right before the last layer). We suppose that the middle layers are the most important for the associative memory, as these layers contain the more abstract representation than the lower or higher layers. On the lower layers, these representations are still not formed, while on the higher ones, they are close to the target tokens. We also compared the training dynamics during curriculum learning with top-4 associative blocks

For a fairer comparison between ARMT fine-tuning and the fine-tuning of other pre-trained models, we also conducted a small-scale pre-training of Gemma-3-1B-IT augmented with ARMT. The model was trained on 19B tokens from the FineWeb-Edu dataset, using concatenated text samples of length 8,192 tokens, which were split into 8 segments of 1,024 tokens each. Both ARMT and Gemma parameters were fully trained. We used a batch size of 0.5M tokens with a learning rate of 1e-5 and a linear warmup scheduler with 5k warmup steps. In Figure 4, we show the ARMT pre-training convergence compared to the loss values computed from 1, 2, 4, and 8 segments using the final checkpoint. The pre-training took approximately 100 hours on an 8×H100-80GB GPU cluster. 13

0 GR

0.35

Number of synthetic train samples 50000 GR-100

100000 GR-100+

150000

learning setup. Due to limited resources, we pretrain the model only on the two segments of 1024 tokens. The results are presented in Table 22. While this training setup allows us to achieve higher performance on in-domain lengths compared to the non-pretrained ARMT, the performance on long-context samples is worse. We hypothesize that this performance drop could arise from the limited length of texts in pre-training. However, this training setup shows promising results and is left for future research.

200000 GR-100++

0.30

Rouge-L

0.25 0.20 0.15

ARMT, Full (0k-65k) ARMT, Long-OOD (32k-65k) Base, Full (0k-65k) Base, Long-OOD (32k-65k)

0.10 0.05 0

50000

100000

150000

Number of train samples

200000

D.6

Figure 5: Scaling synthetic data for the GR dataset using ARMT-Gemma-1B-IT (ARMT) and Gemma-1BIT (Base), trained up to context length 8192. The best performance is achieved with the GR-100+ dataset.

Tables 23 and 25 show the main results on GR-100+ with ARMT model with SmolLM-2-360M-IT backbone and associative layer ablation for this model, while Tables 24 and 26 show the same results on the MT dataset. We also expanded the associative memory ablation study on GR-100+ dataset on SmolLM-2360M-IT model, the results are presented in Tables 27 to 30. As one can see, the importance of associative layers also depends on the model, showing that for SmolLM-2-360M-IT on GR-100+ the 0-th layer is the most important one, while the 23rd is the most important on MT. However, the top-4 associative layers by importance are 0-th, 23-rd, 30th, and 20-th for GR-100+ and 23-rd, 24-th, 17-th, and 14-th for MT, so the importance of the middle layer representations is still high. As was done for Gemma-3-1B-IT, we also show the training dynamics during the curriculum learning for ARMT with the SmolLM-2-360M-IT backbone with pre-selected associative blocks, the results are presented in Tables 31 and 32. The preselected blocks use the same universal pattern as described for Gemma-3-1B-IT, adding associative layers on last layer, two middle layers, two layers between third and fourth quarter, and one layer between first and second quarter. Here we added one more layer between the third and fourth quarters, as the SmolLM-2-360M-IT model has 32 layers instead of 26 for Gemma-3-1B-IT, resulting in associative layers on layers numbered 8, 16, 17, 23, 24, 31. Although this pattern does not add an associative layer on top-1 layer by importance for GR-100+, it still achieves a performance comparable to the full model, using only 20% of associative layers. Moreover, on MT this approach improves performance on long-context samples compared to the full ARMT model.

and with pre-selected associative blocks, the results are presented in Tables 19 and 20. The pre-selected blocks show better final performance while not require additional layers ablation from a fully trained model. We suggest that this universal recipe for associative block selection before training can reduce the number of trained parameters almost without loss of performance. D.4

Ablation Study for Synthetic Training Data Generation

As mentioned in the description of the GR dataset, the base version of the GR dataset is too small to finetune the ARMT model from scratch. In this section, we provide additional results with the base GR dataset without additional synthetic data; the results are provided in Figure 5 and in Table 21. One can notice that the base version of the GR dataset contains enough samples to finetune the Gemma-3-1B-IT model, but not enough to finetune the ARMT model. D.5

Ablation Study for ARMT Backbone

Ablation Study for SFT

We also conducted additional SFT experiments for the ARMT model. We hypothesize that training with the cold-start weights for associative blocks and memory tokens could limit the performance of the model, which is trained on limited data for fine-tuning. To check this, we created an additional synthetic QA dataset with LLM’s generated QA pairs over natural long-context samples. We continuously pre-trained the ARMT model with the Gemma-3-1B-IT backbone and finetuned it on the GR-100+ dataset with the standard curriculum 14

Main results on MT dataset

Main results on GR-100+ dataset

Max. train length & Max. base model context

0.8

Max. train length & Max. base model context

0.35

0.6

0.25

ROUGE-L

Exact match (EM)

0.30

0.4 Base ARMT (ours) ARMT, pre-selected layers, trained (ours) Base, mean ARMT (ours), mean ARMT, pre-selected layers, trained (ours), mean

0.2

0.0 103

0.20 0.15

Base ARMT (ours) ARMT, pre-selected layers, trained (ours) Base, mean ARMT (ours), mean ARMT, pre-selected layers, trained (ours), mean

0.10 0.05 104

103

Context length

(a) MT results.

104

Context length

(b) GR-100+ results.

Figure 6: Main results with SmolLM-2-360M-IT model. (1) SmolLM-2-360M-IT (Base) with full attention trained up to context length 8192 drops in performance after maximum model length of 8k; (2) ARMT model with SmolLM2-360M-IT backbone (ARMT) outperforms base model on OOD context length (after 8k); (3) ARMT fine-tuned with only six pre-selected associative layers (approximately 20%) shows comparable or better performance than ARMT model with all associative layers.

Training Hyperparameters

We also conducted the standard set of experiments on the ContractNLI dataset (Koreeda and Manning, 2021) from the SCROLLS benchmark (Shaham et al., 2022). The results are presented in Table 37. Overall, the proposed training approach demonstrates robustness across different datasets and can be effectively applied to small real-world datasets without requiring additional synthetic data.

The hyperparameters used during training are described in Table 33. All models were trained using two NVIDIA H100-80GB GPUs. The full ARMT fine-tuning process takes approximately 48 hours per run.

Memory Size Ablations

In our experiments, the ARMT model without continued pretraining used 16 memory tokens, whereas the pretrained variant used 32 memory tokens. We hypothesize that a larger memory size is beneficial during pretraining, while 16 memory tokens are sufficient for direct fine-tuning. The comparison between ARMT models with 16 and 32 memory tokens, both trained without pretraining, is presented in Tables 34 and 35. Overall, the model with 16 memory tokens achieves better performance across all data splits and demonstrates stronger generalization during curriculum learning.

ContractNLI Results

BABILong Results

To further validate the proposed ARMT fine-tuning approach, we conducted experiments on the BABILong benchmark (Kuratov et al., 2024), which provides a training set suitable for fine-tuning. Due to the absence of a train set, we did not use other popular long-context benchmarks, such as RULER (Hsieh et al., 2024) and LongBench (Bai et al., 2024). We trained both the base Gemma-31B-IT model and the ARMT model on sequences of up to 5120 tokens. For ARMT, we used a curriculum with 2, 3, and 5 segments of 1024 tokens each, while the base model was trained directly on sequences of length 5120. All models were trained jointly on tasks QA1-QA5 from the BABILong benchmark. The results are presented in Tables 38 and 39. While the base model without fine-tuning achieves near-zero performance on long-context inputs, the fine-tuned base model generalizes to context lengths of up to 32k tokens. In contrast, the ARMT model achieves stronger performance on

RMT Comparison

To compare the ARMT architecture with RMT under identical conditions on a real-world task, we trained the RMT model on the GR-100+ dataset using the same setup as the ARMT model. The results, presented in Table 36, show that ARMT consistently outperforms RMT in all data splits, according to the prior results on synthetic data (Rodkin et al., 2024). 15

the long-context splits and generalizes to context lengths of up to 64k tokens.

Long-context Baselines

We compared the ARMT model with other long-context models, such as Mamba, Mamba-2, DeltaNet, and xLSTM; the results are presented in Table 40. The ARMT model shows comparable or better generalization on OOD and Long-OOD splits, outperforming other long-context baselines. However, Mamba, Mamba-2, DeltaNet, and xLSTM show better in-domain performance, likely due to their substantially more extensive pretraining and a larger number of parameters compared to ARMT. In contrast, the ARMT model is used only with relatively small language modeling pretraining. We also conducted experiments with YaRN (Peng et al., 2024) for context extension. Following the original paper setup, we applied YaRN to the SmolLM-2-360M-IT model and first adapted it on the PG19 dataset (Rae et al., 2020) with 32k token length for 400 steps with a total batch size of 64. After that we fine-tuned the adapted model on the GR-100+ and MT datasets; the results are presented in Tables 41 and 42. The ARMT model outperforms the model with YaRN on both MT and GR-100+, while the base model outperforms the model with YaRN on MT dataset.

Example

Text

Context

from typing import Any import typing [docstring] from alembic import op import sqlalchemy as sa

from sqlalchemy . ext . declarative import declarative_base from sqlalchemy . orm import sessionmaker , relationship [comment] revision = [string] down_revision = [string] Base = declarative_base ( ) db = sa db . Model = Base db . relationship = relationship

def create_session ( ) : connection = op . get_bind ( ) session_maker = sa . orm . sessionmaker ( ) session = session_maker ( bind = connection ) db . session = session

def upgrade ( ) : create_session ( ) [comment] op . alter_column ( [string] , [string] , type_ = sa . Text , existing_type = sa . String ) [comment] def downgrade ( ) : create_session ( ) [comment] op . alter_column ( [string] , [string] , type_ = sa . String , existing_type = sa . Text ) [comment]

[comment] Question Answer

What is the type of variable down_revision? builtins.str

Table 6: Example from MT dataset. The question during training and evaluation is placed before context (i. e. at the start of the prompt).

Example

Text

Context

Economic Significance of Intellectual Property Protection and Theft As we reported in April 2010, IP is an important component of the U.S. economy and IP-related industries pay higher wages and contribute a significant percentage to the U.S. economy. However, the U.S. economy as a whole may grow at a slower pace than it otherwise would because of counterfeiting and piracy’s effect on U.S. industries, government, and consumers. Quantifying Economic Impacts Is Difficult, However Industry Research Suggests the Impacts Are Sizable Generally, as we reported in April 2010, the illicit nature of counterfeiting and piracy makes estimating the economic impact of IP infringements extremely difficult, so assumptions must be used to offset the lack of data. Efforts to estimate losses involve assumptions such as the rate at which consumers would substitute counterfeit for legitimate products, which can have enormous impacts on the resulting estimates. Because of the significant differences in types of counterfeited and pirated goods and industries involved, no single method can be used to develop estimates. Each method has limitations, and most experts observed that it is difficult, if not impossible, to quantify the economy-wide impacts. Nonetheless, research in specific industries suggests that the problem is sizeable. What makes cost-estimates of IP infringements difficult to calculate? Generally, as GAO reported in April 2010, the illicit nature of counterfeiting and piracy makes estimating the economic impact of IP infringements extremely difficult.

Question Answer

Table 7: Example from GR dataset. The question during training and evaluation is placed before context (i. e. at the start of the prompt).

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

0.000 0.000 0.000 0.000 0.000 0.000 0.011 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.001 0.000 0.001

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.782 0.724 0.689 0.616 0.609 0.506 0.374 0.404 0.337 0.229 0.198 0.139 0.089 0.685 0.271 0.116 0.419

0.782 0.714 0.698 0.644 0.717 0.753 0.670 0.702 0.640 0.646 0.584 0.609 0.604 0.711 0.647 0.607 0.670

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

Table 8: Best results on the MT dataset for Gemma-3-1B-IT model, metric - EM. ARMT shows comparable with the base model overall performance, and outperforms it on OOD and Long-OOD.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

0.188 0.161 0.140 0.114 0.113 0.106 0.101 0.091 0.088 0.083 0.078 0.078 0.103

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.357 0.296 0.264 0.240 0.219 0.205 0.213 0.193 0.180 0.115 0.016 0.002 0.000

0.357 0.294 0.275 0.286 0.238 0.233 0.221 0.225 0.217 0.200 0.168 0.128 0.040

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.143

0.351

0.275

0.290

0.299

0.092

0.266

0.162

0.211

0.238

0.084

0.175

0.002

0.108

0.215

0.116

0.306

0.215

0.248

0.266

Table 9: Best results on the GR-100+ dataset for Gemma-3-1B-IT model, ROUGE-L. ARMT shows slightly lower overall performance than the base model, but outperforms it on Long-OOD.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

ARMT, GR-100+, 2k, pretrain

ARMT, GR-100+, 4k, pretrain

ARMT, GR-100+, 8k, pretrain

ARMT w/o CL, GR-100+, 8k, pretrain

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.357 0.296 0.264 0.240 0.219 0.205 0.213 0.193 0.180 0.115 0.016 0.002 0.000

0.357 0.294 0.275 0.286 0.238 0.233 0.221 0.225 0.217 0.200 0.168 0.128 0.040

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.371 0.310 0.268 0.254 0.222 0.135 0.052 0.033 0.026 0.010 0.001 0.000 0.000

0.384 0.320 0.301 0.293 0.256 0.228 0.231 0.237 0.224 0.246 0.194 0.199 0.289

0.370 0.316 0.293 0.303 0.260 0.245 0.260 0.252 0.262 0.260 0.221 0.207 0.300

0.382 0.295 0.302 0.315 0.261 0.238 0.260 0.251 0.244 0.273 0.209 0.223 0.285

0.351

0.275

0.290

0.299

0.285

0.311

0.308

0.311

0.266

0.162

0.211

0.238

0.045

0.231

0.252

0.250

0.175

0.002

0.108

0.215

0.000

0.220

0.228

0.237

0.306

0.215

0.248

0.266

0.157

0.268

0.278

Table 10: Results on the GR-100+ dataset for Gemma-3-1B-IT model with ARMT after continuous pretraining, ROUGE-L. ARMT after continuous pretraining shows significantly better results on all splits.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

ARMT, MT, 2k, pretrain

ARMT, MT, 4k, pretrain

ARMT, MT, 8k, pretrain

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.782 0.724 0.689 0.616 0.609 0.506 0.374 0.404 0.337 0.229 0.198 0.139 0.089 0.685 0.271 0.116 0.419

0.782 0.714 0.698 0.644 0.717 0.753 0.670 0.702 0.640 0.646 0.584 0.609 0.604 0.711 0.647 0.607 0.670

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.769 0.752 0.708 0.671 0.793 0.753 0.692 0.745 0.629 0.632 0.584 0.530 0.426 0.740 0.618 0.481 0.661

0.756 0.743 0.708 0.671 0.804 0.790 0.747 0.734 0.640 0.639 0.673 0.609 0.545 0.738 0.665 0.579 0.691

0.756 0.743 0.736 0.740 0.859 0.827 0.868 0.766 0.697 0.757 0.812 0.791 0.762 0.767 0.783 0.777 0.777

Table 11: Results on the MT dataset for Gemma-3-1B-IT model with ARMT after continuous pretraining, EM. ARMT after continuous pretraining shows significantly better results on all splits.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, MT, 8k

ARMT, MT, 8k

ARMT, MT, 8k, w/o layers 0-6

ARMT, MT, 8k, w/o layers 7-12

ARMT, MT, 8k, w/o layers 13-18

ARMT, MT, 8k, w/o layers 19-25

ARMT, MT, 8k, only top-1 layer

ARMT, MT, 8k, only top-4 layers

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713

0.756 0.733 0.736 0.699 0.815 0.753 0.780 0.766 0.674 0.750 0.703 0.757 0.723

0.756 0.743 0.726 0.699 0.804 0.765 0.791 0.766 0.674 0.764 0.703 0.765 0.723

0.756 0.571 0.443 0.329 0.315 0.395 0.341 0.277 0.337 0.306 0.287 0.226 0.356

0.756 0.676 0.717 0.671 0.739 0.716 0.736 0.809 0.708 0.722 0.683 0.739 0.683

0.756 0.590 0.594 0.521 0.598 0.531 0.505 0.457 0.472 0.458 0.436 0.504 0.505

0.756 0.724 0.708 0.712 0.793 0.765 0.758 0.755 0.663 0.729 0.693 0.765 0.713

0.797

0.749

0.747

0.482

0.711

0.610

0.738

0.709

0.741

0.739

0.745

0.311

0.724

0.481

0.730

0.463

0.741

0.745

0.287

0.713

0.504

0.741

0.744

0.743

0.746

0.372

0.720

0.527

0.733

Table 12: Associative layers ablation on the MT dataset for Gemma-3-1B-IT model, metric - EM. Middle and upper layers representations are the most important for associative memory; ARMT with only top-4 associative blocks keeps almost the same performance as the full ARMT.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k, w/o layers 0-6

ARMT, GR-100+, 8k, w/o layers 7-12

ARMT, GR-100+, 8k, w/o layers 13-18

ARMT, GR-100+, 8k, w/o layers 19-25

ARMT, GR-100+, 8k, only top-1 layer

ARMT, GR-100+, 8k, only top-4 layers

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189 0.351 0.266 0.175 0.306

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.306 0.277 0.295 0.246 0.238 0.237 0.233 0.234 0.248 0.212 0.213 0.327 0.296 0.237 0.239 0.264

0.358 0.301 0.279 0.273 0.243 0.219 0.226 0.223 0.222 0.220 0.206 0.220 0.286 0.291 0.222 0.235 0.254

0.358 0.232 0.186 0.165 0.165 0.155 0.152 0.156 0.154 0.155 0.140 0.152 0.233 0.220 0.154 0.171 0.185

0.358 0.234 0.212 0.198 0.188 0.190 0.188 0.176 0.169 0.182 0.187 0.160 0.205 0.237 0.181 0.170 0.207

0.358 0.207 0.180 0.160 0.144 0.133 0.129 0.146 0.140 0.121 0.129 0.147 0.207 0.209 0.135 0.161 0.170

0.358 0.299 0.285 0.278 0.244 0.233 0.241 0.230 0.237 0.247 0.217 0.229 0.331 0.293 0.237 0.253 0.263

Table 13: Associative layers ablation on the GR-100+ dataset for Gemma-3-1B-IT model, metric - ROUGE-L. Middle and upper layers representations are the most important for associative memory; ARMT with only top-4 associative blocks keeps almost the same performance as the full ARMT.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k, only top-4 layers

ARMT, GR-100+, 8k, only top-4 layers, trained

ARMT, GR-100+, 8k, pre-selected layers, trained

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.358 0.299 0.285 0.278 0.244 0.233 0.241 0.230 0.237 0.247 0.217 0.229 0.331

0.347 0.291 0.265 0.262 0.241 0.232 0.232 0.222 0.233 0.242 0.187 0.206 0.338

0.363 0.317 0.304 0.287 0.250 0.231 0.231 0.221 0.243 0.259 0.204 0.205 0.246

0.351

0.299

0.293

0.281

0.304

0.266

0.238

0.237

0.229

0.234

0.175

0.215

0.253

0.236

0.214

0.306

0.266

0.263

0.254

0.267

Table 14: Associative layers ablation on the GR-100+ for Gemma-3-1B-IT model, metric - ROUGE-L. ARMT with only top-4 associative blocks keeps almost the same performance as the full ARMT even without training. Model with 5 associative blocks (approximately 20%) achieves the same performance as the full ARMT.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

W/o layer 0

W/o layer 1

W/o layer 2

W/o layer 3

W/o layer 4

W/o layer 5

W/o layer 6

W/o layer 7

W/o layer 8

W/o layer 9

W/o layer 10

W/o layer 11

W/o layer 12

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189 0.351 0.266 0.175 0.306

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.307 0.278 0.298 0.247 0.228 0.231 0.238 0.238 0.250 0.212 0.199 0.271 0.297 0.235 0.216 0.264

0.358 0.300 0.274 0.299 0.253 0.232 0.236 0.237 0.242 0.248 0.203 0.203 0.344 0.296 0.237 0.236 0.265

0.358 0.295 0.279 0.298 0.255 0.234 0.241 0.234 0.236 0.254 0.205 0.206 0.278 0.297 0.237 0.223 0.265

0.358 0.302 0.280 0.297 0.246 0.239 0.244 0.230 0.238 0.257 0.214 0.214 0.284 0.296 0.239 0.230 0.266

0.358 0.296 0.283 0.297 0.255 0.243 0.240 0.233 0.239 0.253 0.204 0.205 0.272 0.297 0.238 0.220 0.266

0.358 0.299 0.283 0.298 0.256 0.233 0.239 0.232 0.240 0.258 0.210 0.211 0.310 0.298 0.238 0.234 0.266

0.358 0.303 0.279 0.301 0.254 0.237 0.244 0.235 0.240 0.249 0.204 0.204 0.264 0.299 0.238 0.218 0.266

0.358 0.294 0.278 0.292 0.246 0.243 0.240 0.237 0.240 0.249 0.205 0.216 0.282 0.293 0.239 0.231 0.264

0.358 0.303 0.277 0.302 0.249 0.237 0.241 0.231 0.243 0.257 0.216 0.223 0.378 0.297 0.241 0.259 0.267

0.358 0.290 0.269 0.280 0.240 0.211 0.227 0.216 0.218 0.228 0.190 0.170 0.170 0.287 0.216 0.170 0.249

0.358 0.302 0.282 0.299 0.252 0.236 0.239 0.234 0.240 0.254 0.206 0.198 0.264 0.298 0.237 0.213 0.266

0.358 0.303 0.280 0.301 0.254 0.234 0.242 0.235 0.240 0.251 0.204 0.204 0.264 0.299 0.237 0.218 0.266

0.358 0.302 0.296 0.285 0.254 0.231 0.242 0.227 0.239 0.246 0.216 0.226 0.284 0.299 0.236 0.239 0.265

Table 15: Ablation for all associative layers on the GovReport-100+ dataset for ARMT with Gemma-3-1B-IT model, metric - ROUGE-L. Ablated associative layers from 0 to 12.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

W/o layer 13

W/o layer 14

W/o layer 15

W/o layer 16

W/o layer 17

W/o layer 18

W/o layer 19

W/o layer 20

W/o layer 21

W/o layer 22

W/o layer 23

W/o layer 24

W/o layer 25

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189 0.351 0.266 0.175 0.306

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.301 0.277 0.270 0.239 0.222 0.224 0.215 0.213 0.211 0.209 0.189 0.230 0.289 0.215 0.198 0.250

0.358 0.239 0.219 0.197 0.184 0.176 0.179 0.176 0.177 0.177 0.163 0.149 0.239 0.239 0.176 0.170 0.205

0.358 0.304 0.279 0.302 0.254 0.235 0.242 0.237 0.241 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.267

0.358 0.303 0.279 0.303 0.253 0.241 0.239 0.237 0.241 0.251 0.207 0.200 0.264 0.299 0.238 0.215 0.267

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.289 0.274 0.287 0.263 0.232 0.232 0.234 0.233 0.249 0.209 0.222 0.273 0.294 0.234 0.234 0.262

0.358 0.291 0.278 0.297 0.253 0.241 0.237 0.236 0.242 0.252 0.205 0.202 0.264 0.295 0.238 0.216 0.265

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.238 0.239 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.203 0.264 0.299 0.238 0.217 0.266

0.358 0.296 0.279 0.301 0.253 0.234 0.241 0.235 0.243 0.251 0.208 0.203 0.264 0.297 0.238 0.217 0.265

0.358 0.226 0.208 0.190 0.196 0.188 0.186 0.174 0.171 0.183 0.190 0.163 0.205 0.235 0.180 0.173 0.206

Table 16: Ablation for all associative layers on the GovReport-100+ dataset for ARMT with Gemma-3-1B-IT model, metric - ROUGE-L. Ablated associative layers from 13 to 25.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 8k

W/o layer 0

W/o layer 1

W/o layer 2

W/o layer 3

W/o layer 4

W/o layer 5

W/o layer 6

W/o layer 7

W/o layer 8

W/o layer 9

W/o layer 10

W/o layer 11

W/o layer 12

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.733 0.736 0.699 0.804 0.753 0.791 0.734 0.685 0.764 0.713 0.757 0.713 0.747 0.740 0.736 0.743

0.756 0.733 0.736 0.699 0.804 0.753 0.769 0.755 0.697 0.764 0.723 0.765 0.723 0.747 0.745 0.745 0.746

0.756 0.733 0.736 0.685 0.815 0.741 0.791 0.734 0.685 0.764 0.703 0.757 0.713 0.747 0.738 0.736 0.741

0.756 0.733 0.736 0.699 0.815 0.765 0.802 0.723 0.674 0.750 0.713 0.757 0.733 0.749 0.740 0.746 0.743

0.756 0.733 0.736 0.699 0.804 0.741 0.791 0.734 0.685 0.764 0.713 0.757 0.723 0.747 0.740 0.741 0.743

0.756 0.733 0.736 0.699 0.815 0.728 0.802 0.734 0.685 0.757 0.703 0.757 0.693 0.749 0.734 0.727 0.739

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.734 0.685 0.764 0.703 0.757 0.713 0.749 0.739 0.736 0.743

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.703 0.757 0.713 0.749 0.739 0.736 0.743

0.756 0.743 0.726 0.699 0.815 0.778 0.813 0.766 0.674 0.764 0.693 0.757 0.743 0.749 0.749 0.750 0.749

0.756 0.733 0.736 0.685 0.804 0.753 0.791 0.745 0.685 0.771 0.713 0.765 0.713 0.744 0.744 0.741 0.744

0.756 0.743 0.736 0.699 0.804 0.741 0.791 0.745 0.685 0.764 0.713 0.757 0.713 0.749 0.740 0.736 0.743

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.733 0.736 0.699 0.804 0.741 0.791 0.745 0.685 0.764 0.703 0.757 0.713 0.747 0.739 0.736 0.742

Table 17: Ablation for all associative layers on the MT dataset for ARMT with Gemma-3-1B-IT model, metric EM. Ablated associative layers from 0 to 12.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 8k

W/o layer 13

W/o layer 14

W/o layer 15

W/o layer 16

W/o layer 17

W/o layer 18

W/o layer 19

W/o layer 20

W/o layer 21

W/o layer 22

W/o layer 23

W/o layer 24

W/o layer 25

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.733 0.736 0.699 0.815 0.741 0.791 0.734 0.685 0.757 0.703 0.757 0.713 0.749 0.737 0.736 0.741

0.756 0.610 0.491 0.384 0.370 0.407 0.330 0.255 0.360 0.333 0.287 0.278 0.356 0.522 0.323 0.314 0.394

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.757 0.713 0.749 0.740 0.736 0.743

0.756 0.695 0.726 0.671 0.772 0.716 0.736 0.660 0.618 0.715 0.634 0.713 0.683 0.724 0.686 0.699 0.700

0.756 0.695 0.736 0.712 0.793 0.765 0.780 0.755 0.697 0.743 0.703 0.765 0.713 0.738 0.740 0.741 0.739

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.764 0.703 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.752 0.726 0.685 0.793 0.765 0.769 0.755 0.719 0.764 0.703 0.765 0.723 0.744 0.746 0.745 0.745

0.756 0.743 0.736 0.699 0.804 0.753 0.769 0.755 0.697 0.764 0.713 0.748 0.713 0.749 0.740 0.732 0.743

0.756 0.733 0.726 0.671 0.783 0.741 0.769 0.734 0.697 0.750 0.723 0.765 0.733 0.735 0.740 0.750 0.739

Table 18: Ablation for all associative layers on the MT dataset for ARMT with Gemma-3-1B-IT model, metric EM. Ablated associative layers from 13 to 25.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k, only top-4 layers

ARMT, GR-100+, 2k, only top-4 layers, trained

ARMT, GR-100+, 4k, only top-4 layers, trained

ARMT, GR-100+, 8k, only top-4 layers, trained

ARMT, GR-100+, 2k, pre-selected layers, trained

ARMT, GR-100+, 4k, pre-selected layers, trained

ARMT, GR-100+, 8k, pre-selected layers, trained

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.357 0.296 0.264 0.240 0.219 0.205 0.213 0.193 0.180 0.115 0.016 0.002 0.000

0.357 0.294 0.275 0.286 0.238 0.233 0.221 0.225 0.217 0.200 0.168 0.128 0.040

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.358 0.299 0.285 0.278 0.244 0.233 0.241 0.230 0.237 0.247 0.217 0.229 0.331

0.361 0.276 0.252 0.220 0.212 0.182 0.188 0.178 0.174 0.179 0.160 0.138 0.210

0.366 0.285 0.265 0.260 0.239 0.229 0.216 0.226 0.218 0.232 0.192 0.183 0.288

0.347 0.291 0.265 0.262 0.241 0.232 0.232 0.222 0.233 0.242 0.187 0.206 0.338

0.351 0.290 0.263 0.237 0.194 0.125 0.108 0.067 0.060 0.062 0.052 0.025 0.022

0.353 0.316 0.297 0.280 0.242 0.234 0.225 0.205 0.208 0.169 0.106 0.088 0.049

0.363 0.317 0.304 0.287 0.250 0.231 0.231 0.221 0.243 0.259 0.204 0.205 0.246

0.351

0.275

0.290

0.299

0.293

0.264

0.283

0.281

0.267

0.297

0.304

0.266

0.162

0.211

0.238

0.237

0.178

0.222

0.229

0.079

0.195

0.234

0.175

0.002

0.108

0.215

0.253

0.155

0.207

0.236

0.024

0.079

0.214

0.306

0.215

0.248

0.266

0.263

0.218

0.250

0.254

0.167

0.243

0.267

Table 19: Associative layers ablation on the GR-100+ dataset for Gemma-3-1B-IT model, metric - ROUGE-L. ARMT with only top-4 associative blocks keeps almost the same performance as the full ARMT even without training. Model with 5 associative blocks (approximately 20%) achieves the same performance as the full ARMT model.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

ARMT, MT, 8k, only top-4 layers

ARMT, MT, 2k, only top-4 layers, trained

ARMT, MT, 4k, only top-4 layers, trained

ARMT, MT, 8k, only top-4 layers, trained

ARMT, MT, 2k, pre-selected layers, trained

ARMT, MT, 4k, pre-selected layers, trained

ARMT, MT, 8k, pre-selected layers, trained

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317

0.782 0.724 0.689 0.616 0.609 0.506 0.374 0.404 0.337 0.229 0.198 0.139 0.089

0.782 0.714 0.698 0.644 0.717 0.753 0.670 0.702 0.640 0.646 0.584 0.609 0.604

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713

0.756 0.724 0.708 0.712 0.793 0.765 0.758 0.755 0.663 0.729 0.693 0.765 0.713

0.782 0.714 0.660 0.452 0.185 0.099 0.022 0.032 0.011 0.000 0.000 0.000 0.000

0.782 0.686 0.717 0.699 0.793 0.790 0.747 0.691 0.584 0.542 0.396 0.261 0.277

0.756 0.714 0.736 0.699 0.815 0.827 0.802 0.777 0.708 0.757 0.762 0.765 0.752

0.782 0.724 0.679 0.397 0.130 0.012 0.000 0.000 0.000 0.000 0.000 0.000 0.000

0.756 0.705 0.717 0.699 0.815 0.691 0.747 0.766 0.584 0.604 0.436 0.435 0.485

0.756 0.705 0.717 0.712 0.837 0.840 0.857 0.777 0.685 0.812 0.762 0.800 0.802

0.797

0.685

0.711

0.749

0.738

0.564

0.734

0.744

0.551

0.738

0.744

0.709

0.271

0.647

0.741

0.730

0.017

0.521

0.767

0.001

0.586

0.793

0.463

0.116

0.607

0.741

0.000

0.268

0.759

0.000

0.458

0.801

0.741

0.419

0.670

0.744

0.733

0.213

0.597

0.759

0.198

0.640

0.776

Table 20: Associative layers ablation on the MT dataset for Gemma-3-1B-IT model, metric - EM. ARMT with only top-4 associative blocks keeps almost the same performance as the full ARMT even without training. Model with 5 associative blocks (approximately 20%) achieves the same performance as the full ARMT model.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, GR, 8k

ARMT, GR, 2k

ARMT, GR, 4k

ARMT, GR, 8k

0.188 0.161 0.140 0.114 0.113 0.106 0.101 0.091 0.088 0.083 0.078 0.078 0.103 0.143 0.092 0.084 0.116

0.347 0.358 0.305 0.288 0.289 0.271 0.232 0.241 0.218 0.222 0.149 0.130 0.072 0.317 0.226 0.117 0.269

0.340 0.175 0.151 0.135 0.146 0.146 0.132 0.120 0.128 0.128 0.111 0.123 0.178 0.189 0.130 0.136 0.157

0.326 0.198 0.184 0.175 0.164 0.152 0.166 0.147 0.149 0.145 0.138 0.135 0.190 0.209 0.150 0.148 0.178

0.313 0.211 0.200 0.180 0.176 0.164 0.177 0.156 0.168 0.164 0.175 0.144 0.201 0.216 0.165 0.157 0.189

Table 21: Best results on the GR dataset for Gemma-3-1B-IT model, metric - ROUGE-L.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, GR-100+, 8k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k, Pretrained

0.188 0.161 0.140 0.114 0.113 0.106 0.101 0.091 0.088 0.083 0.078 0.078 0.103 0.143 0.092 0.084 0.116

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189 0.351 0.266 0.175 0.306

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264 0.299 0.238 0.215 0.266

0.363 0.332 0.282 0.284 0.278 0.244 0.251 0.242 0.252 0.249 0.214 0.183 0.018 0.307 0.241 0.145 0.272

Table 22: Best results on the GovReport-100+ dataset for Gemma-3-1B-IT model, comparison with the pretraining on synthetic QA task, metric - ROUGE-L.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

0.145 0.122 0.113 0.091 0.094 0.089 0.039 0.039 0.039 0.048 0.042 0.044 0.088

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.335 0.289 0.238 0.220 0.212 0.177 0.156 0.137 0.074 0.046 0.011 0.025 0.015

0.332 0.297 0.257 0.245 0.215 0.202 0.199 0.195 0.190 0.140 0.059 0.032 0.011

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.113

0.327

0.258

0.269

0.275

0.050

0.074

0.106

0.170

0.211

0.054

0.072

0.023

0.027

0.154

0.080

0.192

0.178

0.216

0.241

Table 23: Best results on the GR-100+ dataset for SmolLM-2-360M-IT model, ROUGE-L. ARMT shows higher overall performance than the base model, and outperforms it on OOD and Long-OOD.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

0.000 0.000 0.000 0.000 0.000 0.000 0.011 0.000 0.000 0.000 0.020 0.009 0.000 0.000 0.005 0.005 0.003

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000 0.771 0.058 0.000 0.313

0.718 0.638 0.642 0.603 0.598 0.469 0.418 0.138 0.000 0.014 0.000 0.000 0.000 0.639 0.112 0.000 0.300

0.756 0.686 0.717 0.630 0.739 0.802 0.692 0.755 0.663 0.708 0.257 0.000 0.000 0.707 0.473 0.000 0.557

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030 0.714 0.643 0.403 0.668

Table 24: Best results on the MT dataset for SmolLM-2-360M-IT model, metric - EM. ARMT shows better overall performance as the base model, and outperforms it on OOD and Long-OOD.

Model/ Lengths

Base, GR-100+, 8k

ARMT, GR-100+, 8k

ARMT, GR-100+, 8k, w/o layers 0-7

ARMT, GR-100+, 8k, w/o layers 8-15

ARMT, GR-100+, 8k, w/o layers 16-23

ARMT, GR-100+, 8k, w/o layers 24-31

ARMT, GR-100+, 8k, only top-1 layer

ARMT, GR-100+, 8k, only top-4 layers

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.337 0.301 0.246 0.224 0.187 0.166 0.152 0.109 0.092 0.085 0.066 0.065 0.049

0.338 0.299 0.258 0.234 0.218 0.206 0.211 0.201 0.199 0.192 0.168 0.153 0.116

0.334 0.192 0.175 0.175 0.159 0.153 0.156 0.151 0.144 0.146 0.140 0.147 0.190

0.331 0.208 0.175 0.178 0.177 0.163 0.174 0.159 0.153 0.155 0.151 0.141 0.090

0.330 0.142 0.134 0.109 0.120 0.119 0.114 0.113 0.121 0.113 0.114 0.123 0.113

0.332 0.197 0.183 0.155 0.143 0.146 0.153 0.146 0.146 0.138 0.142 0.135 0.158

0.327

0.275

0.259

0.269

0.206

0.213

0.166

0.202

0.074

0.211

0.114

0.197

0.150

0.158

0.116

0.145

0.072

0.154

0.061

0.144

0.157

0.129

0.121

0.140

0.192

0.241

0.182

0.231

0.176

0.184

0.140

0.171

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Table 25: Associative layers ablation on the GR-100+ dataset for SmolLM-2-360M-IT model, metric - ROUGE-L. Middle and upper layers representations are the most important for associative memory, but for SmolLM-2-360M-IT backbone the lower layers are also important. Model/ Lengths

Base, MT, 8k

ARMT, MT, 8k

ARMT, MT, 8k, w/o layers 0-7

ARMT, MT, 8k, w/o layers 8-15

ARMT, MT, 8k, w/o layers 16-23

ARMT, MT, 8k, w/o layers 24-31

ARMT, MT, 8k, only top-1 layer

ARMT, MT, 8k, only top-4 layers

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000 0.771 0.058 0.000 0.313

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030 0.714 0.643 0.403 0.668

0.769 0.695 0.708 0.671 0.717 0.790 0.736 0.755 0.618 0.715 0.703 0.678 0.525 0.711 0.689 0.606 0.697

0.769 0.695 0.698 0.658 0.761 0.778 0.769 0.713 0.652 0.701 0.683 0.730 0.030 0.716 0.631 0.403 0.661

0.756 0.438 0.387 0.342 0.315 0.321 0.374 0.255 0.292 0.243 0.287 0.235 0.030 0.440 0.250 0.139 0.318

0.756 0.619 0.651 0.616 0.750 0.716 0.714 0.734 0.573 0.653 0.644 0.635 0.030 0.676 0.586 0.352 0.618

0.756 0.571 0.613 0.548 0.598 0.519 0.462 0.479 0.360 0.479 0.376 0.461 0.475 0.614 0.452 0.468 0.510

0.756 0.638 0.623 0.658 0.750 0.691 0.725 0.723 0.562 0.653 0.604 0.626 0.644 0.681 0.652 0.634 0.662

Table 26: Associative layers ablation on the MT dataset for SmolLM-2-360M-IT model, metric - EM. Middle and upper layers representations are the most important for associative memory. Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

W/o layer 0

W/o layer 1

W/o layer 2

W/o layer 3

W/o layer 4

W/o layer 5

W/o layer 6

W/o layer 7

W/o layer 8

W/o layer 9

W/o layer 10

W/o layer 11

W/o layer 12

W/o layer 13

W/o layer 14

W/o layer 15

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.337 0.298 0.241 0.237 0.201 0.192 0.169 0.155 0.131 0.107 0.068 0.063 0.062

0.338 0.301 0.259 0.257 0.221 0.219 0.214 0.217 0.214 0.218 0.199 0.192 0.159

0.337 0.303 0.253 0.255 0.223 0.216 0.215 0.220 0.218 0.213 0.192 0.189 0.099

0.337 0.302 0.254 0.252 0.217 0.224 0.210 0.214 0.211 0.215 0.198 0.152 0.076

0.337 0.298 0.256 0.257 0.225 0.218 0.215 0.215 0.218 0.212 0.196 0.181 0.098

0.337 0.299 0.252 0.259 0.220 0.222 0.217 0.210 0.212 0.215 0.188 0.173 0.123

0.337 0.304 0.258 0.260 0.228 0.229 0.216 0.223 0.214 0.223 0.201 0.175 0.136

0.337 0.290 0.260 0.252 0.219 0.218 0.214 0.213 0.211 0.216 0.203 0.169 0.095

0.338 0.290 0.256 0.250 0.223 0.222 0.203 0.213 0.212 0.213 0.205 0.181 0.074

0.337 0.298 0.256 0.248 0.218 0.214 0.221 0.219 0.216 0.218 0.183 0.189 0.118

0.337 0.298 0.259 0.243 0.219 0.218 0.209 0.209 0.217 0.205 0.166 0.148 0.080

0.338 0.292 0.256 0.254 0.221 0.213 0.211 0.213 0.211 0.218 0.196 0.170 0.085

0.337 0.298 0.259 0.259 0.220 0.216 0.215 0.219 0.218 0.217 0.189 0.182 0.093

0.337 0.302 0.260 0.261 0.224 0.222 0.217 0.217 0.210 0.220 0.196 0.174 0.128

0.337 0.299 0.263 0.255 0.223 0.219 0.219 0.213 0.219 0.227 0.194 0.198 0.081

0.337 0.295 0.261 0.249 0.225 0.217 0.214 0.219 0.211 0.222 0.183 0.154 0.073

0.327

0.275

0.262

0.275

0.274

0.272

0.274

0.273

0.277

0.271

0.272

0.274

0.276

0.275

0.273

0.074

0.211

0.141

0.214

0.213

0.210

0.212

0.211

0.217

0.211

0.210

0.213

0.205

0.209

0.213

0.216

0.211

0.072

0.154

0.063

0.184

0.168

0.134

0.162

0.161

0.166

0.152

0.156

0.173

0.132

0.150

0.161

0.163

0.171

0.135

0.192

0.241

0.198

0.242

0.241

0.239

0.241

0.240

0.245

0.239

0.238

0.240

0.236

0.239

0.242

0.243

0.240

Table 27: Ablation for all associative layers on the GovReport-100+ dataset for ARMT with SmolLM-2-360M-IT model, metric - ROUGE-L. Ablated associative layers from 0 to 15.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 8k

W/o layer 16

W/o layer 17

W/o layer 18

W/o layer 19

W/o layer 20

W/o layer 21

W/o layer 22

W/o layer 23

W/o layer 24

W/o layer 25

W/o layer 26

W/o layer 27

W/o layer 28

W/o layer 29

W/o layer 30

W/o layer 31

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.337 0.296 0.251 0.253 0.222 0.223 0.211 0.223 0.212 0.223 0.201 0.189 0.106

0.337 0.296 0.256 0.253 0.229 0.224 0.210 0.208 0.213 0.222 0.186 0.176 0.078

0.337 0.295 0.260 0.253 0.216 0.221 0.219 0.215 0.211 0.221 0.196 0.168 0.109

0.338 0.287 0.257 0.239 0.226 0.204 0.208 0.206 0.213 0.209 0.201 0.172 0.080

0.337 0.285 0.258 0.240 0.231 0.208 0.198 0.197 0.202 0.209 0.196 0.131 0.054

0.338 0.286 0.259 0.253 0.216 0.217 0.206 0.208 0.209 0.213 0.192 0.155 0.092

0.335 0.273 0.255 0.235 0.231 0.210 0.214 0.214 0.213 0.211 0.185 0.181 0.087

0.337 0.239 0.212 0.227 0.213 0.193 0.188 0.180 0.203 0.191 0.164 0.151 0.109

0.336 0.288 0.241 0.246 0.217 0.210 0.201 0.214 0.211 0.208 0.197 0.173 0.144

0.337 0.298 0.260 0.258 0.230 0.215 0.210 0.215 0.214 0.216 0.195 0.163 0.113

0.337 0.303 0.262 0.263 0.226 0.216 0.216 0.212 0.206 0.216 0.198 0.176 0.090

0.337 0.277 0.246 0.237 0.215 0.208 0.208 0.211 0.211 0.216 0.182 0.171 0.094

0.338 0.298 0.256 0.260 0.226 0.220 0.211 0.219 0.212 0.222 0.189 0.161 0.093

0.336 0.281 0.248 0.246 0.226 0.207 0.213 0.202 0.213 0.203 0.185 0.159 0.086

0.337 0.267 0.236 0.240 0.224 0.202 0.212 0.207 0.200 0.191 0.185 0.155 0.114

0.337 0.291 0.264 0.253 0.223 0.216 0.211 0.214 0.210 0.220 0.192 0.164 0.103

0.327

0.275

0.271

0.274

0.272

0.269

0.270

0.266

0.245

0.265

0.276

0.278

0.262

0.275

0.267

0.260

0.273

0.074

0.211

0.215

0.211

0.213

0.205

0.198

0.206

0.208

0.187

0.206

0.210

0.206

0.212

0.203

0.199

0.210

0.072

0.154

0.170

0.153

0.154

0.151

0.113

0.140

0.159

0.141

0.166

0.151

0.156

0.153

0.145

0.142

0.146

0.150

0.192

0.241

0.242

0.240

0.241

0.235

0.232

0.236

0.235

0.214

0.234

0.241

0.242

0.232

0.242

0.233

0.227

0.240

Table 28: Ablation for all associative layers on the GovReport-100+ dataset for ARMT with SmolLM-2-360M-IT model, metric - ROUGE-L. Ablated associative layers from 16 to 31.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, MT, 8k

ARMT, MT, 8k

W/o layer 0

W/o layer 1

W/o layer 2

W/o layer 3

W/o layer 4

W/o layer 5

W/o layer 6

W/o layer 7

W/o layer 8

W/o layer 9

W/o layer 10

W/o layer 11

W/o layer 12

W/o layer 13

W/o layer 14

W/o layer 15

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030

0.769 0.695 0.689 0.644 0.761 0.815 0.736 0.766 0.629 0.701 0.703 0.722 0.634

0.769 0.705 0.689 0.658 0.761 0.815 0.747 0.755 0.640 0.729 0.733 0.739 0.040

0.769 0.705 0.698 0.644 0.761 0.802 0.747 0.755 0.629 0.708 0.723 0.722 0.050

0.769 0.705 0.679 0.644 0.772 0.802 0.758 0.766 0.618 0.722 0.733 0.713 0.059

0.769 0.705 0.698 0.644 0.750 0.802 0.747 0.755 0.629 0.722 0.752 0.730 0.040

0.769 0.705 0.698 0.630 0.772 0.790 0.747 0.755 0.663 0.722 0.733 0.730 0.040

0.769 0.705 0.698 0.644 0.772 0.815 0.747 0.755 0.629 0.722 0.743 0.730 0.030

0.769 0.695 0.708 0.658 0.761 0.790 0.758 0.755 0.640 0.722 0.743 0.722 0.040

0.769 0.705 0.698 0.644 0.761 0.802 0.736 0.755 0.629 0.701 0.723 0.739 0.059

0.769 0.705 0.689 0.644 0.761 0.802 0.747 0.745 0.663 0.722 0.733 0.713 0.040

0.769 0.705 0.689 0.644 0.772 0.815 0.736 0.755 0.652 0.715 0.743 0.730 0.030

0.769 0.695 0.689 0.644 0.772 0.790 0.747 0.745 0.640 0.715 0.733 0.739 0.030

0.769 0.705 0.698 0.644 0.761 0.802 0.747 0.755 0.629 0.722 0.743 0.722 0.040

0.769 0.705 0.698 0.630 0.761 0.802 0.747 0.755 0.629 0.715 0.723 0.730 0.030

0.769 0.705 0.698 0.658 0.750 0.765 0.736 0.713 0.652 0.715 0.733 0.730 0.030

0.769 0.705 0.698 0.658 0.761 0.802 0.747 0.755 0.629 0.722 0.743 0.739 0.030

0.771

0.714

0.711

0.716

0.714

0.716

0.718

0.716

0.714

0.716

0.714

0.716

0.714

0.716

0.718

0.058

0.643

0.711

0.649

0.641

0.646

0.647

0.646

0.642

0.645

0.646

0.642

0.645

0.641

0.635

0.646

0.000

0.403

0.681

0.412

0.408

0.407

0.403

0.421

0.398

0.403

0.407

0.403

0.407

0.313

0.668

0.711

0.673

0.668

0.670

0.671

0.672

0.668

0.669

0.671

0.668

0.670

0.667

0.664

0.672

Table 29: Ablation for all associative layers on the MT dataset for ARMT with SmolLM-2-360M-IT model, metric EM. Ablated associative layers from 0 to 15.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, MT, 8k

ARMT, MT, 8k

W/o layer 16

W/o layer 17

W/o layer 18

W/o layer 19

W/o layer 20

W/o layer 21

W/o layer 22

W/o layer 23

W/o layer 24

W/o layer 25

W/o layer 26

W/o layer 27

W/o layer 28

W/o layer 29

W/o layer 30

W/o layer 31

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030

0.769 0.705 0.698 0.644 0.761 0.815 0.736 0.755 0.629 0.715 0.733 0.722 0.030

0.769 0.705 0.689 0.658 0.717 0.741 0.703 0.755 0.652 0.681 0.693 0.670 0.030

0.769 0.705 0.698 0.658 0.761 0.802 0.747 0.755 0.629 0.722 0.733 0.722 0.030

0.769 0.705 0.698 0.644 0.761 0.802 0.736 0.755 0.618 0.708 0.743 0.713 0.030

0.769 0.705 0.698 0.658 0.761 0.815 0.747 0.755 0.629 0.708 0.743 0.730 0.059

0.769 0.705 0.698 0.644 0.804 0.802 0.736 0.755 0.629 0.708 0.733 0.730 0.030

0.769 0.705 0.689 0.630 0.761 0.802 0.747 0.755 0.629 0.715 0.723 0.730 0.050

0.756 0.448 0.415 0.370 0.348 0.296 0.407 0.298 0.326 0.236 0.277 0.243 0.030

0.769 0.667 0.698 0.616 0.696 0.753 0.725 0.745 0.640 0.681 0.693 0.704 0.030

0.769 0.714 0.698 0.644 0.761 0.815 0.725 0.745 0.629 0.715 0.723 0.730 0.030

0.769 0.705 0.698 0.630 0.761 0.815 0.725 0.766 0.629 0.722 0.723 0.730 0.030

0.769 0.667 0.698 0.658 0.750 0.802 0.769 0.755 0.618 0.722 0.733 0.730 0.030

0.756 0.686 0.670 0.658 0.815 0.815 0.725 0.734 0.663 0.715 0.713 0.722 0.030

0.769 0.657 0.708 0.658 0.783 0.802 0.758 0.766 0.629 0.708 0.733 0.730 0.040

0.756 0.686 0.698 0.630 0.761 0.815 0.736 0.766 0.629 0.729 0.733 0.730 0.030

0.769 0.695 0.698 0.630 0.761 0.802 0.758 0.745 0.629 0.715 0.733 0.739 0.030

0.771

0.714

0.716

0.707

0.718

0.716

0.718

0.725

0.712

0.460

0.689

0.718

0.714

0.707

0.716

0.714

0.707

0.711

0.058

0.643

0.641

0.614

0.642

0.637

0.647

0.640

0.643

0.259

0.620

0.638

0.642

0.644

0.639

0.644

0.646

0.643

0.000

0.403

0.398

0.371

0.398

0.394

0.416

0.403

0.412

0.143

0.389

0.403

0.398

0.407

0.403

0.407

0.313

0.668

0.647

0.669

0.665

0.672

0.670

0.668

0.331

0.645

0.667

0.668

0.667

0.666

0.669

0.668

Table 30: Ablation for all associative layers on the MT dataset for ARMT with SmolLM-2-360M-IT model, metric EM. Ablated associative layers from 16 to 31.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

ARMT, MT, 2k, pre-selected layers, trained

ARMT, MT, 4k, pre-selected layers, trained

ARMT, MT, 8k, pre-selected layers, trained

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000 0.771 0.058 0.000 0.313

0.718 0.638 0.642 0.603 0.598 0.469 0.418 0.138 0.000 0.014 0.000 0.000 0.000 0.639 0.112 0.000 0.300

0.756 0.686 0.717 0.630 0.739 0.802 0.692 0.755 0.663 0.708 0.257 0.000 0.000 0.707 0.473 0.000 0.557

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030 0.714 0.643 0.403 0.668

0.731 0.648 0.642 0.562 0.543 0.469 0.429 0.372 0.202 0.146 0.129 0.078 0.119 0.626 0.227 0.097 0.369

0.744 0.667 0.698 0.658 0.728 0.704 0.692 0.702 0.629 0.625 0.624 0.617 0.485 0.698 0.631 0.555 0.655

0.744 0.667 0.698 0.644 0.728 0.765 0.758 0.734 0.640 0.653 0.723 0.696 0.535 0.696 0.684 0.621 0.688

Table 31: Associative layers ablation on the MT dataset for SmolLM-2-360M-IT model, metric - EM. Model with 6 associative blocks (approximately 20%) achieves slightly lower performance as the full ARMT model.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

ARMT, GR-100+, 2k, pre-selected layers, trained

ARMT, GR-100+, 4k, pre-selected layers, trained

ARMT, GR-100+, 8k, pre-selected layers, trained

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.335 0.289 0.238 0.220 0.212 0.177 0.156 0.137 0.074 0.046 0.011 0.025 0.015

0.332 0.297 0.257 0.245 0.215 0.202 0.199 0.195 0.190 0.140 0.059 0.032 0.011

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.345 0.261 0.230 0.204 0.180 0.157 0.159 0.134 0.157 0.156 0.165 0.170 0.165

0.346 0.257 0.242 0.226 0.219 0.186 0.202 0.186 0.202 0.203 0.169 0.164 0.214

0.329 0.252 0.248 0.236 0.220 0.199 0.199 0.199 0.213 0.208 0.178 0.166 0.176

0.327

0.258

0.269

0.275

0.244

0.258

0.257

0.074

0.106

0.170

0.211

0.154

0.193

0.200

0.072

0.023

0.027

0.154

0.169

0.176

0.168

0.192

0.178

0.216

0.241

0.196

0.223

0.227

Table 32: Associative layers ablation on the GR-100+ dataset for SmolLM-2-360M-IT model, metric - ROUGE-L. Model with 6 associative blocks (approximately 20%) achieves slightly lower performance as the full ARMT model.

Dataset MT MT MT MT GR GR GR GR GR-100+ GR-100+ GR-100+ GR-100+ ContractNLI ContractNLI ContractNLI ContractNLI

Model

Length

Learning Rate

Training Steps

Total Batch Size

Gemma-3-1B-IT ARMT ARMT ARMT Gemma-3-1B-IT ARMT ARMT ARMT Gemma-3-1B-IT ARMT ARMT ARMT Gemma-3-1B-IT ARMT ARMT ARMT

8k 2k 4k 8k 8k 2k 4k 8k 8k 2k 4k 8k 8k 2k 4k 8k

1e-4 1e-4 1e-4 3e-5 1e-4 1e-4 3e-5 1e-5 1e-4 1e-4 3e-5 1e-5 1e-4 1e-4 3e-5 1e-5

10.000 10.000 10.000 10.000 8.000 4.000 4.000 4.000 10.000 5.000 5.000 5.000 500 250 250 250

64 64 64 64 64 64 64 64 64 64 64 64 64 64 64 64

Table 33: Training hyperparameters. All models were trained with LoRA (Hu et al., 2021) on all linear layers with rank=64, α=128 and dropout=0.1. We used gradient clipping to the maximum value of 1.0 during training, and weight decay of 0.01. For ARMT model, we used 16 memory tokens and associative memory size of 64.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

ARMT, GR-100+, 2k, 32 mem

ARMT, GR-100+, 4k, 32 mem

ARMT, GR-100+, 8k, 32 mem

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.357 0.296 0.264 0.240 0.219 0.205 0.213 0.193 0.180 0.115 0.016 0.002 0.000

0.357 0.294 0.275 0.286 0.238 0.233 0.221 0.225 0.217 0.200 0.168 0.128 0.040

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.339 0.296 0.249 0.235 0.227 0.190 0.159 0.105 0.043 0.016 0.000 0.000 0.000

0.349 0.299 0.290 0.263 0.239 0.238 0.227 0.201 0.157 0.046 0.001 0.000 0.000

0.346 0.305 0.277 0.290 0.248 0.244 0.225 0.232 0.240 0.241 0.187 0.199 0.194

0.351

0.275

0.290

0.299

0.269

0.288

0.293

0.266

0.162

0.211

0.238

0.090

0.154

0.232

0.175

0.002

0.108

0.215

0.000

0.198

0.306

0.215

0.248

0.266

0.174

0.217

0.260

Table 34: Memory size ablation on the GR-100+ for Gemma-3-1B-IT model, metric - ROUGE-L. ARMT in base setup (with 16 memory tokens) performs better than model with 32 memory tokens.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, MT, 8k

ARMT, MT, 2k

ARMT, MT, 4k

ARMT, MT, 8k

ARMT, MT, 2k, 32 mem

ARMT, MT, 4k, 32 mem

ARMT, MT, 8k, 32 mem

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317

0.782 0.724 0.689 0.616 0.609 0.506 0.374 0.404 0.337 0.229 0.198 0.139 0.089

0.782 0.714 0.698 0.644 0.717 0.753 0.670 0.702 0.640 0.646 0.584 0.609 0.604

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713

0.808 0.695 0.642 0.603 0.326 0.123 0.165 0.096 0.045 0.062 0.040 0.043 0.010

0.808 0.676 0.670 0.644 0.728 0.654 0.725 0.649 0.618 0.583 0.515 0.478 0.505

0.782 0.695 0.708 0.699 0.804 0.840 0.824 0.723 0.685 0.722 0.693 0.748 0.733

0.797

0.685

0.711

0.749

0.612

0.703

0.736

0.709

0.271

0.647

0.741

0.070

0.584

0.743

0.463

0.116

0.607

0.741

0.028

0.491

0.741

0.419

0.670

0.744

0.264

0.627

0.740

Table 35: Memory size ablation on the MT for Gemma-3-1B-IT model, metric - EM. ARMT in base setup (with 16 memory tokens) performs better than model with 32 memory tokens.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, GR-100+, 8k

ARMT, GR-100+, 2k

ARMT, GR-100+, 4k

ARMT, GR-100+, 8k

RMT, GR-100+, 2k

RMT, GR-100+, 4k

RMT, GR-100+, 8k

0.380 0.385 0.352 0.344 0.296 0.306 0.269 0.280 0.231 0.269 0.247 0.171 0.189

0.357 0.296 0.264 0.240 0.219 0.205 0.213 0.193 0.180 0.115 0.016 0.002 0.000

0.357 0.294 0.275 0.286 0.238 0.233 0.221 0.225 0.217 0.200 0.168 0.128 0.040

0.358 0.303 0.279 0.303 0.254 0.235 0.241 0.237 0.240 0.252 0.207 0.200 0.264

0.218 0.198 0.132 0.138 0.122 0.104 0.119 0.119 0.114 0.106 0.101 0.106 0.094

0.242 0.224 0.168 0.170 0.165 0.152 0.157 0.127 0.117 0.103 0.098 0.114 0.016

0.196 0.263 0.220 0.214 0.226 0.167 0.190 0.171 0.154 0.141 0.117 0.097 0.144

0.351

0.275

0.290

0.299

0.161

0.193

0.224

0.266

0.162

0.211

0.238

0.111

0.127

0.159

0.175

0.002

0.108

0.215

0.103

0.091

0.108

0.306

0.215

0.248

0.266

0.135

0.158

0.189

Table 36: Comparison of RMT and ARMT on the GR-100+ dataset for Gemma-3-1B-IT model, metric - ROUGE-L. ARMT consistently outperforms RMT model.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k Full (0k-8k)

Base, No FineTuning

Base, CNLI, 8k

ARMT, CNLI, 2k

ARMT, CNLI, 4k

ARMT, CNLI, 8k

ARMT, CNLI, 2k, pretrain

ARMT, CNLI, 4k, pretrain

ARMT, CNLI, 8k, pretrain

0.427 0.089 0.046 0.012 0.196

0.703 0.690 0.700 0.806 0.922

0.615 0.672 0.512 0.347 0.255

0.678 0.691 0.680 0.747 0.765

0.710 0.701 0.730 0.776 0.882

0.689 0.671 0.542 0.359 0.255

0.675 0.724 0.703 0.747 0.784

0.703 0.737 0.727 0.765 0.824

0.119

0.710

0.579

0.692

0.721

0.599

0.714

0.734

Table 37: Results on the ContractNLI (CNLI) dataset for Gemma-3-1B-IT model with ARMT after continuous pretraining, ROUGE-L. ARMT after continuous pretraining shows significantly better results on all splits.

Task

QA1

QA2

QA3

Model/ Length

16k

32k

64k

16k

32k

64k

16k

32k

64k

Gemma-3-1B-IT Gemma-3-1B-IT, BABILong, 5k ARMT, BABILong, 2k ARMT, BABILong, 3k ARMT, BABILong, 5k

83 100 100 100 100

60 100 100 100 100

60 100 97 99 100

40 100 69 85 100

8 99 43 61 98

4 96 28 44 100

0 94 27 17 100

0 3 4 13 99

48 93 100 99 100

47 99 100 99 100

23 100 85 92 98

8 97 39 76 98

5 91 27 38 96

1 85 28 34 87

0 67 17 18 87

0 2 4 12 84

38 89 98 95 99

34 95 99 98 100

30 95 96 95 98

19 93 62 81 95

8 91 38 27 91

13 89 21 18 87

0 81 21 7 79

0 9 5 7 69

Table 38: Results on the BABILong dataset for Gemma-3-1B-IT model with ARMT, tasks QA1-QA3. ARMT outperforms base model.

Task

QA4

QA5

Avg.

Model/ Length

16k

32k

64k

16k

32k

64k

Avg. (QA1-QA5)

Gemma-3-1B-IT Gemma-3-1B-IT, BABILong, 5k ARMT, BABILong, 2k ARMT, BABILong, 3k ARMT, BABILong, 5k

30 95 100 100 100

12 100 100 100 100

11 100 92 93 90

7 99 67 77 83

6 97 44 66 82

14 81 39 62 79

1 65 42 36 79

0 5 6 10 82

70 97 99 99 99

39 100 99 99 99

54 99 98 99 99

40 98 86 92 99

31 97 55 71 98

21 95 45 36 98

7 91 36 14 97

0 25 3 2 97

22 83 58 62 94

Table 39: Results on the BABILong dataset for Gemma-3-1B-IT model with ARMT, tasks QA4-QA5. ARMT outperforms base model.

Model/ Lengths

Base, MT, 8k

ARMT, MT, 8k

ARMT, MT, 8k, pretrain

Mamba-1.4B, MT, 8k

Mamba-2-1.3B, MT, 8k

DeltaNet-1.3B, MT, 8k

xLSTM-1.4B, MT, 8k

0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

0.795 0.790 0.774 0.767 0.859 0.852 0.868 0.787 0.764 0.812 0.713 0.591 0.317 0.797 0.709 0.463 0.741

0.756 0.743 0.736 0.699 0.804 0.753 0.791 0.745 0.685 0.757 0.713 0.765 0.713 0.749 0.741 0.741 0.744

0.756 0.743 0.736 0.740 0.859 0.827 0.868 0.766 0.697 0.757 0.812 0.791 0.762 0.767 0.783 0.777 0.777

0.846 0.857 0.802 0.836 0.902 0.877 0.879 0.777 0.764 0.882 0.802 0.783 0.525 0.848 0.788 0.662 0.810

0.795 0.886 0.802 0.781 0.902 0.827 0.890 0.766 0.764 0.833 0.762 0.730 0.644 0.837 0.777 0.690 0.798

0.769 0.819 0.726 0.726 0.891 0.864 0.846 0.755 0.708 0.757 0.673 0.565 0.426 0.788 0.694 0.500 0.727

0.769 0.800 0.755 0.822 0.870 0.864 0.890 0.819 0.742 0.819 0.743 0.530 0.426 0.802 0.724 0.481 0.752

Table 40: Results on the MT dataset for Gemma-3-1B-IT model with ARMT and for various baselines, EM. ARMT shows better length generalization than Mamba and DeltaNet and do not require training from scratch.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, GR-100+, 8k

Base, YaRN, GR-100+, 8k

ARMT, GR-100+, 8k

0.145 0.122 0.113 0.091 0.094 0.089 0.039 0.039 0.039 0.048 0.042 0.044 0.088

0.349 0.358 0.330 0.314 0.286 0.233 0.021 0.035 0.049 0.037 0.054 0.067 0.088

0.362 0.357 0.347 0.312 0.289 0.239 0.177 0.067 0.024 0.017 0.003 0.000 0.000

0.337 0.302 0.261 0.253 0.225 0.221 0.211 0.215 0.210 0.219 0.185 0.171 0.096

0.113

0.327

0.334

0.275

0.050

0.074

0.091

0.211

0.054

0.072

0.000

0.154

0.080

0.192

0.205

0.241

Table 41: Results on the GR-100+ dataset for SmolLM-2-360M-IT model with YaRN, ROUGE-L. Model with YaRN shows worse performance than the ARMT model.

Model/ Lengths 0k-1k 1k-2k 2k-4k 4k-6k 6k-8k 8k-10k 10k-12k 12k-14k 14k-16k 16k-24k 24k-32k 32k-49k 49k-65k In-Domain (0k-8k) OOD (8k-65k) Long-OOD (32k-65k) Full (0k-65k)

Base, No FineTuning

Base, MT, 8k

Base, YaRN, MT, 8k

ARMT, MT, 8k

0.000 0.000 0.000 0.000 0.000 0.000 0.011 0.000 0.000 0.000 0.020 0.009 0.000

0.795 0.790 0.736 0.685 0.837 0.568 0.000 0.000 0.000 0.000 0.010 0.000 0.000

0.333 0.257 0.264 0.301 0.250 0.272 0.176 0.106 0.000 0.000 0.000 0.000 0.000

0.769 0.705 0.698 0.630 0.761 0.815 0.747 0.755 0.629 0.715 0.733 0.730 0.030

0.000

0.771

0.277

0.714

0.005

0.058

0.059

0.643

0.005

0.000

0.403

0.003

0.313

0.137

0.668

Table 42: Results on the MT dataset for SmolLM-2-360M-IT model, metric - EM. Model with YaRN shows worse performance than the base model and the ARMT model.

Record · ID 363297 · SHA-256 90f087cfe00bd9af
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.