Paper deep dive

Extending LLM Context via Associative Recurrent Memory

Gleb Kuzmin, Ivan Rodkin, Aydar Bulatov, Yuri Kuratov, Lyudmila Rvanova, Mikhail Katkov, Ilia Sochenkov, Misha Tsodyks, Timothy Baldwin, Mikhail Burtsev, Artem Shelmanov

Year: 2026Venue: arXiv preprintArea: cs.CLType: PreprintEmbeddings: 242

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%

Last extracted: 7/18/2026, 2:49:37 PM

Summary

This paper investigates the Associative Recurrent Memory Transformer (ARMT) as a method to extend the context length of Large Language Models (LLMs) beyond standard transformer limitations. The authors propose a training recipe involving continued pre-training, synthetic data generation, curriculum learning, and selective layer integration. Experiments on Gemma-3-1B-IT and SmolLM-2-360M-IT demonstrate that ARMT-augmented models can process inputs well beyond original context limits, generalize to out-of-distribution lengths, and reduce FLOPs by 30% while maintaining performance. Two new domain-specific datasets, ManyTypes-long and GovReport-long, are introduced for evaluation.

Entities (8)

ARMT · model-architecture · 98%Gemma 3 1B-IT · model · 95%SmolLM-2-360M-IT · model · 95%GovReport-long · dataset · 92%ManyTypes-long · dataset · 92%Curriculum Learning · training-technique · 90%Synthetic Data Generation · training-technique · 90%FLOPs · metric · 85%

Relation Signals (7)

ARMT → improves → Context Length

confidence 95% · ARMT-augmented models: (i) process inputs well beyond their original context limits without degrading performance

ARMT → uses → Associative Memory

confidence 95% · At the core of ARMT lies the associative block, a layerwise memory module that retrieves and updates key-value associations

Gemma 3 1B-IT → isbasemodelfor → ARMT

confidence 90% · We experimented with Gemma-3-1B-IT ... as base LLMs ... ARMT-augmented models

Curriculum Learning → ispartof → ARMT Training Recipe

confidence 90% · We propose a comprehensive training recipe for ARMT-based context extension, combining ... curriculum learning

ARMT → reduces → FLOPs

confidence 90% · need 30% less FLOPs while preserving baseline performance within the original context window.

ManyTypes-long → isusedfor → Evaluation

confidence 85% · We construct two new datasets ... ManyTypes-long (MT) targets variable type prediction

GovReport-long → isusedfor → Evaluation

confidence 85% · GovReport-long (GR) focuses on long-document question answering.

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract: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.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

240,694 characters extracted from source content.

Expand or collapse full text

Extending LLM Context via Associative Recurrent Memory 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 1FusionBrain Lab 2MBZUAI 3Cognitive AI Systems Lab 4RUDN 5London Institute for Mathematical Sciences 6MIRAI 7Lomonosov Moscow State University 8Laboratory for Analysis and Controllable Text Generation Technologies RAS 9School of Natural Sciences, Institute for Advanced Study, Princeton 10Department of Brain Sciences, Weizmann Institute of Science kuzmin.gyu@gmail.com artem.shelmanov@mbzuai.ac.ae Abstract 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; (i) generalize more effectively to out-of-distribution context lengths; and (i) need 30% less FLOPs while preserving baseline performance within the original context window. Extending LLM Context via Associative Recurrent Memory 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 1FusionBrain Lab 2MBZUAI 3Cognitive AI Systems Lab 4RUDN 5London Institute for Mathematical Sciences 6MIRAI 7Lomonosov Moscow State University 8Laboratory for Analysis and Controllable Text Generation Technologies RAS 9School of Natural Sciences, Institute for Advanced Study, Princeton 10Department of Brain Sciences, Weizmann Institute of Science kuzmin.gyu@gmail.com artem.shelmanov@mbzuai.ac.ae 1 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., 2024). Therefore, since the introduction of the 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 segment-level 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 middle ground, enabling linear-compute, constant-memory long-context processing while maintaining strong performance on real-world tasks. 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. 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. 2 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. Fully-recurrent models exhibit strong performance on 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 ability of recurrent architectures to propagate information across arbitrarily long contexts (in principle) and the strong performance of full self-attention 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). 3 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 language-model 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 performs 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 attention 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. 4 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 long-context 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 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. ((a)) GPU memory usage. ((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. 5 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 self-attention. 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-3-1B-IT (Team et al., 2025) and SmolLM-2-360M-IT (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. SmolLM-2-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 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 pre-training on a synthetic long-context QA task. Training hyperparameters are provided in Appendix E. 5.2 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 token 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/ST/S; for T=32,768T=32,768 and S=1,024S=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). 5.3 Analysis of Model Performance ((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 Base, No Fine- Tuning Base, MT, 8k ARMT, MT, 8k ARMT, MT, 8k pre-selected layers, trained ARMT, MT, 8k, pretrain 0k-1k [HTML]F0000 0.000 [HTML]2A9500 0.795 [HTML]42A100 0.756 [HTML]42A100 0.756 [HTML]42A100 0.756 1k-2k [HTML]F0000 0.000 [HTML]2E9700 0.790 [HTML]48A400 0.743 [HTML]60B000 0.705 [HTML]48A400 0.743 2k-4k [HTML]F0000 0.000 [HTML]369B00 0.774 [HTML]4CA600 0.736 [HTML]58AC00 0.717 [HTML]4CA600 0.736 4k-6k [HTML]F0000 0.000 [HTML]3A9D00 0.767 [HTML]62B100 0.699 [HTML]5CAE00 0.712 [HTML]4A500 0.740 6k-8k [HTML]F0000 0.000 [HTML]048200 0.859 [HTML]249200 0.804 [HTML]128900 0.837 [HTML]048200 0.859 8k-10k [HTML]F0000 0.000 [HTML]088400 0.852 [HTML]42A100 0.753 [HTML]108800 0.840 [HTML]188C00 0.827 10k-12k [HTML]F0600 0.011 [HTML]008000 0.868 [HTML]2C9600 0.791 [HTML]068300 0.857 [HTML]008000 0.868 12k-14k [HTML]F0000 0.000 [HTML]2E9700 0.787 [HTML]48A400 0.745 [HTML]349A00 0.777 [HTML]3C9E00 0.766 14k-16k [HTML]F0000 0.000 [HTML]3C9E00 0.764 [HTML]6AB500 0.685 [HTML]6AB500 0.685 [HTML]64B200 0.697 16k-24k [HTML]F0000 0.000 [HTML]209000 0.812 [HTML]40A000 0.757 [HTML]209000 0.812 [HTML]40A000 0.757 24k-32k [HTML]F0000 0.000 [HTML]5AAD00 0.713 [HTML]5AAD00 0.713 [HTML]3E9F00 0.762 [HTML]209000 0.812 32k-49k [HTML]F0000 0.000 [HTML]A2D100 0.591 [HTML]3C9E00 0.765 [HTML]289400 0.800 [HTML]2C9600 0.791 49k-65k [HTML]F0000 0.000 [HTML]FFBA00 0.317 [HTML]5AAD00 0.713 [HTML]269300 0.802 [HTML]3E9F00 0.762 In-Domain (0k-8k) [HTML]F0000 0.000 [HTML]289400 0.797 [HTML]46A300 0.749 [HTML]48A400 0.744 [HTML]3A9D00 0.767 OOD (8k-65k) [HTML]F0000 0.001 [HTML]5CAE00 0.709 [HTML]4A500 0.741 [HTML]2C9600 0.793 [HTML]329900 0.783 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]EEF700 0.463 [HTML]4A500 0.741 [HTML]269300 0.801 [HTML]349A00 0.777 Full (0k-65k) [HTML]F0000 0.001 [HTML]4A500 0.741 [HTML]48A400 0.744 [HTML]369B00 0.776 [HTML]349A00 0.777 Table 1: Best results on the MT dataset for Gemma-3-1B-IT model, metric = EM. ARMT performs comparably with the base model overall, and outperforms it on OOD and Long-OOD. Model/ Lengths Base, No Fine- Tuning Base, GR-100+, 8k ARMT, GR-100+, 8k ARMT, GR-100+, 8k pre-selected layers, trained ARMT, GR-100+, 8k, pretrain 0k-1k [HTML]FFB600 0.188 [HTML]088400 0.380 [HTML]2C9600 0.358 [HTML]249200 0.363 [HTML]188C00 0.370 1k-2k [HTML]F8A00 0.161 [HTML]008000 0.385 [HTML]88C400 0.303 [HTML]70B800 0.317 [HTML]72B900 0.316 2k-4k [HTML]F6600 0.140 [HTML]369B00 0.352 [HTML]B0D800 0.279 [HTML]86C300 0.304 [HTML]98C00 0.293 4k-6k [HTML]F3C00 0.114 [HTML]44A200 0.344 [HTML]88C400 0.303 [HTML]A2D100 0.287 [HTML]88C400 0.303 6k-8k [HTML]F3A00 0.113 [HTML]94CA00 0.296 [HTML]DAED00 0.254 [HTML]E0F000 0.250 [HTML]D0E800 0.260 8k-10k [HTML]F2E00 0.106 [HTML]82C100 0.306 [HTML]FAFD00 0.235 [HTML]FFFE00 0.231 [HTML]E8F400 0.245 10k-12k [HTML]F2600 0.101 [HTML]C0E000 0.269 [HTML]F0F800 0.241 [HTML]FFFE00 0.231 [HTML]D0E800 0.260 12k-14k [HTML]F1400 0.091 [HTML]AED700 0.280 [HTML]F6FB00 0.237 [HTML]FFEE00 0.221 [HTML]DCEE00 0.252 14k-16k [HTML]F1000 0.088 [HTML]FFFE00 0.231 [HTML]F0F800 0.240 [HTML]ECF600 0.243 [HTML]CCE600 0.262 16k-24k [HTML]F0800 0.083 [HTML]C0E000 0.269 [HTML]DCEE00 0.252 [HTML]D2E900 0.259 [HTML]D0E800 0.260 24k-32k [HTML]F0000 0.078 [HTML]E6F300 0.247 [HTML]FFD600 0.207 [HTML]FFD200 0.204 [HTML]FFEE00 0.221 32k-49k [HTML]F0000 0.078 [HTML]F9A00 0.171 [HTML]FFCA00 0.200 [HTML]FFD200 0.205 [HTML]FFD600 0.207 49k-65k [HTML]F2800 0.103 [HTML]FFB800 0.189 [HTML]C8E400 0.264 [HTML]E6F300 0.246 [HTML]8C600 0.300 In-Domain (0k-8k) [HTML]F6C00 0.143 [HTML]389C00 0.351 [HTML]8EC700 0.299 [HTML]86C300 0.304 [HTML]80C000 0.308 OOD (8k-65k) [HTML]F1600 0.092 [HTML]C6E300 0.266 [HTML]F4FA00 0.238 [HTML]FCFE00 0.234 [HTML]DCEE00 0.252 Long-OOD (32k-65k) [HTML]F0800 0.084 [HTML]FFA200 0.175 [HTML]FFE400 0.215 [HTML]FFE200 0.214 [HTML]FFFA00 0.228 Full (0k-65k) [HTML]F3E00 0.116 [HTML]84C200 0.306 [HTML]C4E200 0.266 [HTML]C4E200 0.267 [HTML]B0D800 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. Extending the context window of base LLM. The main results with Gemma-3 on the MT dataset are shown in Figure˜3(a) 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 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˜3(b), 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 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. 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 [HTML]42A100 0.795 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 1k-2k [HTML]48A400 0.790 [HTML]74BA00 0.743 [HTML]84C200 0.724 [HTML]8EC700 0.714 [HTML]96CB00 0.705 2k-4k [HTML]56AB00 0.774 [HTML]7ABD00 0.736 [HTML]94CA00 0.708 [HTML]7ABD00 0.736 [HTML]8C600 0.717 4k-6k [HTML]5CAE00 0.767 [HTML]9CCE00 0.699 [HTML]90C800 0.712 [HTML]9CCE00 0.699 [HTML]90C800 0.712 6k-8k [HTML]088400 0.859 [HTML]3A9D00 0.804 [HTML]44A200 0.793 [HTML]309800 0.815 [HTML]1C8E00 0.837 8k-10k [HTML]0E8700 0.852 [HTML]6AB500 0.753 [HTML]5EAF00 0.765 [HTML]269300 0.827 [HTML]1A8D00 0.840 10k-12k [HTML]008000 0.868 [HTML]46A300 0.791 [HTML]66B300 0.758 [HTML]3C9E00 0.802 [HTML]0A8500 0.857 12k-14k [HTML]4A500 0.787 [HTML]72B900 0.745 [HTML]68B400 0.755 [HTML]54A00 0.777 [HTML]54A00 0.777 14k-16k [HTML]60B000 0.764 [HTML]AAD500 0.685 [HTML]BEDF00 0.663 [HTML]94CA00 0.708 [HTML]AAD500 0.685 16k-24k [HTML]349A00 0.812 [HTML]66B300 0.757 [HTML]80C000 0.729 [HTML]66B300 0.757 [HTML]349A00 0.812 24k-32k [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]A2D100 0.693 [HTML]62B100 0.762 [HTML]62B100 0.762 32k-49k [HTML]FFFE00 0.591 [HTML]5EAF00 0.765 [HTML]5EAF00 0.765 [HTML]5EAF00 0.765 [HTML]3E9F00 0.800 49k-65k [HTML]F0000 0.317 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]6AB500 0.752 [HTML]3C9E00 0.802 In-Domain (0k-8k) [HTML]40A000 0.797 [HTML]6EB700 0.749 [HTML]78BC00 0.738 [HTML]72B900 0.744 [HTML]72B900 0.744 OOD (8k-65k) [HTML]92C900 0.709 [HTML]74BA00 0.741 [HTML]80C000 0.730 [HTML]5CAE00 0.767 [HTML]44A200 0.793 Long-OOD (32k-65k) [HTML]F8600 0.463 [HTML]76B00 0.741 [HTML]76B00 0.741 [HTML]64B200 0.759 [HTML]3E9F00 0.801 Full (0k-65k) [HTML]76B00 0.741 [HTML]72B900 0.744 [HTML]7CBE00 0.733 [HTML]64B200 0.759 [HTML]54A00 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 Model Num. of Trainable Parameters GPU RAM Usage ↓ Steps per Second ↑ 2 ARMT 90.5M 45.4GB 2.86 ARMT, pre-selected layers 59.6M 44.5GB 3.64 4 ARMT 90.5M 71.8GB 1.62 ARMT, pre-selected layers 59.6M 69.8GB 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. 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 long-context data consistently improves performance up 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 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. 6 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, (i) synthetic long-context data generation, (i) 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 privacy-preserving 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%. Limitations 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 (i) 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. Ethical Considerations 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 long-context 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. References L. B. Allal, A. Lozhkov, E. Bakouch, G. M. Blázquez, G. Penedo, L. Tunstall, A. Marafioti, H. Kydlícek, A. P. Lajarín, V. Srivastav, J. Lochner, C. Fahlgren, X. Nguyen, C. Fourrier, B. Burtenshaw, H. Larcher, H. Zhao, C. Zakka, M. Morlon, C. Raffel, L. von Werra, and T. Wolf (2025) SmolLM2: when Smol goes big - data-centric training of a small language model. arXiv preprint arXiv:2502.02737. External Links: Link, Document, 2502.02737 Cited by: §5.1. Anthropic (2024) Introducing the next generation of Claude. Note: urlhttps://w.anthropic.com/news/claude-3-family Cited by: §2. Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, et al. (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), p. 3119–3137. External Links: Link Cited by: Appendix I. M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter (2024) xLSTM: extended long short-term memory. In Proceedings of the 38th International Conference on Neural Information Processing Systems, p. 107547–107603. External Links: Link Cited by: §2, §5.3. I. Beltagy, M. E. Peters, and A. Cohan (2020) Longformer: the long-document transformer. External Links: 2004.05150, Link Cited by: §1. Y. Bengio, J. Louradour, R. Collobert, and J. Weston (2009) Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning - ICML ’09, Montreal, Quebec, Canada, p. 1–8 (en). External Links: ISBN 978-1-60558-516-1, Document, Link Cited by: §3. A. Bulatov, Y. Kuratov, Y. Kapushev, and M. Burtsev (2024) Beyond attention: breaking the limits of transformer context length with recurrent memory. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 38, p. 17700–17708. External Links: Link Cited by: §1. A. Bulatov, Y. Kuratov, and M. Burtsev (2022) Recurrent memory transformer. In Advances in Neural Information Processing Systems, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (Eds.), Vol. 35, p. 11079–11091. External Links: Link Cited by: §1, §5.3. S. Cao and L. 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), S. Muresan, P. Nakov, and A. Villavicencio (Eds.), Dublin, Ireland, p. 786–807. External Links: Link, Document Cited by: §4. R. Child, S. Gray, A. Radford, and I. Sutskever (2019) Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. External Links: Link Cited by: §2. N. Cowan (2008) What are the differences between long-term, short-term, and working memory?. Progress in Brain Research 169, p. 323–338. External Links: Link Cited by: §1. Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. Le, and R. 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, p. 2978–2988. External Links: Link Cited by: §2. T. Dao and A. 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, p. 10041–10071. External Links: Link Cited by: §5.3. Y. Fu, X. Dong, S. Diao, H. Ye, W. Byeon, Y. Karnati, L. Liebenwein, M. Khadkevich, A. Keller, J. Kautz, et al. (2025) Nemotron-Flash: towards latency-optimal hybrid small language models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, External Links: Link Cited by: §2. T. Gao, A. Wettig, H. Yen, and D. 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), p. 7376–7399. External Links: Link Cited by: §4. A. Gu and T. Dao (2024) Mamba: linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, External Links: Link Cited by: §1, §2, §2, §5.3. C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, and B. Ginsburg (2024) RULER: what’s the real context size of your long-context language models?. In First Conference on Language Modeling, External Links: Link Cited by: Appendix I. E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, and W. Chen (2021) LoRA: low-rank adaptation of large language models. CoRR abs/2106.09685. External Links: Link, 2106.09685 Cited by: Table 33, §5.1. S. Jelassi, D. Brandfonbrener, S. M. Kakade, and E. Malach (2024) Repeat after me: transformers are better than state space models at copying. In Proceedings of the 41st International Conference on Machine Learning, p. 21502–21521. External Links: Link Cited by: §1, §2. A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret (2020) Transformers are RNNs: fast autoregressive transformers with linear attention. In International Conference on Machine Learning, p. 5156–5165. External Links: Link Cited by: §1, §2. Y. Koreeda and C. Manning (2021) ContractNLI: a dataset for document-level natural language inference for contracts. In Findings of the Association for Computational Linguistics: EMNLP 2021, Punta Cana, Dominican Republic, p. 1907–1919. External Links: Link, Document Cited by: Appendix H. Y. Kuratov, A. Bulatov, P. Anokhin, I. Rodkin, D. Sorokin, A. Sorokin, and M. Burtsev (2024) BABILong: testing the limits of LLMs with long context reasoning-in-a-haystack. Advances in Neural Information Processing Systems 37, p. 106519–106554. External Links: Link Cited by: Appendix I, §1, §2, §2, §5.3. B. Lenz, O. Lieber, A. Arazi, A. Bergman, A. Manevich, B. Peleg, B. Aviram, C. Almagor, C. Fridman, D. Padnos, et al. (2025) Jamba: hybrid Transformer-Mamba language models. In International Conference on Learning Representations, Vol. 2025, p. 67959–67984. External Links: Link Cited by: §2. C. Lin (2004) ROUGE: a package for automatic evaluation of summaries. In Text Summarization Branches Out, Barcelona, Spain, p. 74–81. External Links: Link Cited by: §5.1. N. F. Liu, K. Lin, J. Hewitt, A. Paranjape, M. Bevilacqua, F. Petroni, and P. Liang (2024) Lost in the middle: how language models use long contexts. Transactions of the Association for Computational Linguistics 12, p. 157–173. External Links: Link Cited by: §1. W. Merrill, J. Petty, and A. Sabharwal (2024) The illusion of state in state-space models. In International Conference on Machine Learning, p. 35492–35506. External Links: Link Cited by: §2. 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), p. 585–589. External Links: Link Cited by: §4, §4. OpenAI (2023) New models and developer products announced at DevDay. Note: https://openai.com/index/new-models-and-developer-products-announced-at-devday/ Cited by: §2. J. Park, J. Park, Z. Xiong, N. Lee, J. Cho, S. Oymak, K. Lee, and D. Papailiopoulos (2024) Can Mamba learn how to learn? a comparative study on in-context learning tasks. In Forty-first International Conference on Machine Learning, External Links: Link Cited by: §1, §2. G. Penedo, H. Kydlíček, L. B. Allal, A. Lozhkov, M. Mitchell, C. Raffel, L. Von Werra, and T. 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. External Links: ISBN 9798331314385, Link Cited by: §5.1, §5.3. B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, X. Du, M. Grella, K. Gv, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. Wind, S. Woźniak, Z. Zhang, Q. Zhou, J. Zhu, and R. Zhu (2023) RWKV: reinventing RNNs for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, H. Bouamor, J. Pino, and K. Bali (Eds.), Singapore, p. 14048–14077. External Links: Link Cited by: §1, §2, §2. B. Peng, J. Quesnelle, H. Fan, and E. Shippole (2024) YaRN: efficient context window extension of large language models. In International Conference on Learning Representations, B. Kim, Y. Yue, S. Chaudhuri, K. Fragkiadaki, M. Khan, and Y. Sun (Eds.), Vol. 2024, p. 31932–31951. External Links: Link Cited by: Appendix J. Qwen Team (2026) Qwen3.5: accelerating productivity with native multimodal agents. Note: https://qwen.ai/blog?id=qwen3.5. Cited by: §2. J. W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap (2020) Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations, External Links: Link Cited by: Appendix J, §2. S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He (2020) Zero: memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, p. 1–16. External Links: Link Cited by: §3. J. Rasley, S. Rajbhandari, O. Ruwase, and Y. 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, p. 3505–3506. External Links: Link Cited by: §3. M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al. (2024) Gemini 1.5: unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530. External Links: Link Cited by: §2. I. Rodkin, Y. Kuratov, A. Bulatov, and M. Burtsev (2024) Associative recurrent memory transformer. arXiv preprint arXiv:2407.04841. External Links: Link Cited by: Appendix G, §1, §1, §2, §2. U. Shaham, E. Segal, M. Ivgi, A. Efrat, O. Yoran, A. Haviv, A. Gupta, W. Xiong, M. Geva, J. Berant, and O. Levy (2022) SCROLLS: standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, Abu Dhabi, United Arab Emirates, p. 12007–12021. External Links: Link Cited by: Appendix H, §5.3. D. Sivtsov, I. Rodkin, G. Kuzmin, Y. Kuratov, and I. Oseledets (2025) Diagonal batching unlocks parallelism in recurrent memory transformers for long contexts. arXiv preprint arXiv:2506.05229. External Links: Link Cited by: §5.2. G. Team, A. Kamath, J. Ferret, S. Pathak, N. Vieillard, R. Merhej, S. Perrin, T. Matejovicova, A. Ramé, M. Rivière, et al. (2025) Gemma 3 technical report. arXiv preprint arXiv:2503.19786. External Links: Link Cited by: §3, §5.1, §5.1. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Advances in Neural Information Processing Systems 30. External Links: Link Cited by: §1. S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma (2020) Linformer: self-attention with linear complexity. External Links: 2006.04768 Cited by: §2. P. Xu, W. Ping, X. Wu, C. Xu, Z. Liu, M. Shoeybi, and B. Catanzaro (2024) ChatQA 2: bridging the gap to proprietary LLMs in long context and RAG capabilities. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §4. S. Yang, J. Kautz, and A. Hatamizadeh (2024a) Gated delta networks: improving Mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations, External Links: Link Cited by: §2. S. Yang, B. Wang, Y. Zhang, Y. Shen, and Y. Kim (2024b) Parallelizing linear transformers with the delta rule over sequence length. Advances in Neural Information Processing Systems 37, p. 115491–115522. External Links: Link Cited by: §5.3. M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, and A. Ahmed (2020) Big Bird: transformers for longer sequences. In Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (Eds.), Vol. 33, p. 17283–17297. External Links: Link Cited by: §2. Appendix A Formal Description of the ARMT Architecture For the hidden states of segment s in layer l HslH_s^l, memory tokens from the previous layer Msl−1M_s^l-1, the associative matrix AslA_s^l, and normalization zslz_s^l are updated as follows: Msl=mi,Msl=TransLayer(Hsl−1,Msl−1) M_s^l=\m_i\, M_s^l=TransLayer(H^l-1_s,M_s^l-1) (1) ki,vi=WKmi,WVmi;βi=σ(Wβmi); k_i,v_i=W_Km_i,W_Vm_i; _i=σ(W_βm_i); (2) A0l=0→;z0l=0→; A_0^l= 0; z_0^l= 0; (3) v¯i=As−1lϕ(ki)(zs−1)Tϕ(ki);γi=1−(zs−1)Tϕ(ki)‖ϕ(ki)‖2; v_i= A_s-1^lφ(k_i)(z_s-1)^Tφ(k_i); _i=1- (z_s-1)^Tφ(k_i)\|φ(k_i)\|^2; (4) Asl=As−1l+∑iβi(vi−v¯i)⊗ϕ(ki); A_s^l=A_s-1^l+ _i _i(v_i- v_i) φ(k_i); (5) zsl=zs−1l+∑iγiϕ(ki). z^l_s=z^l_s-1+ _i _iφ(k_i). (6) Reading from memory in the following segments for embedding xjx_j from Hs+1l−1H_s+1^l-1: qj=WQxj;yj=Aslϕ(qj)(zsl)Tϕ(qj), q_j=W_Qx_j; y_j= A_s^lφ(q_j)(z^l_s)^Tφ(q_j), (7) where yjy_j is an association for xjx_j. A.1 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, NgN_g global-attention layers, NlN_l local-attention layers, hidden dimension d, head dimension dhd_h, and H attention heads. Gemma-3. For a full sequence of length T, the attention cost per layer is 4HT2dh4HT^2d_h for global attention and 4HTWdh4HTWd_h for local attention with window size W. For long sequences, the sequence-length-dependent attention cost is dominated by global attention, giving FGemma_Attn∼4NgHdhT2.F_Gemma\_Attn 4N_gHd_hT^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 FFFN=(2×2×6+2×3)NTd2=30NTd2.F_FFN=(2× 2× 6+2× 3)NTd^2=30NTd^2. This term is independent of the segmentation and is therefore approximately equal for Gemma and ARMT when the memory-token overhead is neglected. ARMT. ARMT splits the sequence into TS TS segments. Each segment is processed with attention context length S+MS+M, where M is the number of memory tokens. Therefore, the global-attention FLOP cost becomes FARMT_Attn∼TS4NgHdh(S+M)2.F_ARMT\_Attn TS4N_gHd_h(S+M)^2. When M≪SM S, this simplifies to FARMT_Attn≈4NgHdhST.F_ARMT\_Attn≈ 4N_gHd_hST. Thus, for global attention alone, the reduction factor is approximately FGemma_AttnFARMT_Attn≈TS. F_Gemma\_AttnF_ARMT\_Attn≈ TS. For T=32,768T=32,768 and S=1,024S=1,024, this gives TS=32, TS=32, i.e., ARMT achieves a ×32× reduction in global-attention FLOPs when memory-token overhead is neglected. 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 FGemma≈4NgHdhT2+30NTd2,F_Gemma≈ 4N_gHd_hT^2+30NTd^2, and FARMT≈4NgHdhST+30NTd2.F_ARMT≈ 4N_gHd_hST+30NTd^2. Dividing both numerator and denominator by 2T2T gives FARMTFGemma≈2NgHdhS+15Nd22NgHdhT+15Nd2. F_ARMTF_Gemma≈ 2N_gHd_hS+15Nd^22N_gHd_hT+15Nd^2. For Gemma-3-1B-IT, substituting the model constants into this expression gives an overall FLOP ratio of approximately 0.670.67. Thus, although ARMT reduces global-attention FLOPs by 32×32× for T=32,768T=32,768 and S=1,024S=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. Appendix B Datasets Statistics and Ablations B.1 Dataset Statistics and Examples 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. Dataset/ Split Train Validation Test MT, 2k 229.4k 3.1k 1.3k MT, 4k 289.2k 4.5k 1.3k MT, 8k 328.3k 5.1k 1.3k GR, 2k 9.9k 0.5k 1.0k GR, 4k 17.4k 0.9k 1.0k GR, 8k 19.7k 1.0k 1.0k GR-100+, 2k 90.5k 0.5k 1.0k GR-100+, 4k 121.7k 0.9k 1.0k GR-100+, 8k 128.4k 1.0k 1.0k ContractNLI (SCROLLS) 7.2k 1.0k 2.1k Table 5: Dataset statistics - number of samples for each dataset and for each split. 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 What is the type of variable down_revision? Answer 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. Question What makes cost-estimates of IP infringements difficult to calculate? Answer Generally, as GAO reported in April 2010, the illicit nature of counterfeiting and piracy makes estimating the economic impact of IP infringements extremely difficult. Table 7: Example from GR dataset. The question during training and evaluation is placed before context (i. e. at the start of the prompt). Appendix C 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. Model/ Lengths Base, No Fine- Tuning Base, MT, 8k ARMT, MT, 2k ARMT, MT, 4k ARMT, MT, 8k 0k-1k [HTML]F0000 0.000 [HTML]2A9500 0.795 [HTML]329900 0.782 [HTML]329900 0.782 [HTML]42A100 0.756 1k-2k [HTML]F0000 0.000 [HTML]2E9700 0.790 [HTML]54A00 0.724 [HTML]5AAD00 0.714 [HTML]48A400 0.743 2k-4k [HTML]F0000 0.000 [HTML]369B00 0.774 [HTML]68B400 0.689 [HTML]64B200 0.698 [HTML]4CA600 0.736 4k-6k [HTML]F0000 0.000 [HTML]3A9D00 0.767 [HTML]94CA00 0.616 [HTML]84C200 0.644 [HTML]62B100 0.699 6k-8k [HTML]F0000 0.000 [HTML]048200 0.859 [HTML]98C00 0.609 [HTML]58AC00 0.717 [HTML]249200 0.804 8k-10k [HTML]F0000 0.000 [HTML]088400 0.852 [HTML]D4EA00 0.506 [HTML]42A100 0.753 [HTML]42A100 0.753 10k-12k [HTML]F0600 0.011 [HTML]008000 0.868 [HTML]FFDC00 0.374 [HTML]74BA00 0.670 [HTML]2C9600 0.791 12k-14k [HTML]F0000 0.000 [HTML]2E9700 0.787 [HTML]FFEE00 0.404 [HTML]60B000 0.702 [HTML]48A400 0.745 14k-16k [HTML]F0000 0.000 [HTML]3C9E00 0.764 [HTML]FFC600 0.337 [HTML]86C300 0.640 [HTML]6AB500 0.685 16k-24k [HTML]F0000 0.000 [HTML]209000 0.812 [HTML]F8600 0.229 [HTML]82C100 0.646 [HTML]40A000 0.757 24k-32k [HTML]F0000 0.000 [HTML]5AAD00 0.713 [HTML]F7400 0.198 [HTML]A6D300 0.584 [HTML]5AAD00 0.713 32k-49k [HTML]F0000 0.000 [HTML]A2D100 0.591 [HTML]F5000 0.139 [HTML]98C00 0.609 [HTML]3C9E00 0.765 49k-65k [HTML]F0000 0.000 [HTML]FFBA00 0.317 [HTML]F3400 0.089 [HTML]9ACD00 0.604 [HTML]5AAD00 0.713 In-Domain (0k-8k) [HTML]F0000 0.000 [HTML]289400 0.797 [HTML]6AB500 0.685 [HTML]5CAE00 0.711 [HTML]46A300 0.749 OOD (8k-65k) [HTML]F0000 0.001 [HTML]5CAE00 0.709 [HTML]F9E00 0.271 [HTML]82C100 0.647 [HTML]4A500 0.741 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]EEF700 0.463 [HTML]F4400 0.116 [HTML]9ACD00 0.607 [HTML]4A500 0.741 Full (0k-65k) [HTML]F0000 0.001 [HTML]4A500 0.741 [HTML]F600 0.419 [HTML]74BA00 0.670 [HTML]48A400 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 Base, No Fine- Tuning Base, GR-100+, 8k ARMT, GR-100+, 2k ARMT, GR-100+, 4k ARMT, GR-100+, 8k 0k-1k [HTML]FFFA00 0.188 [HTML]068300 0.380 [HTML]249200 0.357 [HTML]249200 0.357 [HTML]229100 0.358 1k-2k [HTML]FFD600 0.161 [HTML]008000 0.385 [HTML]76B00 0.296 [HTML]78BC00 0.294 [HTML]6CB600 0.303 2k-4k [HTML]FFBA00 0.140 [HTML]2A9500 0.352 [HTML]A0D000 0.264 [HTML]92C900 0.275 [HTML]8C600 0.279 4k-6k [HTML]F9600 0.114 [HTML]369B00 0.344 [HTML]C0E000 0.240 [HTML]82C100 0.286 [HTML]6CB600 0.303 6k-8k [HTML]F9600 0.113 [HTML]76B00 0.296 [HTML]DCEE00 0.219 [HTML]C2E100 0.238 [HTML]AED700 0.254 8k-10k [HTML]F8C00 0.106 [HTML]68B400 0.306 [HTML]EEF700 0.205 [HTML]CAE500 0.233 [HTML]C6E300 0.235 10k-12k [HTML]F8600 0.101 [HTML]9ACD00 0.269 [HTML]E4F200 0.213 [HTML]DAED00 0.221 [HTML]BEDF00 0.241 12k-14k [HTML]F7800 0.091 [HTML]8AC500 0.280 [HTML]FEFF00 0.193 [HTML]D4EA00 0.225 [HTML]C4E200 0.237 14k-16k [HTML]F7400 0.088 [HTML]CCE600 0.231 [HTML]FFEE00 0.180 [HTML]DEEF00 0.217 [HTML]C0E000 0.240 16k-24k [HTML]F6E00 0.083 [HTML]9ACD00 0.269 [HTML]F9800 0.115 [HTML]F6FB00 0.200 [HTML]B0D800 0.252 24k-32k [HTML]F6600 0.078 [HTML]B6DB00 0.247 [HTML]F1400 0.016 [HTML]FFDE00 0.168 [HTML]ECF600 0.207 32k-49k [HTML]F6600 0.078 [HTML]FFE200 0.171 [HTML]F0200 0.002 [HTML]FFAA00 0.128 [HTML]F6FB00 0.200 49k-65k [HTML]F8800 0.103 [HTML]FFFA00 0.189 [HTML]F0000 0.000 [HTML]F3400 0.040 [HTML]A0D000 0.264 In-Domain (0k-8k) [HTML]FFBE00 0.143 [HTML]2C9600 0.351 [HTML]92C900 0.275 [HTML]7EBF00 0.290 [HTML]72B900 0.299 OOD (8k-65k) [HTML]F7A00 0.092 [HTML]9ECF00 0.266 [HTML]FFD600 0.162 [HTML]E6F300 0.211 [HTML]C2E100 0.238 Long-OOD (32k-65k) [HTML]F6E00 0.084 [HTML]FFE800 0.175 [HTML]F0200 0.002 [HTML]F8E00 0.108 [HTML]E2F100 0.215 Full (0k-65k) [HTML]F9A00 0.116 [HTML]68B400 0.306 [HTML]E2F100 0.215 [HTML]B6DB00 0.248 [HTML]9CCE00 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. Appendix D Additional Experimental Results D.1 ARMT language modeling pre-training 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. Figure 4: ARMT-augmented Gemma-3-1B-IT pre-training 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. D.2 ARMT Fine-Tuning After Pre-Training Model/ Lengths 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 0k-1k [HTML]068300 0.380 [HTML]249200 0.357 [HTML]249200 0.357 [HTML]229100 0.358 [HTML]128900 0.371 [HTML]008000 0.384 [HTML]128900 0.370 [HTML]028100 0.382 1k-2k [HTML]008000 0.385 [HTML]76B00 0.296 [HTML]78BC00 0.294 [HTML]6CB600 0.303 [HTML]62B100 0.310 [HTML]56AB00 0.320 [HTML]5AAD00 0.316 [HTML]76B00 0.295 2k-4k [HTML]2A9500 0.352 [HTML]A0D000 0.264 [HTML]92C900 0.275 [HTML]8C600 0.279 [HTML]9ACD00 0.268 [HTML]6EB700 0.301 [HTML]7ABD00 0.293 [HTML]6EB700 0.302 4k-6k [HTML]369B00 0.344 [HTML]C0E000 0.240 [HTML]82C100 0.286 [HTML]6CB600 0.303 [HTML]AED700 0.254 [HTML]7ABD00 0.293 [HTML]6CB600 0.303 [HTML]5CAE00 0.315 6k-8k [HTML]76B00 0.296 [HTML]DCEE00 0.219 [HTML]C2E100 0.238 [HTML]AED700 0.254 [HTML]D8EC00 0.222 [HTML]AAD500 0.256 [HTML]A6D300 0.260 [HTML]A4D200 0.261 8k-10k [HTML]68B400 0.306 [HTML]EEF700 0.205 [HTML]CAE500 0.233 [HTML]C6E300 0.235 [HTML]FFB200 0.135 [HTML]D0E800 0.228 [HTML]BADD00 0.245 [HTML]C2E100 0.238 10k-12k [HTML]9ACD00 0.269 [HTML]E4F200 0.213 [HTML]DAED00 0.221 [HTML]BEDF00 0.241 [HTML]F4400 0.052 [HTML]CCE600 0.231 [HTML]A6D300 0.260 [HTML]A6D300 0.260 12k-14k [HTML]8AC500 0.280 [HTML]FEFF00 0.193 [HTML]D4EA00 0.225 [HTML]C4E200 0.237 [HTML]F2A00 0.033 [HTML]C4E200 0.237 [HTML]B0D800 0.252 [HTML]B2D900 0.251 14k-16k [HTML]CCE600 0.231 [HTML]FFEE00 0.180 [HTML]DEEF00 0.217 [HTML]C0E000 0.240 [HTML]F2200 0.026 [HTML]D6EB00 0.224 [HTML]A2D100 0.262 [HTML]BADD00 0.244 16k-24k [HTML]9ACD00 0.269 [HTML]F9800 0.115 [HTML]F6FB00 0.200 [HTML]B0D800 0.252 [HTML]F0C00 0.010 [HTML]B8DC00 0.246 [HTML]A6D300 0.260 [HTML]94CA00 0.273 24k-32k [HTML]B6DB00 0.247 [HTML]F1400 0.016 [HTML]FFDE00 0.168 [HTML]ECF600 0.207 [HTML]F0000 0.001 [HTML]FEFF00 0.194 [HTML]DAED00 0.221 [HTML]EAF500 0.209 32k-49k [HTML]FFE200 0.171 [HTML]F0200 0.002 [HTML]FFAA00 0.128 [HTML]F6FB00 0.200 [HTML]F0000 0.000 [HTML]F6FB00 0.199 [HTML]ECF600 0.207 [HTML]D6EB00 0.223 49k-65k [HTML]FFFA00 0.189 [HTML]F0000 0.000 [HTML]F3400 0.040 [HTML]A0D000 0.264 [HTML]F0000 0.000 [HTML]7EBF00 0.289 [HTML]70B800 0.300 [HTML]84C200 0.285 In-Domain (0k-8k) [HTML]2C9600 0.351 [HTML]92C900 0.275 [HTML]7EBF00 0.290 [HTML]72B900 0.299 [HTML]84C200 0.285 [HTML]62B100 0.311 [HTML]66B300 0.308 [HTML]62B100 0.311 OOD (8k-65k) [HTML]9ECF00 0.266 [HTML]FFD600 0.162 [HTML]E6F300 0.211 [HTML]C2E100 0.238 [HTML]F3C00 0.045 [HTML]CCE600 0.231 [HTML]B0D800 0.252 [HTML]B2D900 0.250 Long-OOD (32k-65k) [HTML]FFE800 0.175 [HTML]F0200 0.002 [HTML]F8E00 0.108 [HTML]E2F100 0.215 [HTML]F0000 0.000 [HTML]DAED00 0.220 [HTML]D0E800 0.228 [HTML]C4E200 0.237 Full (0k-65k) [HTML]68B400 0.306 [HTML]E2F100 0.215 [HTML]B6DB00 0.248 [HTML]9CCE00 0.266 [HTML]FFD000 0.157 [HTML]9ACD00 0.268 [HTML]8C600 0.278 [HTML]8C600 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 [HTML]2E9700 0.795 [HTML]389C00 0.782 [HTML]389C00 0.782 [HTML]48A400 0.756 [HTML]40A000 0.769 [HTML]48A400 0.756 [HTML]48A400 0.756 1k-2k [HTML]329900 0.790 [HTML]5EAF00 0.724 [HTML]64B200 0.714 [HTML]52A900 0.743 [HTML]4CA600 0.752 [HTML]52A900 0.743 [HTML]52A900 0.743 2k-4k [HTML]3C9E00 0.774 [HTML]74BA00 0.689 [HTML]6EB700 0.698 [HTML]56AB00 0.736 [HTML]68B400 0.708 [HTML]68B400 0.708 [HTML]56AB00 0.736 4k-6k [HTML]42A100 0.767 [HTML]A4D200 0.616 [HTML]92C900 0.644 [HTML]6EB700 0.699 [HTML]80C000 0.671 [HTML]80C000 0.671 [HTML]54A00 0.740 6k-8k [HTML]048200 0.859 [HTML]AAD500 0.609 [HTML]62B100 0.717 [HTML]2A9500 0.804 [HTML]309800 0.793 [HTML]2A9500 0.804 [HTML]048200 0.859 8k-10k [HTML]0A8500 0.852 [HTML]ECF600 0.506 [HTML]4A500 0.753 [HTML]4A500 0.753 [HTML]4A500 0.753 [HTML]329900 0.790 [HTML]1A8D00 0.827 10k-12k [HTML]008000 0.868 [HTML]FFBA00 0.374 [HTML]82C100 0.670 [HTML]329900 0.791 [HTML]72B900 0.692 [HTML]4EA700 0.747 [HTML]008000 0.868 12k-14k [HTML]349A00 0.787 [HTML]FFCE00 0.404 [HTML]6CB600 0.702 [HTML]50A800 0.745 [HTML]50A800 0.745 [HTML]58AC00 0.734 [HTML]42A100 0.766 14k-16k [HTML]44A200 0.764 [HTML]FFA200 0.337 [HTML]94CA00 0.640 [HTML]78BC00 0.685 [HTML]9CCE00 0.629 [HTML]94CA00 0.640 [HTML]70B800 0.697 16k-24k [HTML]249200 0.812 [HTML]F5C00 0.229 [HTML]90C800 0.646 [HTML]48A400 0.757 [HTML]9ACD00 0.632 [HTML]96CB00 0.639 [HTML]48A400 0.757 24k-32k [HTML]64B200 0.713 [HTML]F4600 0.198 [HTML]BADD00 0.584 [HTML]64B200 0.713 [HTML]BADD00 0.584 [HTML]80C000 0.673 [HTML]249200 0.812 32k-49k [HTML]B6DB00 0.591 [HTML]F2000 0.139 [HTML]AAD500 0.609 [HTML]42A100 0.765 [HTML]DEEF00 0.530 [HTML]AAD500 0.609 [HTML]329900 0.791 49k-65k [HTML]F9400 0.317 [HTML]F0000 0.089 [HTML]ACD600 0.604 [HTML]64B200 0.713 [HTML]FFDC00 0.426 [HTML]D4EA00 0.545 [HTML]44A200 0.762 In-Domain (0k-8k) [HTML]2E9700 0.797 [HTML]78BC00 0.685 [HTML]66B300 0.711 [HTML]4EA700 0.749 [HTML]54A00 0.740 [HTML]54A00 0.738 [HTML]42A100 0.767 OOD (8k-65k) [HTML]68B400 0.709 [HTML]F7600 0.271 [HTML]90C800 0.647 [HTML]52A900 0.741 [HTML]A4D200 0.618 [HTML]84C200 0.665 [HTML]369B00 0.783 Long-OOD (32k-65k) [HTML]F400 0.463 [HTML]F1000 0.116 [HTML]AAD500 0.607 [HTML]52A900 0.741 [HTML]FEFF00 0.481 [HTML]BCDE00 0.579 [HTML]3A9D00 0.777 Full (0k-65k) [HTML]52A900 0.741 [HTML]FFD800 0.419 [HTML]82C100 0.670 [HTML]50A800 0.744 [HTML]86C300 0.661 [HTML]74BA00 0.691 [HTML]3A9D00 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. After ARMT pre-training, we also conducted fine-tuning 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. D.3 Ablation Study for Associative Memory Layers Model/ Lengths 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 0k-1k [HTML]3A9D00 0.795 [HTML]58AC00 0.756 [HTML]58AC00 0.756 [HTML]58AC00 0.756 [HTML]58AC00 0.756 [HTML]58AC00 0.756 [HTML]58AC00 0.756 [HTML]58AC00 0.756 1k-2k [HTML]3E9F00 0.790 [HTML]62B100 0.743 [HTML]6AB500 0.733 [HTML]62B100 0.743 [HTML]ECF600 0.571 [HTML]98C00 0.676 [HTML]DCEE00 0.590 [HTML]72B900 0.724 2k-4k [HTML]4A500 0.774 [HTML]68B400 0.736 [HTML]68B400 0.736 [HTML]70B800 0.726 [HTML]FFAC00 0.443 [HTML]78BC00 0.717 [HTML]DAED00 0.594 [HTML]7EBF00 0.708 4k-6k [HTML]50A800 0.767 [HTML]86C300 0.699 [HTML]86C300 0.699 [HTML]86C300 0.699 [HTML]F5200 0.329 [HTML]9CCE00 0.671 [HTML]FFEA00 0.521 [HTML]7CBE00 0.712 6k-8k [HTML]068300 0.859 [HTML]329900 0.804 [HTML]2A9500 0.815 [HTML]329900 0.804 [HTML]F4600 0.315 [HTML]66B300 0.739 [HTML]D6EB00 0.598 [HTML]3A9D00 0.793 8k-10k [HTML]0C8600 0.852 [HTML]5AAD00 0.753 [HTML]5AAD00 0.753 [HTML]52A900 0.765 [HTML]F8600 0.395 [HTML]78BC00 0.716 [HTML]F200 0.531 [HTML]52A900 0.765 10k-12k [HTML]008000 0.868 [HTML]3C9E00 0.791 [HTML]46A300 0.780 [HTML]3C9E00 0.791 [HTML]F5A00 0.341 [HTML]68B400 0.736 [HTML]FFDE00 0.505 [HTML]56AB00 0.758 12k-14k [HTML]40A000 0.787 [HTML]62B100 0.745 [HTML]50A800 0.766 [HTML]50A800 0.766 [HTML]F2800 0.277 [HTML]2E9700 0.809 [HTML]FFB800 0.457 [HTML]5AAD00 0.755 14k-16k [HTML]52A900 0.764 [HTML]90C800 0.685 [HTML]9ACD00 0.674 [HTML]9ACD00 0.674 [HTML]F5800 0.337 [HTML]7EBF00 0.708 [HTML]FFC400 0.472 [HTML]A2D100 0.663 16k-24k [HTML]2C9600 0.812 [HTML]58AC00 0.757 [HTML]5EAF00 0.750 [HTML]52A900 0.764 [HTML]F3E00 0.306 [HTML]74BA00 0.722 [HTML]FFB800 0.458 [HTML]6EB700 0.729 24k-32k [HTML]7ABD00 0.713 [HTML]7ABD00 0.713 [HTML]82C100 0.703 [HTML]82C100 0.703 [HTML]F3000 0.287 [HTML]92C900 0.683 [HTML]FFA600 0.436 [HTML]8AC500 0.693 32k-49k [HTML]DCEE00 0.591 [HTML]52A900 0.765 [HTML]58AC00 0.757 [HTML]52A900 0.765 [HTML]F0000 0.226 [HTML]66B300 0.739 [HTML]FFDC00 0.504 [HTML]52A900 0.765 49k-65k [HTML]F4800 0.317 [HTML]7ABD00 0.713 [HTML]72B900 0.723 [HTML]72B900 0.723 [HTML]F6600 0.356 [HTML]92C900 0.683 [HTML]FFDE00 0.505 [HTML]7ABD00 0.713 In-Domain (0k-8k) [HTML]389C00 0.797 [HTML]5EAF00 0.749 [HTML]5EAF00 0.749 [HTML]60B000 0.747 [HTML]FFCC00 0.482 [HTML]7CBE00 0.711 [HTML]CCE600 0.610 [HTML]66B300 0.738 OOD (8k-65k) [HTML]7EBF00 0.709 [HTML]64B200 0.741 [HTML]66B300 0.739 [HTML]62B100 0.745 [HTML]F4400 0.311 [HTML]72B900 0.724 [HTML]FFCA00 0.481 [HTML]6CB600 0.730 Long-OOD (32k-65k) [HTML]FFBC00 0.463 [HTML]64B200 0.741 [HTML]64B200 0.741 [HTML]60B000 0.745 [HTML]F3000 0.287 [HTML]7ABD00 0.713 [HTML]FFDE00 0.504 [HTML]64B200 0.741 Full (0k-65k) [HTML]64B200 0.741 [HTML]62B100 0.744 [HTML]64B200 0.743 [HTML]60B000 0.746 [HTML]F7400 0.372 [HTML]76B00 0.720 [HTML]F000 0.527 [HTML]6AB500 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 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 0k-1k [HTML]088400 0.380 [HTML]349A00 0.358 [HTML]349A00 0.358 [HTML]349A00 0.358 [HTML]349A00 0.358 [HTML]349A00 0.358 [HTML]349A00 0.358 [HTML]349A00 0.358 1k-2k [HTML]008000 0.385 [HTML]9ECF00 0.303 [HTML]98C00 0.306 [HTML]A2D100 0.301 [HTML]FFD600 0.232 [HTML]FFDA00 0.234 [HTML]FFA600 0.207 [HTML]A6D300 0.299 2k-4k [HTML]40A000 0.352 [HTML]CCE600 0.279 [HTML]D0E800 0.277 [HTML]CCE600 0.279 [HTML]F7E00 0.186 [HTML]FFB000 0.212 [HTML]F7200 0.180 [HTML]C0E000 0.285 4k-6k [HTML]4EA700 0.344 [HTML]9ECF00 0.303 [HTML]AED700 0.295 [HTML]D8EC00 0.273 [HTML]F5400 0.165 [HTML]F9400 0.198 [HTML]F4A00 0.160 [HTML]CEE700 0.278 6k-8k [HTML]ACD600 0.296 [HTML]FEFF00 0.254 [HTML]F200 0.246 [HTML]FFEC00 0.243 [HTML]F5400 0.165 [HTML]F8000 0.188 [HTML]F2C00 0.144 [HTML]FFEE00 0.244 8k-10k [HTML]98C00 0.306 [HTML]FFDC00 0.235 [HTML]FFE200 0.238 [HTML]FFBE00 0.219 [HTML]F4000 0.155 [HTML]F8400 0.190 [HTML]F1600 0.133 [HTML]FFD800 0.233 10k-12k [HTML]E0F000 0.269 [HTML]FFE800 0.241 [HTML]FFE000 0.237 [HTML]FFCA00 0.226 [HTML]F3C00 0.152 [HTML]F8000 0.188 [HTML]F0E00 0.129 [HTML]FFE800 0.241 12k-14k [HTML]CAE500 0.280 [HTML]FFE000 0.237 [HTML]FFD800 0.233 [HTML]FFC400 0.223 [HTML]F4200 0.156 [HTML]F6A00 0.176 [HTML]F3000 0.146 [HTML]FFD200 0.230 14k-16k [HTML]FFD400 0.231 [HTML]FFE600 0.240 [HTML]FFDA00 0.234 [HTML]FFC200 0.222 [HTML]F4000 0.154 [HTML]F5C00 0.169 [HTML]F2400 0.140 [HTML]FFE000 0.237 16k-24k [HTML]E0F000 0.269 [HTML]FFFE00 0.252 [HTML]F600 0.248 [HTML]FFC000 0.220 [HTML]F4000 0.155 [HTML]F7600 0.182 [HTML]F0000 0.121 [HTML]F400 0.247 24k-32k [HTML]F400 0.247 [HTML]FFA600 0.207 [HTML]FFB000 0.212 [HTML]FFA400 0.206 [HTML]F2400 0.140 [HTML]F8000 0.187 [HTML]F0E00 0.129 [HTML]FFBA00 0.217 32k-49k [HTML]F6000 0.171 [HTML]F9800 0.200 [HTML]FFB200 0.213 [HTML]FFC000 0.220 [HTML]F3C00 0.152 [HTML]F4A00 0.160 [HTML]F3200 0.147 [HTML]FFD000 0.229 49k-65k [HTML]F8200 0.189 [HTML]EAF500 0.264 [HTML]70B800 0.327 [HTML]C0E000 0.286 [HTML]FFD800 0.233 [HTML]FFA200 0.205 [HTML]FFA600 0.207 [HTML]68B400 0.331 In-Domain (0k-8k) [HTML]40A000 0.351 [HTML]A6D300 0.299 [HTML]ACD600 0.296 [HTML]B6DB00 0.291 [HTML]FFC000 0.220 [HTML]FFE000 0.237 [HTML]FFAA00 0.209 [HTML]B2D900 0.293 OOD (8k-65k) [HTML]E6F300 0.266 [HTML]FFE200 0.238 [HTML]FFE000 0.237 [HTML]FFC200 0.222 [HTML]F4000 0.154 [HTML]F7200 0.181 [HTML]F1A00 0.135 [HTML]FFE000 0.237 Long-OOD (32k-65k) [HTML]F6800 0.175 [HTML]FFB400 0.215 [HTML]FFE400 0.239 [HTML]FFDC00 0.235 [HTML]F6000 0.171 [HTML]F5E00 0.170 [HTML]F4C00 0.161 [HTML]FFFE00 0.253 Full (0k-65k) [HTML]98C00 0.306 [HTML]E4F200 0.266 [HTML]E8F400 0.264 [HTML]FEFF00 0.254 [HTML]F7C00 0.185 [HTML]FFA600 0.207 [HTML]F5E00 0.170 [HTML]ECF600 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 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 0k-1k [HTML]0A8500 0.380 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]5AAD00 0.347 [HTML]349A00 0.363 1k-2k [HTML]008000 0.385 [HTML]C4E200 0.303 [HTML]CCE600 0.299 [HTML]E0F000 0.291 [HTML]A2D100 0.317 2k-4k [HTML]4EA700 0.352 [HTML]FCFE00 0.279 [HTML]EEF700 0.285 [HTML]FFE000 0.265 [HTML]C0E000 0.304 4k-6k [HTML]62B100 0.344 [HTML]C4E200 0.303 [HTML]FEFF00 0.278 [HTML]FFD800 0.262 [HTML]EAF500 0.287 6k-8k [HTML]D4EA00 0.296 [HTML]FFC600 0.254 [HTML]FFAE00 0.244 [HTML]FFA600 0.241 [HTML]FFBC00 0.250 8k-10k [HTML]BCDE00 0.306 [HTML]F9800 0.235 [HTML]F9400 0.233 [HTML]F9000 0.232 [HTML]F8E00 0.231 10k-12k [HTML]FFEA00 0.269 [HTML]FFA600 0.241 [HTML]FFA600 0.241 [HTML]F9000 0.232 [HTML]F8E00 0.231 12k-14k [HTML]FAFD00 0.280 [HTML]F9C00 0.237 [HTML]F8C00 0.230 [HTML]F7A00 0.222 [HTML]F7600 0.221 14k-16k [HTML]F8E00 0.231 [HTML]FFA400 0.240 [HTML]F9C00 0.237 [HTML]F9400 0.233 [HTML]FFAC00 0.243 16k-24k [HTML]FFEA00 0.269 [HTML]FFC000 0.252 [HTML]FFB400 0.247 [HTML]FFA800 0.242 [HTML]FFD200 0.259 24k-32k [HTML]FFB400 0.247 [HTML]F5600 0.207 [HTML]F6E00 0.217 [HTML]F2600 0.187 [HTML]F4E00 0.204 32k-49k [HTML]F0000 0.171 [HTML]F4400 0.200 [HTML]F8A00 0.229 [HTML]F5200 0.206 [HTML]F5000 0.205 49k-65k [HTML]F2A00 0.189 [HTML]FFDE00 0.264 [HTML]80C000 0.331 [HTML]70B800 0.338 [HTML]FFB200 0.246 In-Domain (0k-8k) [HTML]50A800 0.351 [HTML]CCE600 0.299 [HTML]DCEE00 0.293 [HTML]F8FC00 0.281 [HTML]C0E000 0.304 OOD (8k-65k) [HTML]FFE200 0.266 [HTML]F9E00 0.238 [HTML]F9C00 0.237 [HTML]F8A00 0.229 [HTML]F9600 0.234 Long-OOD (32k-65k) [HTML]F0800 0.175 [HTML]F6800 0.215 [HTML]FFC200 0.253 [HTML]F9C00 0.236 [HTML]F6600 0.214 Full (0k-65k) [HTML]BCDE00 0.306 [HTML]FFE400 0.266 [HTML]FFDA00 0.263 [HTML]FFC400 0.254 [HTML]FFE400 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. 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, 16, 17 and 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. Model/ Lengths 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 0k-1k [HTML]0A8500 0.380 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 [HTML]40A000 0.358 1k-2k [HTML]008000 0.385 [HTML]C2E100 0.303 [HTML]B8DC00 0.307 [HTML]CAE500 0.300 [HTML]D6EB00 0.295 [HTML]C4E200 0.302 [HTML]D2E900 0.296 [HTML]CCE600 0.299 [HTML]C2E100 0.303 [HTML]D8EC00 0.294 [HTML]C2E100 0.303 [HTML]E2F100 0.290 [HTML]C4E200 0.302 [HTML]C2E100 0.303 [HTML]C4E200 0.302 2k-4k [HTML]4EA700 0.352 [HTML]FCFE00 0.279 [HTML]FEFF00 0.278 [HTML]F600 0.274 [HTML]FCFE00 0.279 [HTML]FAFD00 0.280 [HTML]F2F900 0.283 [HTML]F2F900 0.283 [HTML]FCFE00 0.279 [HTML]FEFF00 0.278 [HTML]FFFE00 0.277 [HTML]FFEA00 0.269 [HTML]F4FA00 0.282 [HTML]FAFD00 0.280 [HTML]D2E900 0.296 4k-6k [HTML]60B000 0.344 [HTML]C2E100 0.303 [HTML]CEE700 0.298 [HTML]CCE600 0.299 [HTML]CEE700 0.298 [HTML]D0E800 0.297 [HTML]D0E800 0.297 [HTML]CEE700 0.298 [HTML]C8E400 0.301 [HTML]DCEE00 0.292 [HTML]C4E200 0.302 [HTML]FAFD00 0.280 [HTML]CCE600 0.299 [HTML]C8E400 0.301 [HTML]EEF700 0.285 6k-8k [HTML]D2E900 0.296 [HTML]FFC800 0.254 [HTML]FFB600 0.247 [HTML]FFC400 0.253 [HTML]FFCA00 0.255 [HTML]FFB400 0.246 [HTML]FFCA00 0.255 [HTML]FFCC00 0.256 [HTML]FFC800 0.254 [HTML]FFB400 0.246 [HTML]FFBC00 0.249 [HTML]FFA600 0.240 [HTML]FFC200 0.252 [HTML]FFC800 0.254 [HTML]FFC800 0.254 8k-10k [HTML]BCDE00 0.306 [HTML]F9A00 0.235 [HTML]F8A00 0.228 [HTML]F9200 0.232 [HTML]F9800 0.234 [HTML]FFA400 0.239 [HTML]FFAC00 0.243 [HTML]F9600 0.233 [HTML]F9E00 0.237 [HTML]FFAC00 0.243 [HTML]F9E00 0.237 [HTML]F6000 0.211 [HTML]F9C00 0.236 [HTML]F9800 0.234 [HTML]F9000 0.231 10k-12k [HTML]FFEA00 0.269 [HTML]FFA800 0.241 [HTML]F9000 0.231 [HTML]F9C00 0.236 [HTML]FFA800 0.241 [HTML]FFB000 0.244 [HTML]FFA600 0.240 [HTML]FFA400 0.239 [HTML]FFB000 0.244 [HTML]FFA600 0.240 [HTML]FFA800 0.241 [HTML]F8600 0.227 [HTML]FFA400 0.239 [HTML]FFAA00 0.242 [HTML]FFAA00 0.242 12k-14k [HTML]FAFD00 0.280 [HTML]F9E00 0.237 [HTML]FFA000 0.238 [HTML]F9E00 0.237 [HTML]F9800 0.234 [HTML]F8E00 0.230 [HTML]F9600 0.233 [HTML]F9200 0.232 [HTML]F9A00 0.235 [HTML]F9E00 0.237 [HTML]F9000 0.231 [HTML]F6C00 0.216 [HTML]F9800 0.234 [HTML]F9A00 0.235 [HTML]F8600 0.227 14k-16k [HTML]F9000 0.231 [HTML]FFA600 0.240 [HTML]FFA000 0.238 [HTML]FFAA00 0.242 [HTML]F9C00 0.236 [HTML]FFA000 0.238 [HTML]FFA400 0.239 [HTML]FFA600 0.240 [HTML]FFA600 0.240 [HTML]FFA600 0.240 [HTML]FFAC00 0.243 [HTML]F7200 0.218 [HTML]FFA600 0.240 [HTML]FFA600 0.240 [HTML]FFA400 0.239 16k-24k [HTML]FFEA00 0.269 [HTML]FFC200 0.252 [HTML]FFBE00 0.250 [HTML]FFB800 0.248 [HTML]FFC800 0.254 [HTML]FFCE00 0.257 [HTML]FFC400 0.253 [HTML]FFD000 0.258 [HTML]FFBC00 0.249 [HTML]FFBC00 0.249 [HTML]FFCE00 0.257 [HTML]F8A00 0.228 [HTML]FFC800 0.254 [HTML]FFC000 0.251 [HTML]FFB400 0.246 24k-32k [HTML]FFB600 0.247 [HTML]F5800 0.207 [HTML]F6400 0.212 [HTML]F4E00 0.203 [HTML]F5200 0.205 [HTML]F6800 0.214 [HTML]F5000 0.204 [HTML]F5E00 0.210 [HTML]F5000 0.204 [HTML]F5200 0.205 [HTML]F6C00 0.216 [HTML]F2E00 0.190 [HTML]F5400 0.206 [HTML]F5000 0.204 [HTML]F6C00 0.216 32k-49k [HTML]F0200 0.171 [HTML]F4600 0.200 [HTML]F4400 0.199 [HTML]F4E00 0.203 [HTML]F5400 0.206 [HTML]F6800 0.214 [HTML]F5200 0.205 [HTML]F6000 0.211 [HTML]F5000 0.204 [HTML]F6C00 0.216 [HTML]F7E00 0.223 [HTML]F0000 0.170 [HTML]F4200 0.198 [HTML]F5000 0.204 [HTML]F8400 0.226 49k-65k [HTML]F2C00 0.189 [HTML]FFDE00 0.264 [HTML]F000 0.271 [HTML]60B000 0.344 [HTML]FEFF00 0.278 [HTML]F0F800 0.284 [HTML]F200 0.272 [HTML]B2D900 0.310 [HTML]FFDE00 0.264 [HTML]F4FA00 0.282 [HTML]108800 0.378 [HTML]F0000 0.170 [HTML]FFDE00 0.264 [HTML]FFDE00 0.264 [HTML]F0F800 0.284 In-Domain (0k-8k) [HTML]50A800 0.351 [HTML]CCE600 0.299 [HTML]D0E800 0.297 [HTML]D2E900 0.296 [HTML]D2E900 0.297 [HTML]D2E900 0.296 [HTML]D0E800 0.297 [HTML]CEE700 0.298 [HTML]CCE600 0.299 [HTML]DAED00 0.293 [HTML]D0E800 0.297 [HTML]E8F400 0.287 [HTML]CEE700 0.298 [HTML]CCE600 0.299 [HTML]CCE600 0.299 OOD (8k-65k) [HTML]FFE200 0.266 [HTML]FFA000 0.238 [HTML]F9A00 0.235 [HTML]F9E00 0.237 [HTML]F9E00 0.237 [HTML]FFA400 0.239 [HTML]FFA200 0.238 [HTML]FFA200 0.238 [HTML]FFA000 0.238 [HTML]FFA400 0.239 [HTML]FFA800 0.241 [HTML]F6C00 0.216 [HTML]FFA000 0.237 [HTML]F9E00 0.237 [HTML]F9C00 0.236 Long-OOD (32k-65k) [HTML]F0C00 0.175 [HTML]F6A00 0.215 [HTML]F6C00 0.216 [HTML]F9C00 0.236 [HTML]F7C00 0.223 [HTML]F8E00 0.230 [HTML]F7800 0.220 [HTML]F9800 0.234 [HTML]F7000 0.218 [HTML]F9000 0.231 [HTML]FFD200 0.259 [HTML]F0000 0.170 [HTML]F6600 0.213 [HTML]F7000 0.218 [HTML]FFA400 0.239 Full (0k-65k) [HTML]BCDE00 0.306 [HTML]FFE400 0.266 [HTML]FFDE00 0.264 [HTML]FFE000 0.265 [HTML]FFE000 0.265 [HTML]FFE400 0.266 [HTML]FFE400 0.266 [HTML]FFE400 0.266 [HTML]FFE400 0.266 [HTML]FFE000 0.264 [HTML]FFE600 0.267 [HTML]FFBC00 0.249 [HTML]FFE400 0.266 [HTML]FFE400 0.266 [HTML]FFE200 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 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 0k-1k [HTML]0A8500 0.380 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 [HTML]3A9D00 0.358 1k-2k [HTML]008000 0.385 [HTML]B0D800 0.303 [HTML]B6DB00 0.301 [HTML]FFC200 0.239 [HTML]AED700 0.304 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]D0E800 0.289 [HTML]CAE500 0.291 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]C0E000 0.296 [HTML]FFA600 0.226 2k-4k [HTML]46A300 0.352 [HTML]E4F200 0.279 [HTML]EAF500 0.277 [HTML]F9600 0.219 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]F0F800 0.274 [HTML]E8F400 0.278 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]E4F200 0.279 [HTML]F7E00 0.208 4k-6k [HTML]58AC00 0.344 [HTML]B0D800 0.303 [HTML]F8FC00 0.270 [HTML]F6800 0.197 [HTML]B4DA00 0.302 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]D4EA00 0.287 [HTML]BEDF00 0.297 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]B0D800 0.303 [HTML]B6DB00 0.301 [HTML]F5800 0.190 6k-8k [HTML]C0E000 0.296 [HTML]FFE200 0.254 [HTML]FFC200 0.239 [HTML]F4A00 0.184 [HTML]FFE200 0.254 [HTML]FFE000 0.253 [HTML]FFE200 0.254 [HTML]F600 0.263 [HTML]FFE000 0.253 [HTML]FFE200 0.254 [HTML]FFE200 0.254 [HTML]FFE200 0.254 [HTML]FFE200 0.254 [HTML]FFE000 0.253 [HTML]F6400 0.196 8k-10k [HTML]AAD500 0.306 [HTML]FFBA00 0.235 [HTML]F9E00 0.222 [HTML]F3A00 0.176 [HTML]FFBA00 0.235 [HTML]FFC600 0.241 [HTML]FFBA00 0.235 [HTML]FFB400 0.232 [HTML]FFC600 0.241 [HTML]FFBA00 0.235 [HTML]FFBA00 0.235 [HTML]FFBA00 0.235 [HTML]FFBA00 0.235 [HTML]FFB800 0.234 [HTML]F5400 0.188 10k-12k [HTML]FAFD00 0.269 [HTML]FFC600 0.241 [HTML]FFA200 0.224 [HTML]F4000 0.179 [HTML]FFC800 0.242 [HTML]FFC200 0.239 [HTML]FFC600 0.241 [HTML]FFB400 0.232 [HTML]FFBE00 0.237 [HTML]FFC600 0.241 [HTML]FFC600 0.241 [HTML]FFC600 0.241 [HTML]FFC600 0.241 [HTML]FFC600 0.241 [HTML]F5000 0.186 12k-14k [HTML]E2F100 0.280 [HTML]FFBE00 0.237 [HTML]F8E00 0.215 [HTML]F3A00 0.176 [HTML]FFBE00 0.237 [HTML]FFBE00 0.237 [HTML]FFBE00 0.237 [HTML]FFB800 0.234 [HTML]FFBC00 0.236 [HTML]FFBE00 0.237 [HTML]FFC000 0.238 [HTML]FFBE00 0.237 [HTML]FFBE00 0.237 [HTML]FFBA00 0.235 [HTML]F3600 0.174 14k-16k [HTML]FFB000 0.231 [HTML]FFC400 0.240 [HTML]F8A00 0.213 [HTML]F3C00 0.177 [HTML]FFC600 0.241 [HTML]FFC600 0.241 [HTML]FFC400 0.240 [HTML]FFB600 0.233 [HTML]FFC800 0.242 [HTML]FFC400 0.240 [HTML]FFC200 0.239 [HTML]FFC400 0.240 [HTML]FFC400 0.240 [HTML]FFCA00 0.243 [HTML]F2E00 0.171 16k-24k [HTML]FAFD00 0.269 [HTML]FFDE00 0.252 [HTML]F8600 0.211 [HTML]F3C00 0.177 [HTML]FFDE00 0.252 [HTML]FFDC00 0.251 [HTML]FFDE00 0.252 [HTML]FFD800 0.249 [HTML]FFDE00 0.252 [HTML]FFDE00 0.252 [HTML]FFDE00 0.252 [HTML]FFDE00 0.252 [HTML]FFDE00 0.252 [HTML]FFDC00 0.251 [HTML]F4800 0.183 24k-32k [HTML]FFD400 0.247 [HTML]F7C00 0.207 [HTML]F8200 0.209 [HTML]F1E00 0.163 [HTML]F7C00 0.207 [HTML]F7C00 0.207 [HTML]F7C00 0.207 [HTML]F8200 0.209 [HTML]F7800 0.205 [HTML]F7C00 0.207 [HTML]F7C00 0.207 [HTML]F7C00 0.207 [HTML]F7C00 0.207 [HTML]F7E00 0.208 [HTML]F5800 0.190 32k-49k [HTML]F2E00 0.171 [HTML]F6E00 0.200 [HTML]F5600 0.189 [HTML]F0000 0.149 [HTML]F6E00 0.200 [HTML]F6E00 0.200 [HTML]F6E00 0.200 [HTML]F9E00 0.222 [HTML]F7200 0.202 [HTML]F6E00 0.200 [HTML]F6E00 0.200 [HTML]F6E00 0.200 [HTML]F7400 0.203 [HTML]F7400 0.203 [HTML]F1E00 0.163 49k-65k [HTML]F5600 0.189 [HTML]F800 0.264 [HTML]FFAE00 0.230 [HTML]FFC200 0.239 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F2F900 0.273 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F800 0.264 [HTML]F7800 0.205 In-Domain (0k-8k) [HTML]48A400 0.351 [HTML]BADD00 0.299 [HTML]D0E800 0.289 [HTML]FFC200 0.239 [HTML]BADD00 0.299 [HTML]BADD00 0.299 [HTML]BADD00 0.299 [HTML]C4E200 0.294 [HTML]C2E100 0.295 [HTML]BADD00 0.299 [HTML]BADD00 0.299 [HTML]BADD00 0.299 [HTML]BADD00 0.299 [HTML]BEDF00 0.297 [HTML]FFBA00 0.235 OOD (8k-65k) [HTML]FFFC00 0.266 [HTML]FFC000 0.238 [HTML]F9000 0.215 [HTML]F3A00 0.176 [HTML]FFC000 0.238 [HTML]FFC200 0.238 [HTML]FFC000 0.238 [HTML]FFB800 0.234 [HTML]FFC000 0.238 [HTML]FFC000 0.238 [HTML]FFC000 0.238 [HTML]FFC000 0.238 [HTML]FFC000 0.238 [HTML]FFC000 0.238 [HTML]F4400 0.180 Long-OOD (32k-65k) [HTML]F3800 0.175 [HTML]F8E00 0.215 [HTML]F6A00 0.198 [HTML]F2C00 0.170 [HTML]F8E00 0.215 [HTML]F8E00 0.215 [HTML]F8E00 0.215 [HTML]FFB600 0.234 [HTML]F9200 0.216 [HTML]F8E00 0.215 [HTML]F8E00 0.215 [HTML]F8E00 0.215 [HTML]F9200 0.217 [HTML]F9200 0.217 [HTML]F3200 0.173 Full (0k-65k) [HTML]AAD500 0.306 [HTML]FFFE00 0.266 [HTML]FFDA00 0.250 [HTML]F7A00 0.205 [HTML]FFFE00 0.267 [HTML]FFFE00 0.267 [HTML]FFFE00 0.266 [HTML]F400 0.262 [HTML]FFFA00 0.265 [HTML]FFFE00 0.266 [HTML]FFFE00 0.266 [HTML]FFFE00 0.266 [HTML]FFFE00 0.266 [HTML]FFFC00 0.265 [HTML]F7A00 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 [HTML]42A100 0.795 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 [HTML]68B400 0.756 1k-2k [HTML]48A400 0.790 [HTML]74BA00 0.743 [HTML]7CBE00 0.733 [HTML]7CBE00 0.733 [HTML]7CBE00 0.733 [HTML]7CBE00 0.733 [HTML]7CBE00 0.733 [HTML]7CBE00 0.733 [HTML]74BA00 0.743 [HTML]74BA00 0.743 [HTML]74BA00 0.743 [HTML]7CBE00 0.733 [HTML]74BA00 0.743 [HTML]74BA00 0.743 [HTML]7CBE00 0.733 2k-4k [HTML]56AB00 0.774 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]82C100 0.726 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 4k-6k [HTML]5CAE00 0.767 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]AAD500 0.685 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]AAD500 0.685 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 [HTML]9CCE00 0.699 6k-8k [HTML]088400 0.859 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 [HTML]309800 0.815 [HTML]309800 0.815 [HTML]3A9D00 0.804 [HTML]309800 0.815 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 [HTML]309800 0.815 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 [HTML]3A9D00 0.804 8k-10k [HTML]0E8700 0.852 [HTML]6AB500 0.753 [HTML]6AB500 0.753 [HTML]6AB500 0.753 [HTML]76B00 0.741 [HTML]5EAF00 0.765 [HTML]76B00 0.741 [HTML]82C100 0.728 [HTML]6AB500 0.753 [HTML]6AB500 0.753 [HTML]52A900 0.778 [HTML]6AB500 0.753 [HTML]76B00 0.741 [HTML]6AB500 0.753 [HTML]76B00 0.741 10k-12k [HTML]008000 0.868 [HTML]46A300 0.791 [HTML]46A300 0.791 [HTML]5AAD00 0.769 [HTML]46A300 0.791 [HTML]3C9E00 0.802 [HTML]46A300 0.791 [HTML]3C9E00 0.802 [HTML]46A300 0.791 [HTML]46A300 0.791 [HTML]329900 0.813 [HTML]46A300 0.791 [HTML]46A300 0.791 [HTML]46A300 0.791 [HTML]46A300 0.791 12k-14k [HTML]4A500 0.787 [HTML]72B900 0.745 [HTML]7CBE00 0.734 [HTML]68B400 0.755 [HTML]7CBE00 0.734 [HTML]86C300 0.723 [HTML]7CBE00 0.734 [HTML]7CBE00 0.734 [HTML]7CBE00 0.734 [HTML]72B900 0.745 [HTML]5EAF00 0.766 [HTML]72B900 0.745 [HTML]72B900 0.745 [HTML]72B900 0.745 [HTML]72B900 0.745 14k-16k [HTML]60B000 0.764 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]9ECF00 0.697 [HTML]AAD500 0.685 [HTML]B4DA00 0.674 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]B4DA00 0.674 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]AAD500 0.685 [HTML]AAD500 0.685 16k-24k [HTML]349A00 0.812 [HTML]66B300 0.757 [HTML]60B000 0.764 [HTML]60B000 0.764 [HTML]60B000 0.764 [HTML]6CB600 0.750 [HTML]60B000 0.764 [HTML]66B300 0.757 [HTML]60B000 0.764 [HTML]66B300 0.757 [HTML]60B000 0.764 [HTML]5AAD00 0.771 [HTML]60B000 0.764 [HTML]66B300 0.757 [HTML]60B000 0.764 24k-32k [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]86C300 0.723 [HTML]98C00 0.703 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]98C00 0.703 [HTML]98C00 0.703 [HTML]98C00 0.703 [HTML]A2D100 0.693 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]98C00 0.703 32k-49k [HTML]FFFE00 0.591 [HTML]5EAF00 0.765 [HTML]66B300 0.757 [HTML]5EAF00 0.765 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]66B300 0.757 [HTML]5EAF00 0.765 [HTML]66B300 0.757 [HTML]5EAF00 0.765 [HTML]66B300 0.757 49k-65k [HTML]F0000 0.317 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]86C300 0.723 [HTML]90C800 0.713 [HTML]7CBE00 0.733 [HTML]86C300 0.723 [HTML]A2D100 0.693 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]74BA00 0.743 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]90C800 0.713 [HTML]90C800 0.713 In-Domain (0k-8k) [HTML]40A000 0.797 [HTML]6EB700 0.749 [HTML]70B800 0.747 [HTML]70B800 0.747 [HTML]70B800 0.747 [HTML]6EB700 0.749 [HTML]70B800 0.747 [HTML]6EB700 0.749 [HTML]6EB700 0.749 [HTML]6EB700 0.749 [HTML]6EB700 0.749 [HTML]72B900 0.744 [HTML]6EB700 0.749 [HTML]6EB700 0.749 [HTML]70B800 0.747 OOD (8k-65k) [HTML]92C900 0.709 [HTML]74BA00 0.741 [HTML]76B00 0.740 [HTML]72B900 0.745 [HTML]78BC00 0.738 [HTML]76B00 0.740 [HTML]76B00 0.740 [HTML]7CBE00 0.734 [HTML]76B00 0.739 [HTML]76B00 0.739 [HTML]6EB700 0.749 [HTML]72B900 0.744 [HTML]76B00 0.740 [HTML]74BA00 0.741 [HTML]76B00 0.739 Long-OOD (32k-65k) [HTML]F8600 0.463 [HTML]76B00 0.741 [HTML]7ABD00 0.736 [HTML]70B800 0.745 [HTML]7ABD00 0.736 [HTML]70B800 0.746 [HTML]74BA00 0.741 [HTML]82C100 0.727 [HTML]7ABD00 0.736 [HTML]7ABD00 0.736 [HTML]6CB600 0.750 [HTML]76B00 0.741 [HTML]7ABD00 0.736 [HTML]76B00 0.741 [HTML]7ABD00 0.736 Full (0k-65k) [HTML]76B00 0.741 [HTML]72B900 0.744 [HTML]74BA00 0.743 [HTML]70B800 0.746 [HTML]76B00 0.741 [HTML]72B900 0.743 [HTML]74BA00 0.743 [HTML]76B00 0.739 [HTML]74BA00 0.743 [HTML]74BA00 0.743 [HTML]6EB700 0.749 [HTML]72B900 0.744 [HTML]72B900 0.743 [HTML]72B900 0.744 [HTML]74BA00 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 [HTML]3C9E00 0.795 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 [HTML]5CAE00 0.756 1k-2k [HTML]40A000 0.790 [HTML]68B400 0.743 [HTML]70B800 0.733 [HTML]D6EB00 0.610 [HTML]68B400 0.743 [HTML]68B400 0.743 [HTML]68B400 0.743 [HTML]90C800 0.695 [HTML]90C800 0.695 [HTML]68B400 0.743 [HTML]68B400 0.743 [HTML]68B400 0.743 [HTML]60B000 0.752 [HTML]68B400 0.743 [HTML]70B800 0.733 2k-4k [HTML]4EA700 0.774 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]FFC400 0.491 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]76B00 0.726 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]6EB700 0.736 [HTML]76B00 0.726 [HTML]6EB700 0.736 [HTML]76B00 0.726 4k-6k [HTML]54A00 0.767 [HTML]8C600 0.699 [HTML]8C600 0.699 [HTML]F6A00 0.384 [HTML]8C600 0.699 [HTML]8C600 0.699 [HTML]8C600 0.699 [HTML]A4D200 0.671 [HTML]82C100 0.712 [HTML]8C600 0.699 [HTML]8C600 0.699 [HTML]8C600 0.699 [HTML]98C00 0.685 [HTML]8C600 0.699 [HTML]A4D200 0.671 6k-8k [HTML]068300 0.859 [HTML]349A00 0.804 [HTML]2C9600 0.815 [HTML]F6000 0.370 [HTML]349A00 0.804 [HTML]349A00 0.804 [HTML]349A00 0.804 [HTML]50A800 0.772 [HTML]3E9F00 0.793 [HTML]349A00 0.804 [HTML]349A00 0.804 [HTML]349A00 0.804 [HTML]3E9F00 0.793 [HTML]349A00 0.804 [HTML]46A300 0.783 8k-10k [HTML]0C8600 0.852 [HTML]60B000 0.753 [HTML]6AB500 0.741 [HTML]F7E00 0.407 [HTML]60B000 0.753 [HTML]60B000 0.753 [HTML]60B000 0.753 [HTML]7EBF00 0.716 [HTML]56AB00 0.765 [HTML]60B000 0.753 [HTML]60B000 0.753 [HTML]60B000 0.753 [HTML]56AB00 0.765 [HTML]60B000 0.753 [HTML]6AB500 0.741 10k-12k [HTML]008000 0.868 [HTML]40A000 0.791 [HTML]40A000 0.791 [HTML]F3E00 0.330 [HTML]40A000 0.791 [HTML]40A000 0.791 [HTML]40A000 0.791 [HTML]6EB700 0.736 [HTML]48A400 0.780 [HTML]40A000 0.791 [HTML]40A000 0.791 [HTML]40A000 0.791 [HTML]52A900 0.769 [HTML]52A900 0.769 [HTML]52A900 0.769 12k-14k [HTML]42A100 0.787 [HTML]66B300 0.745 [HTML]6EB700 0.734 [HTML]F0000 0.255 [HTML]66B300 0.745 [HTML]66B300 0.745 [HTML]66B300 0.745 [HTML]ACD600 0.660 [HTML]5EAF00 0.755 [HTML]66B300 0.745 [HTML]66B300 0.745 [HTML]66B300 0.745 [HTML]5EAF00 0.755 [HTML]5EAF00 0.755 [HTML]6EB700 0.734 14k-16k [HTML]56AB00 0.764 [HTML]98C00 0.685 [HTML]98C00 0.685 [HTML]F5600 0.360 [HTML]98C00 0.685 [HTML]98C00 0.685 [HTML]98C00 0.685 [HTML]D0E800 0.618 [HTML]8EC700 0.697 [HTML]98C00 0.685 [HTML]98C00 0.685 [HTML]98C00 0.685 [HTML]7CBE00 0.719 [HTML]8EC700 0.697 [HTML]8EC700 0.697 16k-24k [HTML]2E9700 0.812 [HTML]5CAE00 0.757 [HTML]5CAE00 0.757 [HTML]F4000 0.333 [HTML]5CAE00 0.757 [HTML]5CAE00 0.757 [HTML]5CAE00 0.757 [HTML]7EBF00 0.715 [HTML]68B400 0.743 [HTML]5CAE00 0.757 [HTML]56AB00 0.764 [HTML]5CAE00 0.757 [HTML]56AB00 0.764 [HTML]56AB00 0.764 [HTML]62B100 0.750 24k-32k [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]88C400 0.703 [HTML]F1A00 0.287 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]C2E100 0.634 [HTML]88C400 0.703 [HTML]80C000 0.713 [HTML]88C400 0.703 [HTML]80C000 0.713 [HTML]88C400 0.703 [HTML]80C000 0.713 [HTML]78BC00 0.723 32k-49k [HTML]E6F300 0.591 [HTML]56AB00 0.765 [HTML]5CAE00 0.757 [HTML]F1200 0.278 [HTML]56AB00 0.765 [HTML]56AB00 0.765 [HTML]5CAE00 0.757 [HTML]80C000 0.713 [HTML]56AB00 0.765 [HTML]56AB00 0.765 [HTML]56AB00 0.765 [HTML]56AB00 0.765 [HTML]56AB00 0.765 [HTML]64B200 0.748 [HTML]56AB00 0.765 49k-65k [HTML]F3200 0.317 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]F5400 0.356 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]9ACD00 0.683 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]80C000 0.713 [HTML]78BC00 0.723 [HTML]80C000 0.713 [HTML]70B800 0.733 In-Domain (0k-8k) [HTML]3A9D00 0.797 [HTML]62B100 0.749 [HTML]62B100 0.749 [HTML]FFDE00 0.522 [HTML]62B100 0.749 [HTML]62B100 0.749 [HTML]62B100 0.749 [HTML]76B00 0.724 [HTML]6CB600 0.738 [HTML]62B100 0.749 [HTML]62B100 0.749 [HTML]62B100 0.749 [HTML]66B300 0.744 [HTML]62B100 0.749 [HTML]6EB700 0.735 OOD (8k-65k) [HTML]84C200 0.709 [HTML]68B400 0.741 [HTML]6CB600 0.737 [HTML]F3800 0.323 [HTML]68B400 0.741 [HTML]68B400 0.741 [HTML]6AB500 0.740 [HTML]96CB00 0.686 [HTML]6AB500 0.740 [HTML]68B400 0.741 [HTML]68B400 0.741 [HTML]68B400 0.741 [HTML]64B200 0.746 [HTML]6AB500 0.740 [HTML]6AB500 0.740 Long-OOD (32k-65k) [HTML]FFAC00 0.463 [HTML]6AB500 0.741 [HTML]6CB600 0.736 [HTML]F3000 0.314 [HTML]6AB500 0.741 [HTML]6AB500 0.741 [HTML]6CB600 0.736 [HTML]8C600 0.699 [HTML]6AB500 0.741 [HTML]6AB500 0.741 [HTML]6AB500 0.741 [HTML]6AB500 0.741 [HTML]66B300 0.745 [HTML]70B800 0.732 [HTML]62B100 0.750 Full (0k-65k) [HTML]6AB500 0.741 [HTML]66B300 0.744 [HTML]6AB500 0.741 [HTML]F7400 0.394 [HTML]66B300 0.744 [HTML]66B300 0.744 [HTML]68B400 0.743 [HTML]8C600 0.700 [HTML]6AB500 0.739 [HTML]66B300 0.744 [HTML]66B300 0.744 [HTML]66B300 0.744 [HTML]66B300 0.745 [HTML]68B400 0.743 [HTML]6CB600 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. We also compared the training dynamics during curriculum learning with top-4 associative blocks 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. Model/ Lengths 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 0k-1k [HTML]068300 0.380 [HTML]249200 0.357 [HTML]249200 0.357 [HTML]229100 0.358 [HTML]229100 0.358 [HTML]1E8F00 0.361 [HTML]188C00 0.366 [HTML]329900 0.347 [HTML]2C9600 0.351 [HTML]2A9500 0.353 [HTML]1C8E00 0.363 1k-2k [HTML]008000 0.385 [HTML]76B00 0.296 [HTML]78BC00 0.294 [HTML]6CB600 0.303 [HTML]72B900 0.299 [HTML]90C800 0.276 [HTML]84C200 0.285 [HTML]7CBE00 0.291 [HTML]7EBF00 0.290 [HTML]5AAD00 0.316 [HTML]5AAD00 0.317 2k-4k [HTML]2A9500 0.352 [HTML]A0D000 0.264 [HTML]92C900 0.275 [HTML]8C600 0.279 [HTML]84C200 0.285 [HTML]B0D800 0.252 [HTML]9ECF00 0.265 [HTML]9ECF00 0.265 [HTML]A2D100 0.263 [HTML]74BA00 0.297 [HTML]6AB500 0.304 4k-6k [HTML]369B00 0.344 [HTML]C0E000 0.240 [HTML]82C100 0.286 [HTML]6CB600 0.303 [HTML]8EC700 0.278 [HTML]DAED00 0.220 [HTML]A6D300 0.260 [HTML]A2D100 0.262 [HTML]C4E200 0.237 [HTML]8AC500 0.280 [HTML]82C100 0.287 6k-8k [HTML]76B00 0.296 [HTML]DCEE00 0.219 [HTML]C2E100 0.238 [HTML]AED700 0.254 [HTML]BADD00 0.244 [HTML]E6F300 0.212 [HTML]C2E100 0.239 [HTML]BEDF00 0.241 [HTML]FEFF00 0.194 [HTML]BEDF00 0.242 [HTML]B2D900 0.250 8k-10k [HTML]68B400 0.306 [HTML]EEF700 0.205 [HTML]CAE500 0.233 [HTML]C6E300 0.235 [HTML]CAE500 0.233 [HTML]F200 0.182 [HTML]CEE700 0.229 [HTML]CAE500 0.232 [HTML]FFA600 0.125 [HTML]C8E400 0.234 [HTML]CCE600 0.231 10k-12k [HTML]9ACD00 0.269 [HTML]E4F200 0.213 [HTML]DAED00 0.221 [HTML]BEDF00 0.241 [HTML]BEDF00 0.241 [HTML]FFFA00 0.188 [HTML]E0F000 0.216 [HTML]CAE500 0.232 [HTML]F8E00 0.108 [HTML]D4EA00 0.225 [HTML]CCE600 0.231 12k-14k [HTML]8AC500 0.280 [HTML]FEFF00 0.193 [HTML]D4EA00 0.225 [HTML]C4E200 0.237 [HTML]CEE700 0.230 [HTML]FFEC00 0.178 [HTML]D2E900 0.226 [HTML]D8EC00 0.222 [HTML]F5800 0.067 [HTML]EEF700 0.205 [HTML]DAED00 0.221 14k-16k [HTML]CCE600 0.231 [HTML]FFEE00 0.180 [HTML]DEEF00 0.217 [HTML]C0E000 0.240 [HTML]C4E200 0.237 [HTML]FFE600 0.174 [HTML]DEEF00 0.218 [HTML]CAE500 0.233 [HTML]F4E00 0.060 [HTML]EAF500 0.208 [HTML]BCDE00 0.243 16k-24k [HTML]9ACD00 0.269 [HTML]F9800 0.115 [HTML]F6FB00 0.200 [HTML]B0D800 0.252 [HTML]B6DB00 0.247 [HTML]FFEE00 0.179 [HTML]CAE500 0.232 [HTML]BEDF00 0.242 [HTML]F5200 0.062 [HTML]FFE000 0.169 [HTML]A6D300 0.259 24k-32k [HTML]B6DB00 0.247 [HTML]F1400 0.016 [HTML]FFDE00 0.168 [HTML]ECF600 0.207 [HTML]DEEF00 0.217 [HTML]FFD400 0.160 [HTML]FFFE00 0.192 [HTML]F800 0.187 [HTML]F4400 0.052 [HTML]F8C00 0.106 [HTML]F0F800 0.204 32k-49k [HTML]FFE200 0.171 [HTML]F0200 0.002 [HTML]FFAA00 0.128 [HTML]F6FB00 0.200 [HTML]CEE700 0.229 [HTML]FFB600 0.138 [HTML]F200 0.183 [HTML]EEF700 0.206 [HTML]F2000 0.025 [HTML]F7400 0.088 [HTML]EEF700 0.205 49k-65k [HTML]FFFA00 0.189 [HTML]F0000 0.000 [HTML]F3400 0.040 [HTML]A0D000 0.264 [HTML]46A300 0.331 [HTML]E8F400 0.210 [HTML]80C000 0.288 [HTML]3E9F00 0.338 [HTML]F1C00 0.022 [HTML]F4000 0.049 [HTML]B8DC00 0.246 In-Domain (0k-8k) [HTML]2C9600 0.351 [HTML]92C900 0.275 [HTML]7EBF00 0.290 [HTML]72B900 0.299 [HTML]7ABD00 0.293 [HTML]A0D000 0.264 [HTML]88C400 0.283 [HTML]8AC500 0.281 [HTML]9CCE00 0.267 [HTML]74BA00 0.297 [HTML]6AB500 0.304 OOD (8k-65k) [HTML]9ECF00 0.266 [HTML]FFD600 0.162 [HTML]E6F300 0.211 [HTML]C2E100 0.238 [HTML]C4E200 0.237 [HTML]FFEC00 0.178 [HTML]D8EC00 0.222 [HTML]CEE700 0.229 [HTML]F6800 0.079 [HTML]FCFE00 0.195 [HTML]C8E400 0.234 Long-OOD (32k-65k) [HTML]FFE800 0.175 [HTML]F0200 0.002 [HTML]F8E00 0.108 [HTML]E2F100 0.215 [HTML]B0D800 0.253 [HTML]FFCC00 0.155 [HTML]ECF600 0.207 [HTML]C4E200 0.236 [HTML]F2000 0.024 [HTML]F6800 0.079 [HTML]E2F100 0.214 Full (0k-65k) [HTML]68B400 0.306 [HTML]E2F100 0.215 [HTML]B6DB00 0.248 [HTML]9CCE00 0.266 [HTML]A2D100 0.263 [HTML]DEEF00 0.218 [HTML]B2D900 0.250 [HTML]AED700 0.254 [HTML]FFDC00 0.167 [HTML]BCDE00 0.243 [HTML]9CCE00 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 [HTML]2A9500 0.795 [HTML]329900 0.782 [HTML]329900 0.782 [HTML]42A100 0.756 [HTML]42A100 0.756 [HTML]329900 0.782 [HTML]329900 0.782 [HTML]42A100 0.756 [HTML]329900 0.782 [HTML]42A100 0.756 [HTML]42A100 0.756 1k-2k [HTML]2E9700 0.790 [HTML]54A00 0.724 [HTML]5AAD00 0.714 [HTML]48A400 0.743 [HTML]54A00 0.724 [HTML]5AAD00 0.714 [HTML]6AB500 0.686 [HTML]5AAD00 0.714 [HTML]54A00 0.724 [HTML]60B000 0.705 [HTML]60B000 0.705 2k-4k [HTML]369B00 0.774 [HTML]68B400 0.689 [HTML]64B200 0.698 [HTML]4CA600 0.736 [HTML]5EAF00 0.708 [HTML]7ABD00 0.660 [HTML]58AC00 0.717 [HTML]4CA600 0.736 [HTML]6EB700 0.679 [HTML]58AC00 0.717 [HTML]58AC00 0.717 4k-6k [HTML]3A9D00 0.767 [HTML]94CA00 0.616 [HTML]84C200 0.644 [HTML]62B100 0.699 [HTML]5CAE00 0.712 [HTML]F4FA00 0.452 [HTML]62B100 0.699 [HTML]62B100 0.699 [HTML]FFEA00 0.397 [HTML]62B100 0.699 [HTML]5CAE00 0.712 6k-8k [HTML]048200 0.859 [HTML]98C00 0.609 [HTML]58AC00 0.717 [HTML]249200 0.804 [HTML]2C9600 0.793 [HTML]F6C00 0.185 [HTML]2C9600 0.793 [HTML]1E8F00 0.815 [HTML]F4C00 0.130 [HTML]1E8F00 0.815 [HTML]128900 0.837 8k-10k [HTML]088400 0.852 [HTML]D4EA00 0.506 [HTML]42A100 0.753 [HTML]42A100 0.753 [HTML]3C9E00 0.765 [HTML]F3A00 0.099 [HTML]2E9700 0.790 [HTML]188C00 0.827 [HTML]F0600 0.012 [HTML]68B400 0.691 [HTML]108800 0.840 10k-12k [HTML]008000 0.868 [HTML]FFDC00 0.374 [HTML]74BA00 0.670 [HTML]2C9600 0.791 [HTML]40A000 0.758 [HTML]F0C00 0.022 [HTML]46A300 0.747 [HTML]269300 0.802 [HTML]F0000 0.000 [HTML]46A300 0.747 [HTML]068300 0.857 12k-14k [HTML]2E9700 0.787 [HTML]FFEE00 0.404 [HTML]60B000 0.702 [HTML]48A400 0.745 [HTML]42A100 0.755 [HTML]F1200 0.032 [HTML]68B400 0.691 [HTML]349A00 0.777 [HTML]F0000 0.000 [HTML]3C9E00 0.766 [HTML]349A00 0.777 14k-16k [HTML]3C9E00 0.764 [HTML]FFC600 0.337 [HTML]86C300 0.640 [HTML]6AB500 0.685 [HTML]78BC00 0.663 [HTML]F0600 0.011 [HTML]A6D300 0.584 [HTML]5EAF00 0.708 [HTML]F0000 0.000 [HTML]A6D300 0.584 [HTML]6AB500 0.685 16k-24k [HTML]209000 0.812 [HTML]F8600 0.229 [HTML]82C100 0.646 [HTML]40A000 0.757 [HTML]50A800 0.729 [HTML]F0000 0.000 [HTML]C0E000 0.542 [HTML]40A000 0.757 [HTML]F0000 0.000 [HTML]9ACD00 0.604 [HTML]209000 0.812 24k-32k [HTML]5AAD00 0.713 [HTML]F7400 0.198 [HTML]A6D300 0.584 [HTML]5AAD00 0.713 [HTML]66B300 0.693 [HTML]F0000 0.000 [HTML]FFE800 0.396 [HTML]3E9F00 0.762 [HTML]F0000 0.000 [HTML]FEFF00 0.436 [HTML]3E9F00 0.762 32k-49k [HTML]A2D100 0.591 [HTML]F5000 0.139 [HTML]98C00 0.609 [HTML]3C9E00 0.765 [HTML]3C9E00 0.765 [HTML]F0000 0.000 [HTML]F9800 0.261 [HTML]3C9E00 0.765 [HTML]F0000 0.000 [HTML]FEFF00 0.435 [HTML]289400 0.800 49k-65k [HTML]FFBA00 0.317 [HTML]F3400 0.089 [HTML]9ACD00 0.604 [HTML]5AAD00 0.713 [HTML]5AAD00 0.713 [HTML]F0000 0.000 [HTML]FFA200 0.277 [HTML]44A200 0.752 [HTML]F0000 0.000 [HTML]E0F000 0.485 [HTML]269300 0.802 In-Domain (0k-8k) [HTML]289400 0.797 [HTML]6AB500 0.685 [HTML]5CAE00 0.711 [HTML]46A300 0.749 [HTML]4CA600 0.738 [HTML]B2D900 0.564 [HTML]4EA700 0.734 [HTML]48A400 0.744 [HTML]BADD00 0.551 [HTML]4CA600 0.738 [HTML]48A400 0.744 OOD (8k-65k) [HTML]5CAE00 0.709 [HTML]F9E00 0.271 [HTML]82C100 0.647 [HTML]4A500 0.741 [HTML]50A800 0.730 [HTML]F0A00 0.017 [HTML]CCE600 0.521 [HTML]3A9D00 0.767 [HTML]F0000 0.001 [HTML]A6D300 0.586 [HTML]2C9600 0.793 Long-OOD (32k-65k) [HTML]EEF700 0.463 [HTML]F4400 0.116 [HTML]9ACD00 0.607 [HTML]4A500 0.741 [HTML]4A500 0.741 [HTML]F0000 0.000 [HTML]F9E00 0.268 [HTML]40A000 0.759 [HTML]F0000 0.000 [HTML]F0F800 0.458 [HTML]269300 0.801 Full (0k-65k) [HTML]4A500 0.741 [HTML]F600 0.419 [HTML]74BA00 0.670 [HTML]48A400 0.744 [HTML]4EA700 0.733 [HTML]F7C00 0.213 [HTML]9ECF00 0.597 [HTML]40A000 0.759 [HTML]F7400 0.198 [HTML]86C300 0.640 [HTML]369B00 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. D.4 Ablation Study for Synthetic Training Data Generation Figure 5: Scaling synthetic data for the GR dataset using ARMT-Gemma-1B-IT (ARMT) and Gemma-1B-IT (Base), trained up to context length 8192. The best performance is achieved with the GR-100+ dataset. 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. Model/ Lengths Base, No Fine- Tuning Base, GR, 8k ARMT, GR, 2k ARMT, GR, 4k ARMT, GR, 8k 0k-1k [HTML]FFCE00 0.188 [HTML]128900 0.347 [HTML]209000 0.340 [HTML]389C00 0.326 [HTML]50A800 0.313 1k-2k [HTML]F9E00 0.161 [HTML]008000 0.358 [HTML]FFB800 0.175 [HTML]FFE000 0.198 [HTML]F800 0.211 2k-4k [HTML]F7800 0.140 [HTML]5EAF00 0.305 [HTML]F8C00 0.151 [HTML]FFC800 0.184 [HTML]FFE400 0.200 4k-6k [HTML]F4A00 0.114 [HTML]7CBE00 0.288 [HTML]F7000 0.135 [HTML]FFB800 0.175 [HTML]FFC000 0.180 6k-8k [HTML]F4800 0.113 [HTML]7ABD00 0.289 [HTML]F8400 0.146 [HTML]FFA400 0.164 [HTML]FFBA00 0.176 8k-10k [HTML]F3C00 0.106 [HTML]9ACD00 0.271 [HTML]F8400 0.146 [HTML]F8E00 0.152 [HTML]FFA400 0.164 10k-12k [HTML]F3200 0.101 [HTML]E0F000 0.232 [HTML]F6A00 0.132 [HTML]FFA800 0.166 [HTML]FFBA00 0.177 12k-14k [HTML]F2200 0.091 [HTML]D0E800 0.241 [HTML]F5400 0.120 [HTML]F8600 0.147 [HTML]F9600 0.156 14k-16k [HTML]F1C00 0.088 [HTML]FAFD00 0.218 [HTML]F6400 0.128 [HTML]F8800 0.149 [HTML]FFAA00 0.168 16k-24k [HTML]F1200 0.083 [HTML]F2F900 0.222 [HTML]F6400 0.128 [HTML]F8200 0.145 [HTML]FFA400 0.164 24k-32k [HTML]F0A00 0.078 [HTML]F8800 0.149 [HTML]F4400 0.111 [HTML]F7600 0.138 [HTML]FFB800 0.175 32k-49k [HTML]F0A00 0.078 [HTML]F6600 0.130 [HTML]F5A00 0.123 [HTML]F7000 0.135 [HTML]F8000 0.144 49k-65k [HTML]F3600 0.103 [HTML]F0000 0.072 [HTML]FFBC00 0.178 [HTML]FFD200 0.190 [HTML]FFE600 0.201 In-Domain (0k-8k) [HTML]F7E00 0.143 [HTML]48A400 0.317 [HTML]FFD000 0.189 [HTML]F400 0.209 [HTML]FEFF00 0.216 OOD (8k-65k) [HTML]F2400 0.092 [HTML]ECF600 0.226 [HTML]F6600 0.130 [HTML]F8C00 0.150 [HTML]FFA600 0.165 Long-OOD (32k-65k) [HTML]F1400 0.084 [HTML]F4E00 0.117 [HTML]F7200 0.136 [HTML]F8600 0.148 [HTML]F9800 0.157 Full (0k-65k) [HTML]F4E00 0.116 [HTML]9ECF00 0.269 [HTML]F9800 0.157 [HTML]FFBC00 0.178 [HTML]FFD000 0.189 Table 21: Best results on the GR dataset for Gemma-3-1B-IT model, metric - ROUGE-L. 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 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 learning setup. Due to limited resources, we pre-train the model only on the two segments of 1024 tokens. The results are presented in Table˜22. Model/ Lengths Base, No Fine- Tuning Base, GR-100+, 8k ARMT, GR-100+, 8k ARMT, GR-100+, 8k, Pretrained 0k-1k [HTML]FFEC00 0.188 [HTML]068300 0.380 [HTML]249200 0.358 [HTML]1E8F00 0.363 1k-2k [HTML]FFC600 0.161 [HTML]008000 0.385 [HTML]72B900 0.303 [HTML]48A400 0.332 2k-4k [HTML]FFAA00 0.140 [HTML]2E9700 0.352 [HTML]92C900 0.279 [HTML]8EC700 0.282 4k-6k [HTML]F8400 0.114 [HTML]389C00 0.344 [HTML]72B900 0.303 [HTML]8C600 0.284 6k-8k [HTML]F8400 0.113 [HTML]7CBE00 0.296 [HTML]B6DB00 0.254 [HTML]94CA00 0.278 8k-10k [HTML]F7A00 0.106 [HTML]6EB700 0.306 [HTML]D0E800 0.235 [HTML]C4E200 0.244 10k-12k [HTML]F7200 0.101 [HTML]A0D000 0.269 [HTML]C8E400 0.241 [HTML]BADD00 0.251 12k-14k [HTML]F6400 0.091 [HTML]92C900 0.280 [HTML]CEE700 0.237 [HTML]C6E300 0.242 14k-16k [HTML]F6000 0.088 [HTML]D6EB00 0.231 [HTML]CAE500 0.240 [HTML]B8DC00 0.252 16k-24k [HTML]F5A00 0.083 [HTML]A0D000 0.269 [HTML]B8DC00 0.252 [HTML]BCDE00 0.249 24k-32k [HTML]F5200 0.078 [HTML]C0E000 0.247 [HTML]F8FC00 0.207 [HTML]EEF700 0.214 32k-49k [HTML]F5200 0.078 [HTML]FFD400 0.171 [HTML]FFFC00 0.200 [HTML]FFE600 0.183 49k-65k [HTML]F7600 0.103 [HTML]FFEE00 0.189 [HTML]A8D400 0.264 [HTML]F0000 0.018 In-Domain (0k-8k) [HTML]FFAE00 0.143 [HTML]2E9700 0.351 [HTML]76B00 0.299 [HTML]6CB600 0.307 OOD (8k-65k) [HTML]F6600 0.092 [HTML]A6D300 0.266 [HTML]CCE600 0.238 [HTML]C8E400 0.241 Long-OOD (32k-65k) [HTML]F5A00 0.084 [HTML]FFDA00 0.175 [HTML]ECF600 0.215 [HTML]FFB000 0.145 Full (0k-65k) [HTML]F8800 0.116 [HTML]6EB700 0.306 [HTML]A4D200 0.266 [HTML]9CCE00 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. 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. ((a)) MT results. ((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 SmolLM-2-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. D.6 Ablation Study for ARMT Backbone Model/ Lengths Base, No Fine- Tuning Base, GR-100+, 8k ARMT, GR-100+, 2k ARMT, GR-100+, 4k ARMT, GR-100+, 8k 0k-1k [HTML]FFC400 0.145 [HTML]0C8600 0.349 [HTML]209000 0.335 [HTML]269300 0.332 [HTML]1E8F00 0.337 1k-2k [HTML]FFA200 0.122 [HTML]008000 0.358 [HTML]64B200 0.289 [HTML]5AAD00 0.297 [HTML]52A900 0.302 2k-4k [HTML]F9600 0.113 [HTML]289400 0.330 [HTML]B0D800 0.238 [HTML]94CA00 0.257 [HTML]8EC700 0.261 4k-6k [HTML]F7600 0.091 [HTML]40A000 0.314 [HTML]CAE500 0.220 [HTML]A6D300 0.245 [HTML]9ACD00 0.253 6k-8k [HTML]F7A00 0.094 [HTML]6AB500 0.286 [HTML]D6EB00 0.212 [HTML]D2E900 0.215 [HTML]C4E200 0.225 8k-10k [HTML]F7200 0.089 [HTML]B8DC00 0.233 [HTML]F400 0.177 [HTML]E6F300 0.202 [HTML]CAE500 0.221 10k-12k [HTML]F2800 0.039 [HTML]F0E00 0.021 [HTML]FFD400 0.156 [HTML]EAF500 0.199 [HTML]D8EC00 0.211 12k-14k [HTML]F2800 0.039 [HTML]F2200 0.035 [HTML]FFB800 0.137 [HTML]F0F800 0.195 [HTML]D2E900 0.215 14k-16k [HTML]F2800 0.039 [HTML]F3800 0.049 [HTML]F5C00 0.074 [HTML]F6FB00 0.190 [HTML]DAED00 0.210 16k-24k [HTML]F3600 0.048 [HTML]F2600 0.037 [HTML]F3200 0.046 [HTML]FFBE00 0.140 [HTML]CCE600 0.219 24k-32k [HTML]F2C00 0.042 [HTML]F3E00 0.054 [HTML]F0000 0.011 [HTML]F4600 0.059 [HTML]FEFF00 0.185 32k-49k [HTML]F3000 0.044 [HTML]F5200 0.067 [HTML]F1400 0.025 [HTML]F1E00 0.032 [HTML]FFEC00 0.171 49k-65k [HTML]F7000 0.088 [HTML]F7000 0.088 [HTML]F0400 0.015 [HTML]F0000 0.011 [HTML]F7C00 0.096 In-Domain (0k-8k) [HTML]F9600 0.113 [HTML]2C9600 0.327 [HTML]92C900 0.258 [HTML]82C100 0.269 [HTML]7ABD00 0.275 OOD (8k-65k) [HTML]F3A00 0.050 [HTML]F5C00 0.074 [HTML]F8C00 0.106 [HTML]FFEA00 0.170 [HTML]D8EC00 0.211 Long-OOD (32k-65k) [HTML]F3E00 0.054 [HTML]F5800 0.072 [HTML]F1000 0.023 [HTML]F1600 0.027 [HTML]FFD200 0.154 Full (0k-65k) [HTML]F6400 0.080 [HTML]F4FA00 0.192 [HTML]F400 0.178 [HTML]D0E800 0.216 [HTML]ACD600 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 Base, No Fine- Tuning Base, MT, 8k ARMT, MT, 2k ARMT, MT, 4k ARMT, MT, 8k 0k-1k [HTML]F0000 0.000 [HTML]188C00 0.795 [HTML]48A400 0.718 [HTML]309800 0.756 [HTML]289400 0.769 1k-2k [HTML]F0000 0.000 [HTML]1C8E00 0.790 [HTML]78BC00 0.638 [HTML]5CAE00 0.686 [HTML]50A800 0.705 2k-4k [HTML]F0000 0.000 [HTML]3C9E00 0.736 [HTML]76B00 0.642 [HTML]48A400 0.717 [HTML]54A00 0.698 4k-6k [HTML]F0000 0.000 [HTML]5CAE00 0.685 [HTML]8EC700 0.603 [HTML]7EBF00 0.630 [HTML]7EBF00 0.630 6k-8k [HTML]F0000 0.000 [HTML]008000 0.837 [HTML]92C900 0.598 [HTML]3A9D00 0.739 [HTML]2E9700 0.761 8k-10k [HTML]F0000 0.000 [HTML]A4D200 0.568 [HTML]E0F000 0.469 [HTML]148A00 0.802 [HTML]0C8600 0.815 10k-12k [HTML]F0600 0.011 [HTML]F0000 0.000 [HTML]FFFE00 0.418 [HTML]58AC00 0.692 [HTML]369B00 0.747 12k-14k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F5400 0.138 [HTML]329900 0.755 [HTML]329900 0.755 14k-16k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]6AB500 0.663 [HTML]7EBF00 0.629 16k-24k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0800 0.014 [HTML]4EA700 0.708 [HTML]4A500 0.715 24k-32k [HTML]F0C00 0.020 [HTML]F0600 0.010 [HTML]F0000 0.000 [HTML]F9C00 0.257 [HTML]3E9F00 0.733 32k-49k [HTML]F0400 0.009 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]40A000 0.730 49k-65k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F1200 0.030 In-Domain (0k-8k) [HTML]F0000 0.000 [HTML]289400 0.771 [HTML]78BC00 0.639 [HTML]4EA700 0.707 [HTML]4A500 0.714 OOD (8k-65k) [HTML]F0200 0.005 [HTML]F2200 0.058 [HTML]F4400 0.112 [HTML]DEEF00 0.473 [HTML]76B00 0.643 Long-OOD (32k-65k) [HTML]F0200 0.005 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F600 0.403 Full (0k-65k) [HTML]F0000 0.003 [HTML]FFBE00 0.313 [HTML]FFB600 0.300 [HTML]AAD500 0.557 [HTML]66B300 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 0k-1k [HTML]0C8600 0.349 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]249200 0.334 [HTML]289400 0.331 [HTML]2A9500 0.330 [HTML]269300 0.332 1k-2k [HTML]008000 0.358 [HTML]54A00 0.302 [HTML]56AB00 0.301 [HTML]58AC00 0.299 [HTML]FCFE00 0.192 [HTML]E2F100 0.208 [HTML]FFB600 0.142 [HTML]F4FA00 0.197 2k-4k [HTML]2A9500 0.330 [HTML]92C900 0.261 [HTML]AAD500 0.246 [HTML]96CB00 0.258 [HTML]FFE800 0.175 [HTML]FFE800 0.175 [HTML]FFAA00 0.134 [HTML]F600 0.183 4k-6k [HTML]42A100 0.314 [HTML]9ECF00 0.253 [HTML]CAE500 0.224 [HTML]BCDE00 0.234 [HTML]FFE800 0.175 [HTML]FFEE00 0.178 [HTML]F8400 0.109 [HTML]FFCA00 0.155 6k-8k [HTML]6CB600 0.286 [HTML]CAE500 0.225 [HTML]FFFC00 0.187 [HTML]D4EA00 0.218 [HTML]FFD000 0.159 [HTML]FFEC00 0.177 [HTML]F9600 0.120 [HTML]FFB800 0.143 8k-10k [HTML]BCDE00 0.233 [HTML]D0E800 0.221 [HTML]FFDC00 0.166 [HTML]E6F300 0.206 [HTML]FFC800 0.153 [HTML]FFD600 0.163 [HTML]F9400 0.119 [HTML]FFBC00 0.146 10k-12k [HTML]F0000 0.021 [HTML]DEEF00 0.211 [HTML]FFC600 0.152 [HTML]DEEF00 0.211 [HTML]FFCC00 0.156 [HTML]FFE800 0.174 [HTML]F8C00 0.114 [HTML]FFC800 0.153 12k-14k [HTML]F1400 0.035 [HTML]D8EC00 0.215 [HTML]F8400 0.109 [HTML]EEF700 0.201 [HTML]FFC400 0.151 [HTML]FFD000 0.159 [HTML]F8A00 0.113 [HTML]FFBC00 0.146 14k-16k [HTML]F2A00 0.049 [HTML]E0F000 0.210 [HTML]F6A00 0.092 [HTML]F0F800 0.199 [HTML]FFBA00 0.144 [HTML]FFC800 0.153 [HTML]F9600 0.121 [HTML]FFBC00 0.146 16k-24k [HTML]F1800 0.037 [HTML]D2E900 0.219 [HTML]F6000 0.085 [HTML]FCFE00 0.192 [HTML]FFBC00 0.146 [HTML]FFCA00 0.155 [HTML]F8A00 0.113 [HTML]FFB000 0.138 24k-32k [HTML]F3200 0.054 [HTML]F800 0.185 [HTML]F4400 0.066 [HTML]FFDE00 0.168 [HTML]FFB400 0.140 [HTML]FFC400 0.151 [HTML]F8C00 0.114 [HTML]FFB600 0.142 32k-49k [HTML]F4400 0.067 [HTML]FFE200 0.171 [HTML]F4200 0.065 [HTML]FFC800 0.153 [HTML]FFBE00 0.147 [HTML]FFB600 0.141 [HTML]F9A00 0.123 [HTML]FFAC00 0.135 49k-65k [HTML]F6400 0.088 [HTML]F7000 0.096 [HTML]F2A00 0.049 [HTML]F9000 0.116 [HTML]FEFF00 0.190 [HTML]F6800 0.090 [HTML]F8A00 0.113 [HTML]FFD000 0.158 In-Domain (0k-8k) [HTML]2E9700 0.327 [HTML]7CBE00 0.275 [HTML]96CB00 0.259 [HTML]86C300 0.269 [HTML]E6F300 0.206 [HTML]DCEE00 0.213 [HTML]FFDC00 0.166 [HTML]ECF600 0.202 OOD (8k-65k) [HTML]F4E00 0.074 [HTML]DEEF00 0.211 [HTML]F8C00 0.114 [HTML]F4FA00 0.197 [HTML]FFC200 0.150 [HTML]FFD000 0.158 [HTML]F9000 0.116 [HTML]FFBC00 0.145 Long-OOD (32k-65k) [HTML]F4C00 0.072 [HTML]FFC800 0.154 [HTML]F3C00 0.061 [HTML]FFBA00 0.144 [HTML]FFCE00 0.157 [HTML]FFA400 0.129 [HTML]F9600 0.121 [HTML]FFB400 0.140 Full (0k-65k) [HTML]FAFD00 0.192 [HTML]B0D800 0.241 [HTML]F200 0.182 [HTML]C0E000 0.231 [HTML]FFEA00 0.176 [HTML]F600 0.184 [HTML]FFB400 0.140 [HTML]FFE400 0.171 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 [HTML]188C00 0.795 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]309800 0.756 [HTML]309800 0.756 [HTML]309800 0.756 [HTML]309800 0.756 1k-2k [HTML]1C8E00 0.790 [HTML]50A800 0.705 [HTML]56AB00 0.695 [HTML]56AB00 0.695 [HTML]F4FA00 0.438 [HTML]84C200 0.619 [HTML]A2D100 0.571 [HTML]78BC00 0.638 2k-4k [HTML]3C9E00 0.736 [HTML]54A00 0.698 [HTML]4EA700 0.708 [HTML]54A00 0.698 [HTML]FFEC00 0.387 [HTML]70B800 0.651 [HTML]88C400 0.613 [HTML]82C100 0.623 4k-6k [HTML]5CAE00 0.685 [HTML]7EBF00 0.630 [HTML]64B200 0.671 [HTML]6CB600 0.658 [HTML]FFD000 0.342 [HTML]86C300 0.616 [HTML]B0D800 0.548 [HTML]6CB600 0.658 6k-8k [HTML]008000 0.837 [HTML]2E9700 0.761 [HTML]48A400 0.717 [HTML]2E9700 0.761 [HTML]FFC000 0.315 [HTML]349A00 0.750 [HTML]92C900 0.598 [HTML]349A00 0.750 8k-10k [HTML]A4D200 0.568 [HTML]0C8600 0.815 [HTML]1C8E00 0.790 [HTML]249200 0.778 [HTML]FFC400 0.321 [HTML]4A500 0.716 [HTML]C2E100 0.519 [HTML]58AC00 0.691 10k-12k [HTML]F0000 0.000 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]289400 0.769 [HTML]FFE400 0.374 [HTML]4A500 0.714 [HTML]E4F200 0.462 [HTML]44A200 0.725 12k-14k [HTML]F0000 0.000 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]4A500 0.713 [HTML]F9A00 0.255 [HTML]3E9F00 0.734 [HTML]DAED00 0.479 [HTML]44A200 0.723 14k-16k [HTML]F0000 0.000 [HTML]7EBF00 0.629 [HTML]84C200 0.618 [HTML]70B800 0.652 [HTML]FFB200 0.292 [HTML]A0D000 0.573 [HTML]FFDC00 0.360 [HTML]A8D400 0.562 16k-24k [HTML]F0000 0.000 [HTML]4A500 0.715 [HTML]4A500 0.715 [HTML]52A900 0.701 [HTML]F9400 0.243 [HTML]70B800 0.653 [HTML]DAED00 0.479 [HTML]70B800 0.653 24k-32k [HTML]F0600 0.010 [HTML]3E9F00 0.733 [HTML]50A800 0.703 [HTML]5EAF00 0.683 [HTML]FFAE00 0.287 [HTML]76B00 0.644 [HTML]FFE600 0.376 [HTML]8EC700 0.604 32k-49k [HTML]F0000 0.000 [HTML]40A000 0.730 [HTML]60B000 0.678 [HTML]40A000 0.730 [HTML]F8E00 0.235 [HTML]7ABD00 0.635 [HTML]E6F300 0.461 [HTML]80C000 0.626 49k-65k [HTML]F0000 0.000 [HTML]F1200 0.030 [HTML]BEDF00 0.525 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]DCEE00 0.475 [HTML]76B00 0.644 In-Domain (0k-8k) [HTML]289400 0.771 [HTML]4A500 0.714 [HTML]4CA600 0.711 [HTML]4A500 0.716 [HTML]F2F900 0.440 [HTML]62B100 0.676 [HTML]88C400 0.614 [HTML]5EAF00 0.681 OOD (8k-65k) [HTML]F2200 0.058 [HTML]76B00 0.643 [HTML]5AAD00 0.689 [HTML]7CBE00 0.631 [HTML]F9800 0.250 [HTML]98C00 0.586 [HTML]EAF500 0.452 [HTML]70B800 0.652 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]F600 0.403 [HTML]8C600 0.606 [HTML]F600 0.403 [HTML]F5400 0.139 [HTML]FFD600 0.352 [HTML]E0F000 0.468 [HTML]7ABD00 0.634 Full (0k-65k) [HTML]FFBE00 0.313 [HTML]66B300 0.668 [HTML]54A00 0.697 [HTML]6AB500 0.661 [HTML]FFC200 0.318 [HTML]84C200 0.618 [HTML]C6E300 0.510 [HTML]6AB500 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. 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-2-360M-IT model, the results are presented in Tables˜27, 28, 29 and 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 23-rd is the most important on MT. However, the top-4 associative layers by importance are 0-th, 23-rd, 30-th, 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. Model/ Lengths 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 0k-1k [HTML]0C8600 0.349 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 1k-2k [HTML]008000 0.358 [HTML]54A00 0.302 [HTML]5AAD00 0.298 [HTML]56AB00 0.301 [HTML]52A900 0.303 [HTML]54A00 0.302 [HTML]5AAD00 0.298 [HTML]58AC00 0.299 [HTML]52A900 0.304 [HTML]66B300 0.290 [HTML]66B300 0.290 [HTML]5AAD00 0.298 [HTML]5AAD00 0.298 [HTML]64B200 0.292 [HTML]5AAD00 0.298 [HTML]54A00 0.302 [HTML]58AC00 0.299 [HTML]5EAF00 0.295 2k-4k [HTML]2A9500 0.330 [HTML]92C900 0.261 [HTML]B0D800 0.241 [HTML]96CB00 0.259 [HTML]9ECF00 0.253 [HTML]9ECF00 0.254 [HTML]9ACD00 0.256 [HTML]A0D000 0.252 [HTML]96CB00 0.258 [HTML]94CA00 0.260 [HTML]9ACD00 0.256 [HTML]9ACD00 0.256 [HTML]96CB00 0.259 [HTML]9ACD00 0.256 [HTML]96CB00 0.259 [HTML]94CA00 0.260 [HTML]90C800 0.263 [HTML]92C900 0.261 4k-6k [HTML]42A100 0.314 [HTML]9ECF00 0.253 [HTML]B6DB00 0.237 [HTML]98C00 0.257 [HTML]9CCE00 0.255 [HTML]A0D000 0.252 [HTML]98C00 0.257 [HTML]96CB00 0.259 [HTML]94CA00 0.260 [HTML]A0D000 0.252 [HTML]A4D200 0.250 [HTML]A6D300 0.248 [HTML]AED700 0.243 [HTML]9ECF00 0.254 [HTML]96CB00 0.259 [HTML]92C900 0.261 [HTML]9CCE00 0.255 [HTML]A4D200 0.249 6k-8k [HTML]6CB600 0.286 [HTML]CAE500 0.225 [HTML]EEF700 0.201 [HTML]D0E800 0.221 [HTML]CCE600 0.223 [HTML]D6EB00 0.217 [HTML]CAE500 0.225 [HTML]D0E800 0.220 [HTML]C4E200 0.228 [HTML]D2E900 0.219 [HTML]CCE600 0.223 [HTML]D4EA00 0.218 [HTML]D2E900 0.219 [HTML]D0E800 0.221 [HTML]D0E800 0.220 [HTML]CAE500 0.224 [HTML]CCE600 0.223 [HTML]CAE500 0.225 8k-10k [HTML]BCDE00 0.233 [HTML]D0E800 0.221 [HTML]FCFE00 0.192 [HTML]D2E900 0.219 [HTML]D6EB00 0.216 [HTML]CAE500 0.224 [HTML]D4EA00 0.218 [HTML]CEE700 0.222 [HTML]C2E100 0.229 [HTML]D4EA00 0.218 [HTML]CEE700 0.222 [HTML]DAED00 0.214 [HTML]D4EA00 0.218 [HTML]DCEE00 0.213 [HTML]D6EB00 0.216 [HTML]CEE700 0.222 [HTML]D2E900 0.219 [HTML]D6EB00 0.217 10k-12k [HTML]F0000 0.021 [HTML]DEEF00 0.211 [HTML]FFE000 0.169 [HTML]DAED00 0.214 [HTML]D8EC00 0.215 [HTML]E0F000 0.210 [HTML]D8EC00 0.215 [HTML]D6EB00 0.217 [HTML]D6EB00 0.216 [HTML]DAED00 0.214 [HTML]EAF500 0.203 [HTML]D0E800 0.221 [HTML]E2F100 0.209 [HTML]DEEF00 0.211 [HTML]D8EC00 0.215 [HTML]D6EB00 0.217 [HTML]D2E900 0.219 [HTML]DAED00 0.214 12k-14k [HTML]F1400 0.035 [HTML]D8EC00 0.215 [HTML]FFCA00 0.155 [HTML]D6EB00 0.217 [HTML]D0E800 0.220 [HTML]DAED00 0.214 [HTML]D8EC00 0.215 [HTML]E0F000 0.210 [HTML]CCE600 0.223 [HTML]DCEE00 0.213 [HTML]DCEE00 0.213 [HTML]D2E900 0.219 [HTML]E2F100 0.209 [HTML]DCEE00 0.213 [HTML]D2E900 0.219 [HTML]D6EB00 0.217 [HTML]DCEE00 0.213 [HTML]D2E900 0.219 14k-16k [HTML]F2A00 0.049 [HTML]E0F000 0.210 [HTML]FFA600 0.131 [HTML]DAED00 0.214 [HTML]D4EA00 0.218 [HTML]DEEF00 0.211 [HTML]D4EA00 0.218 [HTML]DCEE00 0.212 [HTML]DAED00 0.214 [HTML]DEEF00 0.211 [HTML]DCEE00 0.212 [HTML]D6EB00 0.216 [HTML]D6EB00 0.217 [HTML]DEEF00 0.211 [HTML]D4EA00 0.218 [HTML]E0F000 0.210 [HTML]D2E900 0.219 [HTML]DEEF00 0.211 16k-24k [HTML]F1800 0.037 [HTML]D2E900 0.219 [HTML]F8200 0.107 [HTML]D4EA00 0.218 [HTML]DCEE00 0.213 [HTML]D8EC00 0.215 [HTML]DCEE00 0.212 [HTML]D8EC00 0.215 [HTML]CCE600 0.223 [HTML]D6EB00 0.216 [HTML]DCEE00 0.213 [HTML]D4EA00 0.218 [HTML]E8F400 0.205 [HTML]D4EA00 0.218 [HTML]D6EB00 0.217 [HTML]D0E800 0.220 [HTML]C6E300 0.227 [HTML]CEE700 0.222 24k-32k [HTML]F3200 0.054 [HTML]F800 0.185 [HTML]F4600 0.068 [HTML]F0F800 0.199 [HTML]FCFE00 0.192 [HTML]F2F900 0.198 [HTML]F6FB00 0.196 [HTML]FFFC00 0.188 [HTML]EEF700 0.201 [HTML]EAF500 0.203 [HTML]E8F400 0.205 [HTML]F600 0.183 [HTML]FFDC00 0.166 [HTML]F6FB00 0.196 [HTML]FFFE00 0.189 [HTML]F6FB00 0.196 [HTML]F8FC00 0.194 [HTML]F600 0.183 32k-49k [HTML]F4400 0.067 [HTML]FFE200 0.171 [HTML]F3E00 0.063 [HTML]FCFE00 0.192 [HTML]FFFE00 0.189 [HTML]FFC600 0.152 [HTML]F200 0.181 [HTML]FFE600 0.173 [HTML]FFE800 0.175 [HTML]FFE000 0.169 [HTML]F200 0.181 [HTML]FFFE00 0.189 [HTML]FFC000 0.148 [HTML]FFE200 0.170 [HTML]F400 0.182 [HTML]FFE800 0.174 [HTML]F2F900 0.198 [HTML]FFCA00 0.154 49k-65k [HTML]F6400 0.088 [HTML]F7000 0.096 [HTML]F3E00 0.062 [HTML]FFD000 0.159 [HTML]F7600 0.099 [HTML]F5200 0.076 [HTML]F7400 0.098 [HTML]F9A00 0.123 [HTML]FFAE00 0.136 [HTML]F7000 0.095 [HTML]F5000 0.074 [HTML]F9200 0.118 [HTML]F5800 0.080 [HTML]F6000 0.085 [HTML]F6C00 0.093 [HTML]FFA200 0.128 [HTML]F5A00 0.081 [HTML]F4E00 0.073 In-Domain (0k-8k) [HTML]2E9700 0.327 [HTML]7CBE00 0.275 [HTML]90C800 0.262 [HTML]7EBF00 0.275 [HTML]7EBF00 0.274 [HTML]82C100 0.272 [HTML]7EBF00 0.274 [HTML]80C000 0.273 [HTML]7ABD00 0.277 [HTML]82C100 0.271 [HTML]84C200 0.271 [HTML]84C200 0.271 [HTML]84C200 0.271 [HTML]82C100 0.272 [HTML]7EBF00 0.274 [HTML]7ABD00 0.276 [HTML]7CBE00 0.275 [HTML]80C000 0.273 OOD (8k-65k) [HTML]F4E00 0.074 [HTML]DEEF00 0.211 [HTML]FFB600 0.141 [HTML]DAED00 0.214 [HTML]DCEE00 0.213 [HTML]E0F000 0.210 [HTML]DEEF00 0.212 [HTML]DEEF00 0.211 [HTML]D4EA00 0.217 [HTML]DEEF00 0.211 [HTML]E0F000 0.210 [HTML]DAED00 0.213 [HTML]E8F400 0.205 [HTML]E0F000 0.209 [HTML]DCEE00 0.213 [HTML]DAED00 0.213 [HTML]D8EC00 0.216 [HTML]DEEF00 0.211 Long-OOD (32k-65k) [HTML]F4C00 0.072 [HTML]FFC800 0.154 [HTML]F3E00 0.063 [HTML]F800 0.184 [HTML]FFDE00 0.168 [HTML]FFAC00 0.134 [HTML]FFD400 0.162 [HTML]FFD400 0.161 [HTML]FFDC00 0.166 [HTML]FFC600 0.152 [HTML]FFCC00 0.156 [HTML]FFE600 0.173 [HTML]FFA800 0.132 [HTML]FFC400 0.150 [HTML]FFD400 0.161 [HTML]FFD800 0.163 [HTML]FFE200 0.171 [HTML]FFAC00 0.135 Full (0k-65k) [HTML]FAFD00 0.192 [HTML]B0D800 0.241 [HTML]F2F900 0.198 [HTML]AED700 0.242 [HTML]B0D800 0.241 [HTML]B4DA00 0.239 [HTML]B0D800 0.241 [HTML]B2D900 0.240 [HTML]AAD500 0.245 [HTML]B4DA00 0.239 [HTML]B4DA00 0.238 [HTML]B2D900 0.240 [HTML]B8DC00 0.236 [HTML]B4DA00 0.239 [HTML]B0D800 0.242 [HTML]AED700 0.243 [HTML]AED700 0.243 [HTML]B2D900 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 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 0k-1k [HTML]0C8600 0.349 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]229100 0.335 [HTML]1E8F00 0.337 [HTML]209000 0.336 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 [HTML]1E8F00 0.338 [HTML]209000 0.336 [HTML]1E8F00 0.337 [HTML]1E8F00 0.337 1k-2k [HTML]008000 0.358 [HTML]54A00 0.302 [HTML]5EAF00 0.296 [HTML]5EAF00 0.296 [HTML]5EAF00 0.295 [HTML]6AB500 0.287 [HTML]6EB700 0.285 [HTML]6CB600 0.286 [HTML]80C000 0.273 [HTML]B4DA00 0.239 [HTML]6AB500 0.288 [HTML]5AAD00 0.298 [HTML]52A900 0.303 [HTML]7ABD00 0.277 [HTML]5AAD00 0.298 [HTML]74BA00 0.281 [HTML]8AC500 0.267 [HTML]64B200 0.291 2k-4k [HTML]2A9500 0.330 [HTML]92C900 0.261 [HTML]A2D100 0.251 [HTML]9ACD00 0.256 [HTML]94CA00 0.260 [HTML]98C00 0.257 [HTML]96CB00 0.258 [HTML]96CB00 0.259 [HTML]9CCE00 0.255 [HTML]DCEE00 0.212 [HTML]B0D800 0.241 [HTML]94CA00 0.260 [HTML]90C800 0.262 [HTML]AAD500 0.246 [HTML]9ACD00 0.256 [HTML]A6D300 0.248 [HTML]B8DC00 0.236 [HTML]8EC700 0.264 4k-6k [HTML]42A100 0.314 [HTML]9ECF00 0.253 [HTML]9ECF00 0.253 [HTML]9ECF00 0.253 [HTML]9ECF00 0.253 [HTML]B4DA00 0.239 [HTML]B2D900 0.240 [HTML]9ECF00 0.253 [HTML]BADD00 0.235 [HTML]C6E300 0.227 [HTML]AAD500 0.246 [HTML]96CB00 0.258 [HTML]90C800 0.263 [HTML]B6DB00 0.237 [HTML]94CA00 0.260 [HTML]AAD500 0.246 [HTML]B2D900 0.240 [HTML]9ECF00 0.253 6k-8k [HTML]6CB600 0.286 [HTML]CAE500 0.225 [HTML]CEE700 0.222 [HTML]C2E100 0.229 [HTML]D6EB00 0.216 [HTML]C8E400 0.226 [HTML]C0E000 0.231 [HTML]D6EB00 0.216 [HTML]C0E000 0.231 [HTML]DCEE00 0.213 [HTML]D6EB00 0.217 [HTML]C2E100 0.230 [HTML]C8E400 0.226 [HTML]D8EC00 0.215 [HTML]C8E400 0.226 [HTML]C8E400 0.226 [HTML]CAE500 0.224 [HTML]CCE600 0.223 8k-10k [HTML]BCDE00 0.233 [HTML]D0E800 0.221 [HTML]CCE600 0.223 [HTML]CAE500 0.224 [HTML]D0E800 0.221 [HTML]E8F400 0.204 [HTML]E2F100 0.208 [HTML]D6EB00 0.217 [HTML]E0F000 0.210 [HTML]FAFD00 0.193 [HTML]E0F000 0.210 [HTML]D8EC00 0.215 [HTML]D6EB00 0.216 [HTML]E2F100 0.208 [HTML]D0E800 0.220 [HTML]E4F200 0.207 [HTML]ECF600 0.202 [HTML]D6EB00 0.216 10k-12k [HTML]F0000 0.021 [HTML]DEEF00 0.211 [HTML]DEEF00 0.211 [HTML]E0F000 0.210 [HTML]D2E900 0.219 [HTML]E2F100 0.208 [HTML]F2F900 0.198 [HTML]E6F300 0.206 [HTML]DAED00 0.214 [HTML]FFFC00 0.188 [HTML]EEF700 0.201 [HTML]E0F000 0.210 [HTML]D6EB00 0.216 [HTML]E2F100 0.208 [HTML]DEEF00 0.211 [HTML]DCEE00 0.213 [HTML]DCEE00 0.212 [HTML]DEEF00 0.211 12k-14k [HTML]F1400 0.035 [HTML]D8EC00 0.215 [HTML]CCE600 0.223 [HTML]E2F100 0.208 [HTML]D8EC00 0.215 [HTML]E6F300 0.206 [HTML]F4FA00 0.197 [HTML]E2F100 0.208 [HTML]DAED00 0.214 [HTML]F000 0.180 [HTML]DAED00 0.214 [HTML]D8EC00 0.215 [HTML]DCEE00 0.212 [HTML]DEEF00 0.211 [HTML]D2E900 0.219 [HTML]ECF600 0.202 [HTML]E4F200 0.207 [HTML]DAED00 0.214 14k-16k [HTML]F2A00 0.049 [HTML]E0F000 0.210 [HTML]DCEE00 0.212 [HTML]DCEE00 0.213 [HTML]DEEF00 0.211 [HTML]DCEE00 0.213 [HTML]ECF600 0.202 [HTML]E2F100 0.209 [HTML]DCEE00 0.213 [HTML]EAF500 0.203 [HTML]DEEF00 0.211 [HTML]DAED00 0.214 [HTML]E6F300 0.206 [HTML]DEEF00 0.211 [HTML]DCEE00 0.212 [HTML]DCEE00 0.213 [HTML]F0F800 0.200 [HTML]E0F000 0.210 16k-24k [HTML]F1800 0.037 [HTML]D2E900 0.219 [HTML]CCE600 0.223 [HTML]CEE700 0.222 [HTML]D0E800 0.221 [HTML]E2F100 0.209 [HTML]E2F100 0.209 [HTML]DCEE00 0.213 [HTML]DEEF00 0.211 [HTML]FCFE00 0.191 [HTML]E2F100 0.208 [HTML]D6EB00 0.216 [HTML]D6EB00 0.216 [HTML]D6EB00 0.216 [HTML]CEE700 0.222 [HTML]EAF500 0.203 [HTML]FCFE00 0.191 [HTML]D0E800 0.220 24k-32k [HTML]F3200 0.054 [HTML]F800 0.185 [HTML]EEF700 0.201 [HTML]FFFA00 0.186 [HTML]F6FB00 0.196 [HTML]EEF700 0.201 [HTML]F6FB00 0.196 [HTML]FCFE00 0.192 [HTML]F800 0.185 [HTML]FFD800 0.164 [HTML]F4FA00 0.197 [HTML]F6FB00 0.195 [HTML]F2F900 0.198 [HTML]F400 0.182 [HTML]FFFE00 0.189 [HTML]F800 0.185 [HTML]F800 0.185 [HTML]FCFE00 0.192 32k-49k [HTML]F4400 0.067 [HTML]FFE200 0.171 [HTML]FFFE00 0.189 [HTML]FFEA00 0.176 [HTML]FFDE00 0.168 [HTML]FFE400 0.172 [HTML]FFA600 0.131 [HTML]FFCA00 0.155 [HTML]F200 0.181 [HTML]FFC400 0.151 [HTML]FFE600 0.173 [HTML]FFD600 0.163 [HTML]FFEA00 0.176 [HTML]FFE200 0.171 [HTML]FFD400 0.161 [HTML]FFD000 0.159 [HTML]FFCA00 0.155 [HTML]FFD800 0.164 49k-65k [HTML]F6400 0.088 [HTML]F7000 0.096 [HTML]F8000 0.106 [HTML]F5600 0.078 [HTML]F8400 0.109 [HTML]F5800 0.080 [HTML]F3200 0.054 [HTML]F6A00 0.092 [HTML]F6400 0.087 [HTML]F8400 0.109 [HTML]FFBA00 0.144 [HTML]F8A00 0.113 [HTML]F6800 0.090 [HTML]F6E00 0.094 [HTML]F6C00 0.093 [HTML]F6200 0.086 [HTML]F8C00 0.114 [HTML]F7C00 0.103 In-Domain (0k-8k) [HTML]2E9700 0.327 [HTML]7CBE00 0.275 [HTML]82C100 0.271 [HTML]7EBF00 0.274 [HTML]82C100 0.272 [HTML]86C300 0.269 [HTML]84C200 0.270 [HTML]84C200 0.270 [HTML]8C600 0.266 [HTML]AAD500 0.245 [HTML]8C600 0.265 [HTML]7CBE00 0.276 [HTML]78BC00 0.278 [HTML]90C800 0.262 [HTML]7CBE00 0.275 [HTML]8AC500 0.267 [HTML]94CA00 0.260 [HTML]80C000 0.273 OOD (8k-65k) [HTML]F4E00 0.074 [HTML]DEEF00 0.211 [HTML]D8EC00 0.215 [HTML]DEEF00 0.211 [HTML]DCEE00 0.213 [HTML]E8F400 0.205 [HTML]F2F900 0.198 [HTML]E6F300 0.206 [HTML]E2F100 0.208 [HTML]FFFA00 0.187 [HTML]E6F300 0.206 [HTML]E0F000 0.210 [HTML]E0F000 0.210 [HTML]E6F300 0.206 [HTML]DCEE00 0.212 [HTML]EAF500 0.203 [HTML]F2F900 0.199 [HTML]E0F000 0.210 Long-OOD (32k-65k) [HTML]F4C00 0.072 [HTML]FFC800 0.154 [HTML]FFE200 0.170 [HTML]FFC800 0.153 [HTML]FFCA00 0.154 [HTML]FFC400 0.151 [HTML]F8C00 0.113 [HTML]FFB400 0.140 [HTML]FFD200 0.159 [HTML]FFB600 0.141 [HTML]FFDC00 0.166 [HTML]FFC600 0.151 [HTML]FFCC00 0.156 [HTML]FFC800 0.153 [HTML]FFBC00 0.145 [HTML]FFB800 0.142 [HTML]FFBC00 0.146 [HTML]FFC200 0.150 Full (0k-65k) [HTML]FAFD00 0.192 [HTML]B0D800 0.241 [HTML]B0D800 0.242 [HTML]B2D900 0.240 [HTML]B2D900 0.241 [HTML]BADD00 0.235 [HTML]BEDF00 0.232 [HTML]B8DC00 0.236 [HTML]BADD00 0.235 [HTML]DAED00 0.214 [HTML]BCDE00 0.234 [HTML]B0D800 0.241 [HTML]B0D800 0.242 [HTML]BEDF00 0.232 [HTML]B0D800 0.242 [HTML]BEDF00 0.233 [HTML]C6E300 0.227 [HTML]B2D900 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 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 0k-1k [HTML]188C00 0.795 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 1k-2k [HTML]1C8E00 0.790 [HTML]50A800 0.705 [HTML]56AB00 0.695 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]56AB00 0.695 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]56AB00 0.695 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 2k-4k [HTML]3C9E00 0.736 [HTML]54A00 0.698 [HTML]5AAD00 0.689 [HTML]5AAD00 0.689 [HTML]54A00 0.698 [HTML]60B000 0.679 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]4EA700 0.708 [HTML]54A00 0.698 [HTML]5AAD00 0.689 [HTML]5AAD00 0.689 [HTML]5AAD00 0.689 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 4k-6k [HTML]5CAE00 0.685 [HTML]7EBF00 0.630 [HTML]76B00 0.644 [HTML]6CB600 0.658 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]7EBF00 0.630 [HTML]76B00 0.644 [HTML]6CB600 0.658 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]76B00 0.644 [HTML]7EBF00 0.630 [HTML]6CB600 0.658 [HTML]6CB600 0.658 6k-8k [HTML]008000 0.837 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]269300 0.772 [HTML]349A00 0.750 [HTML]269300 0.772 [HTML]269300 0.772 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]269300 0.772 [HTML]269300 0.772 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]349A00 0.750 [HTML]2E9700 0.761 8k-10k [HTML]A4D200 0.568 [HTML]0C8600 0.815 [HTML]0C8600 0.815 [HTML]0C8600 0.815 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]1C8E00 0.790 [HTML]0C8600 0.815 [HTML]1C8E00 0.790 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]0C8600 0.815 [HTML]1C8E00 0.790 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]2C9600 0.765 [HTML]148A00 0.802 10k-12k [HTML]F0000 0.000 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]369B00 0.747 [HTML]369B00 0.747 [HTML]309800 0.758 [HTML]369B00 0.747 [HTML]369B00 0.747 [HTML]369B00 0.747 [HTML]309800 0.758 [HTML]3C9E00 0.736 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]369B00 0.747 [HTML]369B00 0.747 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]369B00 0.747 12k-14k [HTML]F0000 0.000 [HTML]329900 0.755 [HTML]2A9500 0.766 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]2A9500 0.766 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]389C00 0.745 [HTML]329900 0.755 [HTML]389C00 0.745 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]4A500 0.713 [HTML]329900 0.755 14k-16k [HTML]F0000 0.000 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]78BC00 0.640 [HTML]7EBF00 0.629 [HTML]84C200 0.618 [HTML]7EBF00 0.629 [HTML]6AB500 0.663 [HTML]7EBF00 0.629 [HTML]78BC00 0.640 [HTML]7EBF00 0.629 [HTML]6AB500 0.663 [HTML]70B800 0.652 [HTML]78BC00 0.640 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]70B800 0.652 [HTML]7EBF00 0.629 16k-24k [HTML]F0000 0.000 [HTML]4A500 0.715 [HTML]52A900 0.701 [HTML]42A100 0.729 [HTML]4EA700 0.708 [HTML]46A300 0.722 [HTML]46A300 0.722 [HTML]46A300 0.722 [HTML]46A300 0.722 [HTML]46A300 0.722 [HTML]52A900 0.701 [HTML]46A300 0.722 [HTML]4A500 0.715 [HTML]4A500 0.715 [HTML]46A300 0.722 [HTML]4A500 0.715 [HTML]4A500 0.715 [HTML]46A300 0.722 24k-32k [HTML]F0600 0.010 [HTML]3E9F00 0.733 [HTML]50A800 0.703 [HTML]3E9F00 0.733 [HTML]44A200 0.723 [HTML]3E9F00 0.733 [HTML]329900 0.752 [HTML]3E9F00 0.733 [HTML]389C00 0.743 [HTML]389C00 0.743 [HTML]44A200 0.723 [HTML]3E9F00 0.733 [HTML]389C00 0.743 [HTML]3E9F00 0.733 [HTML]389C00 0.743 [HTML]44A200 0.723 [HTML]3E9F00 0.733 [HTML]389C00 0.743 32k-49k [HTML]F0000 0.000 [HTML]40A000 0.730 [HTML]46A300 0.722 [HTML]3A9D00 0.739 [HTML]46A300 0.722 [HTML]4A500 0.713 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]46A300 0.722 [HTML]3A9D00 0.739 [HTML]4A500 0.713 [HTML]40A000 0.730 [HTML]3A9D00 0.739 [HTML]46A300 0.722 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]3A9D00 0.739 49k-65k [HTML]F0000 0.000 [HTML]F1200 0.030 [HTML]7CBE00 0.634 [HTML]F1800 0.040 [HTML]F1E00 0.050 [HTML]F2400 0.059 [HTML]F1800 0.040 [HTML]F1800 0.040 [HTML]F1200 0.030 [HTML]F1800 0.040 [HTML]F2400 0.059 [HTML]F1800 0.040 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1800 0.040 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 In-Domain (0k-8k) [HTML]289400 0.771 [HTML]4A500 0.714 [HTML]4CA600 0.711 [HTML]48A400 0.716 [HTML]4A500 0.716 [HTML]4A500 0.714 [HTML]4A500 0.714 [HTML]4A500 0.716 [HTML]48A400 0.718 [HTML]48A400 0.718 [HTML]4A500 0.716 [HTML]4A500 0.714 [HTML]48A400 0.716 [HTML]4A500 0.714 [HTML]4A500 0.716 [HTML]4A500 0.714 [HTML]4A500 0.716 [HTML]48A400 0.718 OOD (8k-65k) [HTML]F2200 0.058 [HTML]76B00 0.643 [HTML]4CA600 0.711 [HTML]72B900 0.649 [HTML]76B00 0.641 [HTML]74BA00 0.646 [HTML]74BA00 0.647 [HTML]74BA00 0.647 [HTML]74BA00 0.646 [HTML]74BA00 0.646 [HTML]76B00 0.642 [HTML]74BA00 0.645 [HTML]74BA00 0.646 [HTML]76B00 0.642 [HTML]74BA00 0.645 [HTML]78BC00 0.641 [HTML]7ABD00 0.635 [HTML]74BA00 0.646 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]F600 0.403 [HTML]5EAF00 0.681 [HTML]FFFC00 0.412 [HTML]F800 0.408 [HTML]F800 0.407 [HTML]F800 0.407 [HTML]F800 0.407 [HTML]F600 0.403 [HTML]F600 0.403 [HTML]FEFF00 0.421 [HTML]F200 0.398 [HTML]F600 0.403 [HTML]F800 0.407 [HTML]F600 0.403 [HTML]F600 0.403 [HTML]F600 0.403 [HTML]F800 0.407 Full (0k-65k) [HTML]FFBE00 0.313 [HTML]66B300 0.668 [HTML]4CA600 0.711 [HTML]64B200 0.673 [HTML]66B300 0.668 [HTML]66B300 0.670 [HTML]64B200 0.671 [HTML]64B200 0.672 [HTML]64B200 0.672 [HTML]64B200 0.672 [HTML]66B300 0.668 [HTML]66B300 0.669 [HTML]64B200 0.671 [HTML]66B300 0.668 [HTML]66B300 0.670 [HTML]68B400 0.667 [HTML]68B400 0.664 [HTML]64B200 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 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 0k-1k [HTML]188C00 0.795 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]309800 0.756 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]289400 0.769 [HTML]309800 0.756 [HTML]289400 0.769 [HTML]309800 0.756 [HTML]289400 0.769 1k-2k [HTML]1C8E00 0.790 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]50A800 0.705 [HTML]ECF600 0.448 [HTML]66B300 0.667 [HTML]4A500 0.714 [HTML]50A800 0.705 [HTML]66B300 0.667 [HTML]5CAE00 0.686 [HTML]6EB700 0.657 [HTML]5CAE00 0.686 [HTML]56AB00 0.695 2k-4k [HTML]3C9E00 0.736 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]5AAD00 0.689 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]5AAD00 0.689 [HTML]FFFC00 0.415 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]54A00 0.698 [HTML]66B300 0.670 [HTML]4EA700 0.708 [HTML]54A00 0.698 [HTML]54A00 0.698 4k-6k [HTML]5CAE00 0.685 [HTML]7EBF00 0.630 [HTML]76B00 0.644 [HTML]6CB600 0.658 [HTML]6CB600 0.658 [HTML]76B00 0.644 [HTML]6CB600 0.658 [HTML]76B00 0.644 [HTML]7EBF00 0.630 [HTML]FFE200 0.370 [HTML]86C300 0.616 [HTML]76B00 0.644 [HTML]7EBF00 0.630 [HTML]6CB600 0.658 [HTML]6CB600 0.658 [HTML]6CB600 0.658 [HTML]7EBF00 0.630 [HTML]7EBF00 0.630 6k-8k [HTML]008000 0.837 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]48A400 0.717 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]148A00 0.804 [HTML]2E9700 0.761 [HTML]FFD400 0.348 [HTML]56AB00 0.696 [HTML]2E9700 0.761 [HTML]2E9700 0.761 [HTML]349A00 0.750 [HTML]0C8600 0.815 [HTML]209000 0.783 [HTML]2E9700 0.761 [HTML]2E9700 0.761 8k-10k [HTML]A4D200 0.568 [HTML]0C8600 0.815 [HTML]0C8600 0.815 [HTML]3A9D00 0.741 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]0C8600 0.815 [HTML]148A00 0.802 [HTML]148A00 0.802 [HTML]FFB400 0.296 [HTML]329900 0.753 [HTML]0C8600 0.815 [HTML]0C8600 0.815 [HTML]148A00 0.802 [HTML]0C8600 0.815 [HTML]148A00 0.802 [HTML]0C8600 0.815 [HTML]148A00 0.802 10k-12k [HTML]F0000 0.000 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]50A800 0.703 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]369B00 0.747 [HTML]3C9E00 0.736 [HTML]369B00 0.747 [HTML]F800 0.407 [HTML]44A200 0.725 [HTML]44A200 0.725 [HTML]44A200 0.725 [HTML]289400 0.769 [HTML]44A200 0.725 [HTML]309800 0.758 [HTML]3C9E00 0.736 [HTML]309800 0.758 12k-14k [HTML]F0000 0.000 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]FFB600 0.298 [HTML]389C00 0.745 [HTML]389C00 0.745 [HTML]2A9500 0.766 [HTML]329900 0.755 [HTML]3E9F00 0.734 [HTML]2A9500 0.766 [HTML]2A9500 0.766 [HTML]389C00 0.745 14k-16k [HTML]F0000 0.000 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]70B800 0.652 [HTML]7EBF00 0.629 [HTML]84C200 0.618 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]FFC600 0.326 [HTML]78BC00 0.640 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]84C200 0.618 [HTML]6AB500 0.663 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 [HTML]7EBF00 0.629 16k-24k [HTML]F0000 0.000 [HTML]4A500 0.715 [HTML]4A500 0.715 [HTML]5EAF00 0.681 [HTML]46A300 0.722 [HTML]4EA700 0.708 [HTML]4EA700 0.708 [HTML]4EA700 0.708 [HTML]4A500 0.715 [HTML]F9000 0.236 [HTML]5EAF00 0.681 [HTML]4A500 0.715 [HTML]46A300 0.722 [HTML]46A300 0.722 [HTML]4A500 0.715 [HTML]4EA700 0.708 [HTML]42A100 0.729 [HTML]4A500 0.715 24k-32k [HTML]F0600 0.010 [HTML]3E9F00 0.733 [HTML]3E9F00 0.733 [HTML]58AC00 0.693 [HTML]3E9F00 0.733 [HTML]389C00 0.743 [HTML]389C00 0.743 [HTML]3E9F00 0.733 [HTML]44A200 0.723 [HTML]FFA800 0.277 [HTML]58AC00 0.693 [HTML]44A200 0.723 [HTML]44A200 0.723 [HTML]3E9F00 0.733 [HTML]4A500 0.713 [HTML]3E9F00 0.733 [HTML]3E9F00 0.733 [HTML]3E9F00 0.733 32k-49k [HTML]F0000 0.000 [HTML]40A000 0.730 [HTML]46A300 0.722 [HTML]66B300 0.670 [HTML]46A300 0.722 [HTML]4A500 0.713 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]F9400 0.243 [HTML]50A800 0.704 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]46A300 0.722 [HTML]40A000 0.730 [HTML]40A000 0.730 [HTML]3A9D00 0.739 49k-65k [HTML]F0000 0.000 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F2400 0.059 [HTML]F1200 0.030 [HTML]F1E00 0.050 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1200 0.030 [HTML]F1800 0.040 [HTML]F1200 0.030 [HTML]F1200 0.030 In-Domain (0k-8k) [HTML]289400 0.771 [HTML]4A500 0.714 [HTML]4A500 0.716 [HTML]4EA700 0.707 [HTML]48A400 0.718 [HTML]4A500 0.716 [HTML]48A400 0.718 [HTML]44A200 0.725 [HTML]4CA600 0.712 [HTML]E6F300 0.460 [HTML]5AAD00 0.689 [HTML]48A400 0.718 [HTML]4A500 0.714 [HTML]4EA700 0.707 [HTML]4A500 0.716 [HTML]4A500 0.714 [HTML]4EA700 0.707 [HTML]4CA600 0.711 OOD (8k-65k) [HTML]F2200 0.058 [HTML]76B00 0.643 [HTML]76B00 0.641 [HTML]88C400 0.614 [HTML]76B00 0.642 [HTML]7ABD00 0.637 [HTML]74BA00 0.647 [HTML]78BC00 0.640 [HTML]76B00 0.643 [HTML]F9E00 0.259 [HTML]84C200 0.620 [HTML]78BC00 0.638 [HTML]76B00 0.642 [HTML]74BA00 0.644 [HTML]78BC00 0.639 [HTML]74BA00 0.644 [HTML]74BA00 0.646 [HTML]76B00 0.643 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]F600 0.403 [HTML]F200 0.398 [HTML]FFE200 0.371 [HTML]F200 0.398 [HTML]F000 0.394 [HTML]FFFE00 0.416 [HTML]F600 0.403 [HTML]FFFC00 0.412 [HTML]F5600 0.143 [HTML]FFEC00 0.389 [HTML]F600 0.403 [HTML]F600 0.403 [HTML]F600 0.403 [HTML]F200 0.398 [HTML]F800 0.407 [HTML]F600 0.403 [HTML]F800 0.407 Full (0k-65k) [HTML]FFBE00 0.313 [HTML]66B300 0.668 [HTML]66B300 0.668 [HTML]72B900 0.647 [HTML]66B300 0.669 [HTML]68B400 0.665 [HTML]64B200 0.672 [HTML]66B300 0.670 [HTML]66B300 0.668 [HTML]FFCA00 0.331 [HTML]74BA00 0.645 [HTML]68B400 0.667 [HTML]66B300 0.668 [HTML]68B400 0.667 [HTML]68B400 0.666 [HTML]66B300 0.669 [HTML]66B300 0.668 [HTML]66B300 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. 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 pre-selected 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. 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 [HTML]188C00 0.795 [HTML]48A400 0.718 [HTML]309800 0.756 [HTML]289400 0.769 [HTML]40A000 0.731 [HTML]389C00 0.744 [HTML]389C00 0.744 1k-2k [HTML]1C8E00 0.790 [HTML]78BC00 0.638 [HTML]5CAE00 0.686 [HTML]50A800 0.705 [HTML]72B900 0.648 [HTML]66B300 0.667 [HTML]66B300 0.667 2k-4k [HTML]3C9E00 0.736 [HTML]76B00 0.642 [HTML]48A400 0.717 [HTML]54A00 0.698 [HTML]76B00 0.642 [HTML]54A00 0.698 [HTML]54A00 0.698 4k-6k [HTML]5CAE00 0.685 [HTML]8EC700 0.603 [HTML]7EBF00 0.630 [HTML]7EBF00 0.630 [HTML]A8D400 0.562 [HTML]6CB600 0.658 [HTML]76B00 0.644 6k-8k [HTML]008000 0.837 [HTML]92C900 0.598 [HTML]3A9D00 0.739 [HTML]2E9700 0.761 [HTML]B2D900 0.543 [HTML]42A100 0.728 [HTML]42A100 0.728 8k-10k [HTML]A4D200 0.568 [HTML]E0F000 0.469 [HTML]148A00 0.802 [HTML]0C8600 0.815 [HTML]E0F000 0.469 [HTML]50A800 0.704 [HTML]2C9600 0.765 10k-12k [HTML]F0000 0.000 [HTML]FFFE00 0.418 [HTML]58AC00 0.692 [HTML]369B00 0.747 [HTML]F8FC00 0.429 [HTML]58AC00 0.692 [HTML]309800 0.758 12k-14k [HTML]F0000 0.000 [HTML]F5400 0.138 [HTML]329900 0.755 [HTML]329900 0.755 [HTML]FFE200 0.372 [HTML]52A900 0.702 [HTML]3E9F00 0.734 14k-16k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]6AB500 0.663 [HTML]7EBF00 0.629 [HTML]F7A00 0.202 [HTML]7EBF00 0.629 [HTML]78BC00 0.640 16k-24k [HTML]F0000 0.000 [HTML]F0800 0.014 [HTML]4EA700 0.708 [HTML]4A500 0.715 [HTML]F5800 0.146 [HTML]80C000 0.625 [HTML]70B800 0.653 24k-32k [HTML]F0600 0.010 [HTML]F0000 0.000 [HTML]F9C00 0.257 [HTML]3E9F00 0.733 [HTML]F4E00 0.129 [HTML]82C100 0.624 [HTML]44A200 0.723 32k-49k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]40A000 0.730 [HTML]F2E00 0.078 [HTML]86C300 0.617 [HTML]56AB00 0.696 49k-65k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F1200 0.030 [HTML]F4800 0.119 [HTML]D6EB00 0.485 [HTML]B8DC00 0.535 In-Domain (0k-8k) [HTML]289400 0.771 [HTML]78BC00 0.639 [HTML]4EA700 0.707 [HTML]4A500 0.714 [HTML]80C000 0.626 [HTML]54A00 0.698 [HTML]56AB00 0.696 OOD (8k-65k) [HTML]F2200 0.058 [HTML]F4400 0.112 [HTML]DEEF00 0.473 [HTML]76B00 0.643 [HTML]F8A00 0.227 [HTML]7CBE00 0.631 [HTML]5CAE00 0.684 Long-OOD (32k-65k) [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F600 0.403 [HTML]F3A00 0.097 [HTML]ACD600 0.555 [HTML]84C200 0.621 Full (0k-65k) [HTML]FFBE00 0.313 [HTML]FFB600 0.300 [HTML]AAD500 0.557 [HTML]66B300 0.668 [HTML]FFE000 0.369 [HTML]6EB700 0.655 [HTML]5AAD00 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 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 0k-1k [HTML]0C8600 0.349 [HTML]209000 0.335 [HTML]269300 0.332 [HTML]1E8F00 0.337 [HTML]128900 0.345 [HTML]108800 0.346 [HTML]2A9500 0.329 1k-2k [HTML]008000 0.358 [HTML]64B200 0.289 [HTML]5AAD00 0.297 [HTML]52A900 0.302 [HTML]8EC700 0.261 [HTML]94CA00 0.257 [HTML]9CCE00 0.252 2k-4k [HTML]289400 0.330 [HTML]B0D800 0.238 [HTML]94CA00 0.257 [HTML]8EC700 0.261 [HTML]BCDE00 0.230 [HTML]AAD500 0.242 [HTML]A2D100 0.248 4k-6k [HTML]40A000 0.314 [HTML]CAE500 0.220 [HTML]A6D300 0.245 [HTML]9ACD00 0.253 [HTML]E2F100 0.204 [HTML]C2E100 0.226 [HTML]B4DA00 0.236 6k-8k [HTML]6AB500 0.286 [HTML]D6EB00 0.212 [HTML]D2E900 0.215 [HTML]C4E200 0.225 [HTML]F800 0.180 [HTML]CCE600 0.219 [HTML]CAE500 0.220 8k-10k [HTML]B8DC00 0.233 [HTML]F400 0.177 [HTML]E6F300 0.202 [HTML]CAE500 0.221 [HTML]FFD600 0.157 [HTML]FCFE00 0.186 [HTML]EAF500 0.199 10k-12k [HTML]F0E00 0.021 [HTML]FFD400 0.156 [HTML]EAF500 0.199 [HTML]D8EC00 0.211 [HTML]FFDA00 0.159 [HTML]E6F300 0.202 [HTML]EAF500 0.199 12k-14k [HTML]F2200 0.035 [HTML]FFB800 0.137 [HTML]F0F800 0.195 [HTML]D2E900 0.215 [HTML]FFB400 0.134 [HTML]FCFE00 0.186 [HTML]EAF500 0.199 14k-16k [HTML]F3800 0.049 [HTML]F5C00 0.074 [HTML]F6FB00 0.190 [HTML]DAED00 0.210 [HTML]FFD600 0.157 [HTML]E6F300 0.202 [HTML]D4EA00 0.213 16k-24k [HTML]F2600 0.037 [HTML]F3200 0.046 [HTML]FFBE00 0.140 [HTML]CCE600 0.219 [HTML]FFD400 0.156 [HTML]E4F200 0.203 [HTML]DCEE00 0.208 24k-32k [HTML]F3E00 0.054 [HTML]F0000 0.011 [HTML]F4600 0.059 [HTML]FEFF00 0.185 [HTML]FFE200 0.165 [HTML]FFE800 0.169 [HTML]F600 0.178 32k-49k [HTML]F5200 0.067 [HTML]F1400 0.025 [HTML]F1E00 0.032 [HTML]FFEC00 0.171 [HTML]FFEA00 0.170 [HTML]FFE000 0.164 [HTML]FFE400 0.166 49k-65k [HTML]F7000 0.088 [HTML]F0400 0.015 [HTML]F0000 0.011 [HTML]F7C00 0.096 [HTML]FFE200 0.165 [HTML]D4EA00 0.214 [HTML]F200 0.176 In-Domain (0k-8k) [HTML]2C9600 0.327 [HTML]92C900 0.258 [HTML]82C100 0.269 [HTML]7ABD00 0.275 [HTML]A8D400 0.244 [HTML]94CA00 0.258 [HTML]94CA00 0.257 OOD (8k-65k) [HTML]F5C00 0.074 [HTML]F8C00 0.106 [HTML]FFEA00 0.170 [HTML]D8EC00 0.211 [HTML]FFD200 0.154 [HTML]F2F900 0.193 [HTML]E8F400 0.200 Long-OOD (32k-65k) [HTML]F5800 0.072 [HTML]F1000 0.023 [HTML]F1600 0.027 [HTML]FFD200 0.154 [HTML]FFE800 0.169 [HTML]F200 0.176 [HTML]FFE800 0.168 Full (0k-65k) [HTML]F4FA00 0.192 [HTML]F400 0.178 [HTML]D0E800 0.216 [HTML]ACD600 0.241 [HTML]EEF700 0.196 [HTML]C6E300 0.223 [HTML]C0E000 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. Appendix E Training Hyperparameters 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. Dataset Model Length Learning Rate Training Steps Total Batch Size MT Gemma-3-1B-IT 8k 1e-4 10.000 64 MT ARMT 2k 1e-4 10.000 64 MT ARMT 4k 1e-4 10.000 64 MT ARMT 8k 3e-5 10.000 64 GR Gemma-3-1B-IT 8k 1e-4 8.000 64 GR ARMT 2k 1e-4 4.000 64 GR ARMT 4k 3e-5 4.000 64 GR ARMT 8k 1e-5 4.000 64 GR-100+ Gemma-3-1B-IT 8k 1e-4 10.000 64 GR-100+ ARMT 2k 1e-4 5.000 64 GR-100+ ARMT 4k 3e-5 5.000 64 GR-100+ ARMT 8k 1e-5 5.000 64 ContractNLI Gemma-3-1B-IT 8k 1e-4 500 64 ContractNLI ARMT 2k 1e-4 250 64 ContractNLI ARMT 4k 3e-5 250 64 ContractNLI ARMT 8k 1e-5 250 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. Appendix F 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. Model/ Lengths 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 0k-1k [HTML]068300 0.380 [HTML]249200 0.357 [HTML]249200 0.357 [HTML]229100 0.358 [HTML]3C9E00 0.339 [HTML]2E9700 0.349 [HTML]329900 0.346 1k-2k [HTML]008000 0.385 [HTML]76B00 0.296 [HTML]78BC00 0.294 [HTML]6CB600 0.303 [HTML]76B00 0.296 [HTML]72B900 0.299 [HTML]6AB500 0.305 2k-4k [HTML]2A9500 0.352 [HTML]A0D000 0.264 [HTML]92C900 0.275 [HTML]8C600 0.279 [HTML]B4DA00 0.249 [HTML]7EBF00 0.290 [HTML]8EC700 0.277 4k-6k [HTML]369B00 0.344 [HTML]C0E000 0.240 [HTML]82C100 0.286 [HTML]6CB600 0.303 [HTML]C6E300 0.235 [HTML]A2D100 0.263 [HTML]7EBF00 0.290 6k-8k [HTML]76B00 0.296 [HTML]DCEE00 0.219 [HTML]C2E100 0.238 [HTML]AED700 0.254 [HTML]D2E900 0.227 [HTML]C2E100 0.239 [HTML]B6DB00 0.248 8k-10k [HTML]68B400 0.306 [HTML]EEF700 0.205 [HTML]CAE500 0.233 [HTML]C6E300 0.235 [HTML]FFFC00 0.190 [HTML]C2E100 0.238 [HTML]BADD00 0.244 10k-12k [HTML]9ACD00 0.269 [HTML]E4F200 0.213 [HTML]DAED00 0.221 [HTML]BEDF00 0.241 [HTML]FFD200 0.159 [HTML]D2E900 0.227 [HTML]D4EA00 0.225 12k-14k [HTML]8AC500 0.280 [HTML]FEFF00 0.193 [HTML]D4EA00 0.225 [HTML]C4E200 0.237 [HTML]F8A00 0.105 [HTML]F4FA00 0.201 [HTML]CAE500 0.232 14k-16k [HTML]CCE600 0.231 [HTML]FFEE00 0.180 [HTML]DEEF00 0.217 [HTML]C0E000 0.240 [HTML]F3800 0.043 [HTML]FFD000 0.157 [HTML]C0E000 0.240 16k-24k [HTML]9ACD00 0.269 [HTML]F9800 0.115 [HTML]F6FB00 0.200 [HTML]B0D800 0.252 [HTML]F1400 0.016 [HTML]F3C00 0.046 [HTML]BEDF00 0.241 24k-32k [HTML]B6DB00 0.247 [HTML]F1400 0.016 [HTML]FFDE00 0.168 [HTML]ECF600 0.207 [HTML]F0000 0.000 [HTML]F0000 0.001 [HTML]F800 0.187 32k-49k [HTML]FFE200 0.171 [HTML]F0200 0.002 [HTML]FFAA00 0.128 [HTML]F6FB00 0.200 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F6FB00 0.199 49k-65k [HTML]FFFA00 0.189 [HTML]F0000 0.000 [HTML]F3400 0.040 [HTML]A0D000 0.264 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]FEFF00 0.194 In-Domain (0k-8k) [HTML]2C9600 0.351 [HTML]92C900 0.275 [HTML]7EBF00 0.290 [HTML]72B900 0.299 [HTML]9ACD00 0.269 [HTML]80C000 0.288 [HTML]7ABD00 0.293 OOD (8k-65k) [HTML]9ECF00 0.266 [HTML]FFD600 0.162 [HTML]E6F300 0.211 [HTML]C2E100 0.238 [HTML]F7800 0.090 [HTML]FFCC00 0.154 [HTML]CAE500 0.232 Long-OOD (32k-65k) [HTML]FFE800 0.175 [HTML]F0200 0.002 [HTML]F8E00 0.108 [HTML]E2F100 0.215 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F8FC00 0.198 Full (0k-65k) [HTML]68B400 0.306 [HTML]E2F100 0.215 [HTML]B6DB00 0.248 [HTML]9CCE00 0.266 [HTML]FFE600 0.174 [HTML]E0F000 0.217 [HTML]A4D200 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 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 0k-1k [HTML]2A9500 0.795 [HTML]329900 0.782 [HTML]329900 0.782 [HTML]42A100 0.756 [HTML]229100 0.808 [HTML]229100 0.808 [HTML]329900 0.782 1k-2k [HTML]2E9700 0.790 [HTML]54A00 0.724 [HTML]5AAD00 0.714 [HTML]4A500 0.743 [HTML]66B300 0.695 [HTML]72B900 0.676 [HTML]66B300 0.695 2k-4k [HTML]389C00 0.774 [HTML]6AB500 0.689 [HTML]64B200 0.698 [HTML]4EA700 0.736 [HTML]86C300 0.642 [HTML]76B00 0.670 [HTML]5EAF00 0.708 4k-6k [HTML]3C9E00 0.767 [HTML]96CB00 0.616 [HTML]84C200 0.644 [HTML]64B200 0.699 [HTML]9ECF00 0.603 [HTML]84C200 0.644 [HTML]64B200 0.699 6k-8k [HTML]048200 0.859 [HTML]9ACD00 0.609 [HTML]5AAD00 0.717 [HTML]269300 0.804 [HTML]FFBC00 0.326 [HTML]52A900 0.728 [HTML]269300 0.804 8k-10k [HTML]088400 0.852 [HTML]D8EC00 0.506 [HTML]44A200 0.753 [HTML]44A200 0.753 [HTML]F4200 0.123 [HTML]7EBF00 0.654 [HTML]108800 0.840 10k-12k [HTML]008000 0.868 [HTML]FFD800 0.374 [HTML]76B00 0.670 [HTML]2C9600 0.791 [HTML]F5C00 0.165 [HTML]54A00 0.725 [HTML]1A8D00 0.824 12k-14k [HTML]309800 0.787 [HTML]FFEA00 0.404 [HTML]62B100 0.702 [HTML]48A400 0.745 [HTML]F3200 0.096 [HTML]82C100 0.649 [HTML]56AB00 0.723 14k-16k [HTML]3E9F00 0.764 [HTML]FFC200 0.337 [HTML]88C400 0.640 [HTML]6CB600 0.685 [HTML]F1400 0.045 [HTML]94CA00 0.618 [HTML]6CB600 0.685 16k-24k [HTML]209000 0.812 [HTML]F8200 0.229 [HTML]84C200 0.646 [HTML]42A100 0.757 [HTML]F1E00 0.062 [HTML]AAD500 0.583 [HTML]56AB00 0.722 24k-32k [HTML]5CAE00 0.713 [HTML]F7000 0.198 [HTML]A8D400 0.584 [HTML]5CAE00 0.713 [HTML]F1000 0.040 [HTML]D2E900 0.515 [HTML]68B400 0.693 32k-49k [HTML]A4D200 0.591 [HTML]F4C00 0.139 [HTML]9ACD00 0.609 [HTML]3C9E00 0.765 [HTML]F1200 0.043 [HTML]E8F400 0.478 [HTML]46A300 0.748 49k-65k [HTML]FFB600 0.317 [HTML]F2E00 0.089 [HTML]9CCE00 0.604 [HTML]5CAE00 0.713 [HTML]F0000 0.010 [HTML]D8EC00 0.505 [HTML]50A800 0.733 In-Domain (0k-8k) [HTML]2A9500 0.797 [HTML]6CB600 0.685 [HTML]5CAE00 0.711 [HTML]46A300 0.749 [HTML]98C00 0.612 [HTML]62B100 0.703 [HTML]4EA700 0.736 OOD (8k-65k) [HTML]5EAF00 0.709 [HTML]F9A00 0.271 [HTML]82C100 0.647 [HTML]4A500 0.741 [HTML]F2200 0.070 [HTML]A8D400 0.584 [HTML]4A500 0.743 Long-OOD (32k-65k) [HTML]F0F800 0.463 [HTML]F3E00 0.116 [HTML]9ACD00 0.607 [HTML]4A500 0.741 [HTML]F0A00 0.028 [HTML]E0F000 0.491 [HTML]4A500 0.741 Full (0k-65k) [HTML]4A500 0.741 [HTML]F200 0.419 [HTML]76B00 0.670 [HTML]48A400 0.744 [HTML]F9600 0.264 [HTML]8EC700 0.627 [HTML]4CA600 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. Appendix G 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). Model/ Lengths 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 0k-1k [HTML]068300 0.380 [HTML]249200 0.357 [HTML]249200 0.357 [HTML]229100 0.358 [HTML]DEEF00 0.218 [HTML]BEDF00 0.242 [HTML]FAFD00 0.196 1k-2k [HTML]008000 0.385 [HTML]76B00 0.296 [HTML]78BC00 0.294 [HTML]6CB600 0.303 [HTML]F8FC00 0.198 [HTML]D6EB00 0.224 [HTML]A2D100 0.263 2k-4k [HTML]2A9500 0.352 [HTML]A0D000 0.264 [HTML]92C900 0.275 [HTML]8C600 0.279 [HTML]FFAE00 0.132 [HTML]FFDE00 0.168 [HTML]DAED00 0.220 4k-6k [HTML]369B00 0.344 [HTML]C0E000 0.240 [HTML]82C100 0.286 [HTML]6CB600 0.303 [HTML]FFB600 0.138 [HTML]FFE200 0.170 [HTML]E2F100 0.214 6k-8k [HTML]76B00 0.296 [HTML]DCEE00 0.219 [HTML]C2E100 0.238 [HTML]AED700 0.254 [HTML]FFA200 0.122 [HTML]FFDA00 0.165 [HTML]D2E900 0.226 8k-10k [HTML]68B400 0.306 [HTML]EEF700 0.205 [HTML]CAE500 0.233 [HTML]C6E300 0.235 [HTML]F8A00 0.104 [HTML]FFCA00 0.152 [HTML]FFDE00 0.167 10k-12k [HTML]9ACD00 0.269 [HTML]E4F200 0.213 [HTML]DAED00 0.221 [HTML]BEDF00 0.241 [HTML]F9E00 0.119 [HTML]FFD000 0.157 [HTML]FFFC00 0.190 12k-14k [HTML]8AC500 0.280 [HTML]FEFF00 0.193 [HTML]D4EA00 0.225 [HTML]C4E200 0.237 [HTML]F9E00 0.119 [HTML]FFA800 0.127 [HTML]FFE200 0.171 14k-16k [HTML]CCE600 0.231 [HTML]FFEE00 0.180 [HTML]DEEF00 0.217 [HTML]C0E000 0.240 [HTML]F9600 0.114 [HTML]F9A00 0.117 [HTML]FFCC00 0.154 16k-24k [HTML]9ACD00 0.269 [HTML]F9800 0.115 [HTML]F6FB00 0.200 [HTML]B0D800 0.252 [HTML]F8C00 0.106 [HTML]F8800 0.103 [HTML]FFBA00 0.141 24k-32k [HTML]B6DB00 0.247 [HTML]F1400 0.016 [HTML]FFDE00 0.168 [HTML]ECF600 0.207 [HTML]F8600 0.101 [HTML]F8200 0.098 [HTML]F9A00 0.117 32k-49k [HTML]FFE200 0.171 [HTML]F0200 0.002 [HTML]FFAA00 0.128 [HTML]F6FB00 0.200 [HTML]F8C00 0.106 [HTML]F9600 0.114 [HTML]F8000 0.097 49k-65k [HTML]FFFA00 0.189 [HTML]F0000 0.000 [HTML]F3400 0.040 [HTML]A0D000 0.264 [HTML]F7C00 0.094 [HTML]F1400 0.016 [HTML]FFBE00 0.144 In-Domain (0k-8k) [HTML]2C9600 0.351 [HTML]92C900 0.275 [HTML]7EBF00 0.290 [HTML]72B900 0.299 [HTML]FFD600 0.161 [HTML]FEFF00 0.193 [HTML]D6EB00 0.224 OOD (8k-65k) [HTML]9ECF00 0.266 [HTML]FFD600 0.162 [HTML]E6F300 0.211 [HTML]C2E100 0.238 [HTML]F9200 0.111 [HTML]FFA800 0.127 [HTML]FFD200 0.159 Long-OOD (32k-65k) [HTML]FFE800 0.175 [HTML]F0200 0.002 [HTML]F8E00 0.108 [HTML]E2F100 0.215 [HTML]F8800 0.103 [HTML]F7800 0.091 [HTML]F8E00 0.108 Full (0k-65k) [HTML]68B400 0.306 [HTML]E2F100 0.215 [HTML]B6DB00 0.248 [HTML]9CCE00 0.266 [HTML]FFB200 0.135 [HTML]FFD000 0.158 [HTML]FFFA00 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. Appendix H ContractNLI Results 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. Model/ Lengths Base, No Fine- Tuning Base, CNLI, 8k ARMT, CNLI, 2k ARMT, CNLI, 4k ARMT, CNLI, 8k ARMT, CNLI, 2k, pretrain ARMT, CNLI, 4k, pretrain ARMT, CNLI, 8k, pretrain 0k-1k [HTML]FFE800 0.427 [HTML]7ABD00 0.703 [HTML]ACD600 0.615 [HTML]88C400 0.678 [HTML]76B00 0.710 [HTML]82C100 0.689 [HTML]8AC500 0.675 [HTML]7ABD00 0.703 1k-2k [HTML]F2A00 0.089 [HTML]82C100 0.690 [HTML]8C600 0.672 [HTML]80C000 0.691 [HTML]7CBE00 0.701 [HTML]8C600 0.671 [HTML]6EB700 0.724 [HTML]68B400 0.737 2k-4k [HTML]F1200 0.046 [HTML]7CBE00 0.700 [HTML]E6F300 0.512 [HTML]88C400 0.680 [HTML]6CB600 0.730 [HTML]D4EA00 0.542 [HTML]7ABD00 0.703 [HTML]6CB600 0.727 4k-6k [HTML]F0000 0.012 [HTML]40A000 0.806 [HTML]FFBC00 0.347 [HTML]62B100 0.747 [HTML]52A900 0.776 [HTML]FFC200 0.359 [HTML]62B100 0.747 [HTML]58AC00 0.765 6k-8k [HTML]F6600 0.196 [HTML]008000 0.922 [HTML]F8800 0.255 [HTML]58AC00 0.765 [HTML]168B00 0.882 [HTML]F8800 0.255 [HTML]4CA600 0.784 [HTML]369B00 0.824 Full (0k-8k) [HTML]F3C00 0.119 [HTML]76B00 0.710 [HTML]C0E000 0.579 [HTML]80C000 0.692 [HTML]70B800 0.721 [HTML]B4DA00 0.599 [HTML]74BA00 0.714 [HTML]68B400 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. Appendix I 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-3-1B-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 the long-context splits and generalizes to context lengths of up to 64k tokens. Task QA1 QA2 QA3 Model/ Length 0k 1k 2k 4k 8k 16k 32k 64k 0k 1k 2k 4k 8k 16k 32k 64k 0k 1k 2k 4k 8k 16k 32k 64k Gemma-3-1B-IT [HTML]56AB00 83 [HTML]CCE600 60 [HTML]CCE600 60 [HTML]FFCC00 40 [HTML]F2800 8 [HTML]F1400 4 [HTML]F0000 0 [HTML]F0000 0 [HTML]F400 48 [HTML]F000 47 [HTML]F7400 23 [HTML]F2800 8 [HTML]F1800 5 [HTML]F0400 1 [HTML]F0000 0 [HTML]F0000 0 [HTML]FFC200 38 [HTML]FFAE00 34 [HTML]F9800 30 [HTML]F6000 19 [HTML]F2800 8 [HTML]F4200 13 [HTML]F0000 0 [HTML]F0000 0 Gemma-3-1B-IT, BABILong, 5k [HTML]008000 100 [HTML]008000 100 [HTML]008000 100 [HTML]008000 100 [HTML]048200 99 [HTML]148A00 96 [HTML]1E8F00 94 [HTML]F0E00 3 [HTML]229100 93 [HTML]048200 99 [HTML]008000 100 [HTML]0E8700 97 [HTML]2E9700 91 [HTML]4CA600 85 [HTML]A8D400 67 [HTML]F0A00 2 [HTML]389C00 89 [HTML]188C00 95 [HTML]188C00 95 [HTML]229100 93 [HTML]2E9700 91 [HTML]389C00 89 [HTML]60B000 81 [HTML]F2E00 9 ARMT, BABILong, 2k [HTML]008000 100 [HTML]008000 100 [HTML]0E8700 97 [HTML]9ECF00 69 [HTML]FFDC00 43 [HTML]F8E00 28 [HTML]F8A00 27 [HTML]F1400 4 [HTML]008000 100 [HTML]008000 100 [HTML]4CA600 85 [HTML]FFC600 39 [HTML]F8A00 27 [HTML]F8E00 28 [HTML]F5600 17 [HTML]F1400 4 [HTML]0A8500 98 [HTML]048200 99 [HTML]148A00 96 [HTML]C2E100 62 [HTML]FFC200 38 [HTML]F6A00 21 [HTML]F6A00 21 [HTML]F1800 5 ARMT, BABILong, 3k [HTML]008000 100 [HTML]008000 100 [HTML]048200 99 [HTML]4CA600 85 [HTML]C6E300 61 [HTML]FFE000 44 [HTML]F5600 17 [HTML]F4200 13 [HTML]048200 99 [HTML]048200 99 [HTML]289400 92 [HTML]7ABD00 76 [HTML]FFC200 38 [HTML]FFAE00 34 [HTML]F5C00 18 [HTML]F3C00 12 [HTML]188C00 95 [HTML]0A8500 98 [HTML]188C00 95 [HTML]60B000 81 [HTML]F8A00 27 [HTML]F5C00 18 [HTML]F2200 7 [HTML]F2200 7 ARMT, BABILong, 5k [HTML]008000 100 [HTML]008000 100 [HTML]008000 100 [HTML]008000 100 [HTML]0A8500 98 [HTML]008000 100 [HTML]008000 100 [HTML]048200 99 [HTML]008000 100 [HTML]008000 100 [HTML]0A8500 98 [HTML]0A8500 98 [HTML]148A00 96 [HTML]42A100 87 [HTML]42A100 87 [HTML]50A800 84 [HTML]048200 99 [HTML]008000 100 [HTML]0A8500 98 [HTML]188C00 95 [HTML]2E9700 91 [HTML]42A100 87 [HTML]6AB500 79 [HTML]9ECF00 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 0k 1k 2k 4k 8k 16k 32k 64k 0k 1k 2k 4k 8k 16k 32k 64k Avg. (QA1-QA5) Gemma-3-1B-IT [HTML]F9800 30 [HTML]F3C00 12 [HTML]F3800 11 [HTML]F2200 7 [HTML]F1E00 6 [HTML]F4600 14 [HTML]F0400 1 [HTML]F0000 0 [HTML]98C00 70 [HTML]FFC600 39 [HTML]EAF500 54 [HTML]FFCC00 40 [HTML]F9E00 31 [HTML]F6A00 21 [HTML]F2200 7 [HTML]F0000 0 [HTML]F6E00 22 Gemma-3-1B-IT, BABILong, 5k [HTML]188C00 95 [HTML]008000 100 [HTML]008000 100 [HTML]048200 99 [HTML]0E8700 97 [HTML]60B000 81 [HTML]B2D900 65 [HTML]F1800 5 [HTML]0E8700 97 [HTML]008000 100 [HTML]048200 99 [HTML]0A8500 98 [HTML]0E8700 97 [HTML]188C00 95 [HTML]2E9700 91 [HTML]F8000 25 [HTML]58AC00 83 ARMT, BABILong, 2k [HTML]008000 100 [HTML]008000 100 [HTML]289400 92 [HTML]A8D400 67 [HTML]FFE000 44 [HTML]FFC600 39 [HTML]FFD600 42 [HTML]F1E00 6 [HTML]048200 99 [HTML]048200 99 [HTML]0A8500 98 [HTML]46A300 86 [HTML]E6F300 55 [HTML]FFE600 45 [HTML]FFB800 36 [HTML]F0E00 3 [HTML]D6EB00 58 ARMT, BABILong, 3k [HTML]008000 100 [HTML]008000 100 [HTML]229100 93 [HTML]74BA00 77 [HTML]AED700 66 [HTML]C2E100 62 [HTML]FFB800 36 [HTML]F3200 10 [HTML]048200 99 [HTML]048200 99 [HTML]048200 99 [HTML]289400 92 [HTML]94CA00 71 [HTML]FFB800 36 [HTML]F4600 14 [HTML]F0A00 2 [HTML]C2E100 62 ARMT, BABILong, 5k [HTML]008000 100 [HTML]008000 100 [HTML]329900 90 [HTML]56AB00 83 [HTML]5CAE00 82 [HTML]6AB500 79 [HTML]6AB500 79 [HTML]5CAE00 82 [HTML]048200 99 [HTML]048200 99 [HTML]048200 99 [HTML]048200 99 [HTML]0A8500 98 [HTML]0A8500 98 [HTML]0E8700 97 [HTML]0E8700 97 [HTML]209000 94 Table 39: Results on the BABILong dataset for Gemma-3-1B-IT model with ARMT, tasks QA4-QA5. ARMT outperforms base model. Appendix J 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 pre-training. 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. 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 [HTML]5CAE00 0.795 [HTML]7EBF00 0.756 [HTML]7EBF00 0.756 [HTML]309800 0.846 [HTML]5CAE00 0.795 [HTML]74BA00 0.769 [HTML]74BA00 0.769 1k-2k [HTML]62B100 0.790 [HTML]8AC500 0.743 [HTML]8AC500 0.743 [HTML]269300 0.857 [HTML]0E8700 0.886 [HTML]48A400 0.819 [HTML]58AC00 0.800 2k-4k [HTML]70B800 0.774 [HTML]90C800 0.736 [HTML]90C800 0.736 [HTML]56AB00 0.802 [HTML]56AB00 0.802 [HTML]9ACD00 0.726 [HTML]80C000 0.755 4k-6k [HTML]76B00 0.767 [HTML]B0D800 0.699 [HTML]8C600 0.740 [HTML]389C00 0.836 [HTML]68B400 0.781 [HTML]9ACD00 0.726 [HTML]46A300 0.822 6k-8k [HTML]249200 0.859 [HTML]54A00 0.804 [HTML]249200 0.859 [HTML]008000 0.902 [HTML]008000 0.902 [HTML]088400 0.891 [HTML]1C8E00 0.870 8k-10k [HTML]2A9500 0.852 [HTML]82C100 0.753 [HTML]40A000 0.827 [HTML]148A00 0.877 [HTML]40A000 0.827 [HTML]209000 0.864 [HTML]209000 0.864 10k-12k [HTML]1C8E00 0.868 [HTML]60B000 0.791 [HTML]1C8E00 0.868 [HTML]148A00 0.879 [HTML]0A8500 0.890 [HTML]309800 0.846 [HTML]0A8500 0.890 12k-14k [HTML]64B200 0.787 [HTML]88C400 0.745 [HTML]76B00 0.766 [HTML]6CB600 0.777 [HTML]76B00 0.766 [HTML]80C000 0.755 [HTML]48A400 0.819 14k-16k [HTML]78BC00 0.764 [HTML]BCDE00 0.685 [HTML]B2D900 0.697 [HTML]78BC00 0.764 [HTML]78BC00 0.764 [HTML]A8D400 0.708 [HTML]8C600 0.742 16k-24k [HTML]4EA700 0.812 [HTML]7EBF00 0.757 [HTML]7EBF00 0.757 [HTML]108800 0.882 [HTML]3C9E00 0.833 [HTML]7EBF00 0.757 [HTML]48A400 0.819 24k-32k [HTML]A4D200 0.713 [HTML]A4D200 0.713 [HTML]4EA700 0.812 [HTML]56AB00 0.802 [HTML]7ABD00 0.762 [HTML]C8E400 0.673 [HTML]8AC500 0.743 32k-49k [HTML]FFEE00 0.591 [HTML]76B00 0.765 [HTML]60B000 0.791 [HTML]68B400 0.783 [HTML]96CB00 0.730 [HTML]FFD800 0.565 [HTML]FFBA00 0.530 49k-65k [HTML]F0000 0.317 [HTML]A4D200 0.713 [HTML]7ABD00 0.762 [HTML]FFB600 0.525 [HTML]E0F000 0.644 [HTML]F5E00 0.426 [HTML]F5E00 0.426 In-Domain (0k-8k) [HTML]5AAD00 0.797 [HTML]86C300 0.749 [HTML]76B00 0.767 [HTML]2E9700 0.848 [HTML]389C00 0.837 [HTML]62B100 0.788 [HTML]56AB00 0.802 OOD (8k-65k) [HTML]A8D400 0.709 [HTML]8C600 0.741 [HTML]68B400 0.783 [HTML]62B100 0.788 [HTML]6CB600 0.777 [HTML]B6DB00 0.694 [HTML]9ACD00 0.724 Long-OOD (32k-65k) [HTML]F7E00 0.463 [HTML]8C600 0.741 [HTML]6CB600 0.777 [HTML]D0E800 0.662 [HTML]B8DC00 0.690 [HTML]FFA000 0.500 [HTML]F8E00 0.481 Full (0k-65k) [HTML]8C600 0.741 [HTML]8AC500 0.744 [HTML]6CB600 0.777 [HTML]50A800 0.810 [HTML]5AAD00 0.798 [HTML]98C00 0.727 [HTML]82C100 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 Base, No Fine- Tuning Base, GR-100+, 8k Base, YaRN, GR-100+, 8k ARMT, GR-100+, 8k 0k-1k [HTML]FFCC00 0.145 [HTML]128900 0.349 [HTML]008000 0.362 [HTML]229100 0.337 1k-2k [HTML]FFAC00 0.122 [HTML]048200 0.358 [HTML]068300 0.357 [HTML]54A00 0.302 2k-4k [HTML]F9E00 0.113 [HTML]2C9600 0.330 [HTML]148A00 0.347 [HTML]8EC700 0.261 4k-6k [HTML]F8000 0.091 [HTML]42A100 0.314 [HTML]46A300 0.312 [HTML]9ACD00 0.253 6k-8k [HTML]F8400 0.094 [HTML]6AB500 0.286 [HTML]66B300 0.289 [HTML]C0E000 0.225 8k-10k [HTML]F7C00 0.089 [HTML]B6DB00 0.233 [HTML]ACD600 0.239 [HTML]C6E300 0.221 10k-12k [HTML]F3600 0.039 [HTML]F1C00 0.021 [HTML]FFFA00 0.177 [HTML]D4EA00 0.211 12k-14k [HTML]F3600 0.039 [HTML]F3000 0.035 [HTML]F5E00 0.067 [HTML]CEE700 0.215 14k-16k [HTML]F3600 0.039 [HTML]F4400 0.049 [HTML]F2000 0.024 [HTML]D6EB00 0.210 16k-24k [HTML]F4200 0.048 [HTML]F3400 0.037 [HTML]F1800 0.017 [HTML]CAE500 0.219 24k-32k [HTML]F3A00 0.042 [HTML]F4C00 0.054 [HTML]F0400 0.003 [HTML]FAFD00 0.185 32k-49k [HTML]F3E00 0.044 [HTML]F5E00 0.067 [HTML]F0000 0.000 [HTML]F000 0.171 49k-65k [HTML]F7C00 0.088 [HTML]F7C00 0.088 [HTML]F0000 0.000 [HTML]F8600 0.096 In-Domain (0k-8k) [HTML]F9E00 0.113 [HTML]309800 0.327 [HTML]289400 0.334 [HTML]7ABD00 0.275 OOD (8k-65k) [HTML]F4600 0.050 [HTML]F6800 0.074 [HTML]F8000 0.091 [HTML]D6EB00 0.211 Long-OOD (32k-65k) [HTML]F4C00 0.054 [HTML]F6400 0.072 [HTML]F0000 0.000 [HTML]FFD800 0.154 Full (0k-65k) [HTML]F7000 0.080 [HTML]EEF700 0.192 [HTML]DEEF00 0.205 [HTML]AAD500 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 Base, No Fine- Tuning Base, MT, 8k Base, YaRN, MT, 8k ARMT, MT, 8k 0k-1k [HTML]F0000 0.000 [HTML]188C00 0.795 [HTML]FFCA00 0.333 [HTML]289400 0.769 1k-2k [HTML]F0000 0.000 [HTML]1C8E00 0.790 [HTML]F9C00 0.257 [HTML]50A800 0.705 2k-4k [HTML]F0000 0.000 [HTML]3C9E00 0.736 [HTML]FFA000 0.264 [HTML]54A00 0.698 4k-6k [HTML]F0000 0.000 [HTML]5CAE00 0.685 [HTML]FFB800 0.301 [HTML]7EBF00 0.630 6k-8k [HTML]F0000 0.000 [HTML]008000 0.837 [HTML]F9800 0.250 [HTML]2E9700 0.761 8k-10k [HTML]F0000 0.000 [HTML]A4D200 0.568 [HTML]FFA600 0.272 [HTML]0C8600 0.815 10k-12k [HTML]F0600 0.011 [HTML]F0000 0.000 [HTML]F6A00 0.176 [HTML]369B00 0.747 12k-14k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F4000 0.106 [HTML]329900 0.755 14k-16k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]7EBF00 0.629 16k-24k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]4A500 0.715 24k-32k [HTML]F0C00 0.020 [HTML]F0600 0.010 [HTML]F0000 0.000 [HTML]3E9F00 0.733 32k-49k [HTML]F0400 0.009 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]40A000 0.730 49k-65k [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F1200 0.030 In-Domain (0k-8k) [HTML]F0000 0.000 [HTML]289400 0.771 [HTML]FFA800 0.277 [HTML]4A500 0.714 OOD (8k-65k) [HTML]F0200 0.005 [HTML]F2200 0.058 [HTML]F2200 0.059 [HTML]76B00 0.643 Long-OOD (32k-65k) [HTML]F0200 0.005 [HTML]F0000 0.000 [HTML]F0000 0.000 [HTML]F600 0.403 Full (0k-65k) [HTML]F0000 0.003 [HTML]FFBE00 0.313 [HTML]F5200 0.137 [HTML]66B300 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.

Instant research discovery