Paper deep dive
MoNe: Modular Neural Memory for Efficient Long Context Inference
Wonguk Cho, Kyubyung Chae, Tribhuvanesh Orekondy, Sunghyun Park, Hyoungwoo Park, Jeongho Kim, Arash Behboodi, Kyuwoong Hwang, Sungrack Yun
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%
Last extracted: 8/19/2026, 5:12:09 AM
Summary
MoNe (Modular Neural Memory) is a lightweight plugin module that attaches to frozen pretrained Transformers to enable efficient long-context inference without retraining. It uses test-time learning with layer-localized gradient updates on fast-weight neural memory networks to process context in fixed-size segments. This decouples inference cost from context length, achieving O(N) preprocessing and O(1) query cost. At 128K tokens, MoNe reduces compute and peak GPU memory by ~80% compared to In-Context Learning (ICL) with only 6.4% parameter overhead, while maintaining strong performance on RULER benchmarks where ICL degrades.
Entities (11)
Relation Signals (10)
MoNe → achievescomplexity → O(N) preprocessing
confidence 95% · achieving O(N) preprocessing and O(1) query cost
MoNe → achievescomplexity → O(1) query cost
confidence 95% · achieving O(N) preprocessing and O(1) query cost
MoNe → isa → modular neural memory
confidence 95% · We present MoNe, a lightweight modular neural memory that attaches to any frozen pretrained Transformer
Qwen2.5-0.5B-Instruct → servesasbackbonefor → MoNe
confidence 95% · We use Qwen2.5-0.5B-Instruct (25) as the frozen backbone for all methods.
MoNe → evaluatedon → RULER
confidence 90% · achieving strong performance on needle-in-a-haystack and word extraction benchmarks from RULER
MoNe → generalizesto → 128K tokens
confidence 90% · MoNe generalizes to 128K (a 32× extrapolation with no additional training)
MoNe → hasparameteroverhead → 6.4%
confidence 90% · with only 6.4% parameter overhead
MoNe → outperforms → ICL
confidence 90% · MoNe reduces both compute and peak GPU memory by approximately 80% compared to ICL
→ →
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We present MoNe, a lightweight modular neural memory that attaches to any frozen pretrained Transformer to enable long-context inference without retraining. MoNe reads context in fixed-size segments via test-time learning of fast-weight neural memory networks with layer-localized gradient updates; at inference, the memory generates keys and values from the query tokens alone, with no context tokens re-read. This two-phase design decouples inference cost from context length, achieving $O(N)$ preprocessing and $O(1)$ query cost with peak GPU memory that does not grow with $N$. At 128K tokens, MoNe reduces both compute and peak GPU memory by approximately 80% compared to ICL with only 6.4% parameter overhead. MoNe generalizes to context lengths far beyond the backbone's native window, achieving strong performance on needle-in-a-haystack and word extraction benchmarks from RULER, where ICL degrades sharply.
Tags
Links
- Source: https://arxiv.org/abs/2608.17616v1
- Canonical: https://arxiv.org/abs/2608.17616v1
Trouble viewing inline? Open PDF directly →
Full Text
41,679 characters extracted from source content.
Expand or collapse full text
MoNe: Modular Neural Memory for Efficient Long Context Inference Wonguk Cho Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Kyubyung Chae Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Tribhuvanesh Orekondy Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Sunghyun Park Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Hyoungwoo Park Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Jeongho Kim Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Arash Behboodi Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Kyuwoong Hwang Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Sungrack Yun Affiliation: Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc Correspondence to: sungrack@qti.qualcomm.com Abstract We present MoNe, a lightweight modular neural memory that attaches to any frozen pretrained Transformer to enable long-context inference without retraining. MoNe reads context in fixed-size segments via test-time learning of fast-weight neural memory networks with layer-localized gradient updates; at inference, the memory generates keys and values from the query tokens alone, with no context tokens re-read. This two-phase design decouples inference cost from context length, achieving O(N)O(N) preprocessing and O(1)O(1) query cost with peak GPU memory that does not grow with N. At 128K tokens, MoNe reduces both compute and peak GPU memory by approximately 80% compared to ICL with only 6.4% parameter overhead. MoNe generalizes to context lengths far beyond the backbone’s native window, achieving strong performance on needle-in-a-haystack and word extraction benchmarks from RULER, where ICL degrades sharply. 1 Introduction Long-context processing has become essential to practical AI applications: personalized assistants that reason over user history, document QA over lengthy contracts or medical records, and agentic systems that must synthesize information spanning tens of thousands of tokens. The dominant approach—feeding the full context as a prompt—scales quadratically in FLOPs (O(N2)O(N^2)) with context length. On resource-constrained hardware such as mobile devices, this cost becomes prohibitive; and for small models that are practical on-device, the problem is compounded: they struggle to reliably extract information from long prompts even when those prompts fit within the context window, particularly when tasks demand complex reasoning across many attended tokens. Figure 1: Comparison of ICL and MoNe. (Left) In ICL, question tokens must attend to all context tokens. (Right) In MoNe, segments are processed sequentially during test-time learning to update neural memory; at inference, question tokens attend only to memory tokens. Figure 2: End-to-end peak GPU memory and total FLOPs for MoNe (test-time learning + inference) at context lengths of 32K and 128K, in comparison with ICL. As an alternative, Retrieval-Augmented Generation (RAG) (20; 13) partially addresses the cost problem by selecting relevant chunks before generation. However, RAG fundamentally limits reasoning to locally-recoverable facts: embedding-based retrieval fails when an answer requires integrating multiple scattered fragments that are individually unremarkable. On the other hand, recurrent models such as Mamba (10), T (31), and Titans (2) achieve linear-complexity context processing, but they require designing and training entirely new architectures from scratch and cannot be seamlessly integrated with existing pretrained Transformers without full retraining. We propose MoNe (Modular Neural Memory), a lightweight plugin module that attaches to any frozen pretrained Transformer and enables efficient long-context reasoning without retraining the backbone. As presented in Fig. 1, MoNe sequentially reads context in fixed-size segments and performs test-time learning of fast-weight neural memory networks via layer-localized gradient updates—without modifying any backbone weights. After all context segments are processed, the updated neural memory generates memory tokens from the question at inference time; these memory tokens serve as keys and values in the frozen backbone’s self-attention, enabling the model to answer from the question alone, with no directly attention to any context tokens. MoNe’s efficiency stems directly from this two-phase design: during test-time learning, context segments of size T are processed sequentially in O(N)O(N) total FLOPs; at inference, the query attends only to the memory tokens it generates—not to any context tokens—so inference cost does not scale with N, compared to O(N2)O(N^2) FLOPs for in-context learning (ICL). The fast weights occupy a constant footprint regardless of context length: at 128K tokens, this reduces both compute and peak GPU memory by approximately 80% compared to ICL, with only 6.4% parameter overhead, as shown in Fig. 2. The final fast-weight state can additionally be reused across multiple queries and incrementally extended as new context arrives, without reprocessing past content. MoNe generalizes to 128K (a 32× extrapolation with no additional training) with strong accuracy on S-NIAH, MK-NIAH, and Frequent Word Extraction benchmarks from RULER, while ICL collapses beyond the model’s native context window. Our contributions can be summarized as follows: • We propose a modular neural memory architecture that plugs into any pretrained Transformer without backbone modification, adding only 6.4% parameter overhead. • We introduce a test-time learning procedure with layer-localized gradient updates that adapts per-layer fast-weight neural memory networks to encode long contexts in O(N)O(N) preprocessing and O(1)O(1) query cost—compared to O(N2)O(N^2) for ICL—with natural support for incremental context extension and multi-query reuse. • We demonstrate that MoNe trained on contexts up to 4K tokens generalizes to 128K (a 32× extrapolation), achieving strong performance on S-NIAH, MK-NIAH, and Frequent Word Extraction where ICL collapses beyond the native context window, while reducing total FLOPs and peak GPU memory by approximately 80% at 128K. 2 Related Work In this section, we discuss how our work on enabling efficient long-context understanding ties with prior literature. Long Context Problem. Reasoning over long contexts remains a central challenge for LLMs, in part because self-attention scales quadratically with sequence length. In practice, long-context settings arise when generation is conditioned on large external context, such as documents in retrieval-augmented generation (RAG) (32; 36), or on extended reasoning traces, such as those produced during mathematical reasoning (11). In this work, we focus on the former setting, where answering a query requires analyzing, understanding, and retrieving relevant information from a long input context. Despite increasingly large context windows, LLMs often fail to robustly exploit long contexts, missing facts buried in documents or failing to preserve information from distant positions (21; 15). This limitation is further reflected in the context rot phenomenon (14), where increasing input length degrades a model’s ability to use relevant information already present in the context, leading to systematic performance drops even when task-relevant evidence is retained. Prior work has sought to improve long-context performance through long-context fine-tuning (4; 9), context compression (33; 19), and memory-augmented architectures or systems (5; 38; 2). Our work belongs to the memory-augmented line of research, but differs by studying a lightweight neural memory model that is trained at inference time. Retrieval-Augmented Generation (RAG). RAG (20; 13) retrieves relevant documents at inference time to extend LLMs beyond their context window (18; 30). In the context of personalization, RAG-based methods select the most relevant fragments of a user’s behavior history to augment the prompt (27; 23), yielding strong results on personalization benchmarks. Despite its effectiveness, RAG faces fundamental limitations for personalized on-device agentic AI. The computational overhead of repeatedly searching through large local user contexts (e.g., lengthy chat histories and app usage logs) at every query is prohibitive on resource-constrained devices. More critically, embedding-based retrieval operates on independently chunked segments, making it ill-suited for long-context reasoning tasks that require identifying and synthesizing multiple fragmented, interdependent facts distributed across the full user context (30; 12). Test-Time Learning with Neural Memory. Test-Time Training (T) (31) has emerged as an alternative paradigm for long-context modeling, where a small sub-network with rapidly adaptable fast weights is updated at inference time to compress and store past context as neural memory (28; 2; 34; 10). These fast weights function similarly to recurrent states in RNNs, enabling sub-quadratic sequence modeling as an alternative to full self-attention. However, existing T methods suffer from low hardware utilization (often below 5% peak FLOPs) due to small mini-batch update sizes. LaCT (37) addresses this by adopting extremely large chunk updates (2K–1M tokens), achieving up to 70% GPU utilization and enabling significantly larger fast-weight state sizes. Despite these advances, all such approaches require training new architectures from scratch, making them difficult to integrate with existing pretrained Transformer-based models. 3 MoNe: Modular Neural Memory for Pretrained Transformer Attention 3.1 Problem Setup We consider a pretrained L-layer Transformer with hidden dimension d whose weights are kept completely frozen. Given a context =(t1,…,tN)C=(t_1,…,t_N) of N tokens and a query q, our goal is to enable efficient and effective long-context inference without modifying any backbone parameter. We partition C into S non-overlapping segments 1,…,Ss_1,…,s_S of T=512T=512 tokens each, with s∈1,…,Ss∈\1,…,S\ and s=0s=0 denoting the initial state before any segment is processed. At each layer l, self-attention projects the segment hidden states (l)∈ℝT×dX^(l) ^T× d into queries, keys, and values via frozen weight matrices qW_q, kW_k, vW_v: =q(l),=k(l),=v(l),Q=W_qX^(l), =W_kX^(l), =W_vX^(l), (1) and computes the attention output as: Attention(,,)=softmax(⊤dh),Attention(Q,K,V)=softmax\! ( QK d_h )V, (2) where dhd_h is the per-head dimension. Processing a long context of N tokens incurs O(N2)O(N^2) FLOPs and requires a KV cache that grows as O(N)O(N), making long-context inference prohibitively expensive on resource-constrained devices. Figure 3: An overview of the update process of MoNe during test-time learning. For each layer, neural memory fast weights are updated using the associative memory loss, which is localized and computed entirely from layer l’s own forward-pass activations, making the memory update both efficient and truly modular. 3.2 Memory Injection into Backbone Attention MoNe attaches a fast-weight neural memory to each decoder layer l (Fig. 3). Following 37, we use a SwiGLU MLP for the memory module of each layer l, defined as ℳ(⋅,(l))M(·\,;\,W^(l)) with fast-weight parameters (l)=in,gate,outW^(l)=\W_in,W_gate,W_out\: ℳ(,(l))=out⊤(SiLU(in)⊙gate),M(x;W^(l))\;=\;W_out \! (SiLU\! (W_in\,x ) _gate\,x ), (3) where all three matrices are ℝdh×dhR^d_h× d_h per fast-weight head and are updated online as context segments are consumed. Using a nonlinear network rather than a single linear map substantially increases the expressive capacity of the stored memory, allowing each layer to encode richer associative structure within the same parameter budget. At generation time, the query tokens are projected via qW_q and used to read from the cached fast-weight state S(l)W^(l)_S (the state after all S segments have been processed), producing layer-wise memory tokens: j=ℳ(qj,S(l)).h_j=M\! (W_q\,x_j;\;W^(l)_S ). (4) The raw output is normalized per fast-weight head with RMSNorm and scaled by a learned per-head gate, producing the memory tokens jh_j used below. These memory tokens are projected into key-value space using the same frozen backbone matrices kW_k and vW_v, and concatenated alongside the standard attention entries derived from the query tokens: out=Attention(,[∥k],[∥v]),out=Attention\! (Q,\;[K\,\|\,W_kh],\;[V\,\|\,W_vh] ), (5) where Q, K, V are the standard attention projections of the query tokens per Eq. (1). While qW_q, kW_k, and vW_v are shared directly from the backbone without duplication, MoNe adds only small low-rank adapters (16) to these frozen weights. The memory KV pair has a fixed size of T entries per layer regardless of N, keeping the KV-cache footprint constant during both test-time learning and inference. 3.3 Associative Memory Loss and Update During test-time learning, the memory module ℳM is optimized via an associative memory loss. The key-value training targets are produced by kW_k and vW_v, which are frozen backbone weights equipped with meta-trained LoRA adapters: s,j(l)=ks,j(l),s,j(l)=vs,j(l),k_s,j^(l)=W_k\,x_s,j^(l), _s,j^(l)=W_v\,x_s,j^(l), (6) where j∈1,…,Tj∈\1,…,T\ indexes tokens within segment s. In practice, s,j(l)k_s,j^(l) is further passed through SiLU activation and per-token L2 normalization after projection, and s,j(l)v_s,j^(l) through SiLU activation; see Appendix A for the full pipeline. Following 37, we define the associative memory loss as the negative inner product between the memory output and the target value: ℓs,j(l)((l))=−(s,j(l))⊤ℳ(s,j(l),(l)), _s,j^(l)\! (W^(l) )\;=\;- (v_s,j^(l) ) \,M\! (k_s,j^(l);\,W^(l) ), (7) with segment-level average ℒs(l)=1T∑j=1Tℓs,j(l)L_s^(l)= 1T _j=1^T _s,j^(l). Minimizing ℓs,j(l) _s,j^(l) directly imprints the association s,j(l)→s,j(l)k_s,j^(l)\!→\!v_s,j^(l) into the fast weights by maximizing the inner product between the memory output and the target value, without requiring a prediction error. The fast weights are updated at each segment via a gradient step: s(l)=s−1(l)−s(l),W^(l)_s\;=\;W^(l)_s-1\;-\; μ_s^(l), (8) where s(l) μ_s^(l) accumulates gradients of ℓs,j(l) _s,j^(l) with per-token learning rates and data-dependent momentum decay (Appendix A). Positional information is incorporated into s,j(l)k_s,j^(l) via segment-local RoPE: each context token is assigned position pjlocal=jmodTp_j^local=j T, so position indices always lie in [0,T)[0,T) regardless of N. At inference time, query tokens receive positions in [0,Q)[0,Q) where Q is the query length; since Q≪TQ T, this falls within the same [0,T)[0,T) range with no positional interpolation required, allowing MoNe to generalize to arbitrarily long contexts. The LoRA adapters and meta-parameters ((l) η^(l), momentum projections, output scale) are trained offline via a generation loss on answer tokens conditioned on the updated memory S(l)W^(l)_S and remain frozen during test-time learning and inference, so that the fast weight update operates as a fixed, pre-learned algorithm at deployment. 3.4 Test-Time Learning and Inference During test-time learning, segments are consumed sequentially and the fast weights at each layer l are updated per Eq. (8): 0(l)→1(l)→⋯→S(l).W^(l)_0\;→\;W^(l)_1\;→\;·s\;→\;W^(l)_S. (9) Each update step computes the per-token gradient of ℓs,j(l) _s,j^(l) locally at layer l: ∇s−1(l)ℓs,j(l)=−s,j(l)∂ℳ(s,j(l),s−1(l))⊤∂s−1(l), _W^(l)_s-1 _s,j^(l)\;=\;-v_s,j^(l)\, ∂\,M\! (k_s,j^(l);W^(l)_s-1 ) ^(l)_s-1, (10) carrying no gradient through any other layer: ∂ℓs,j(l)∂s−1(l′)=∀l′≠l. ∂ _s,j^(l) ^(l )_s-1=0 ∀\,l ≠ l. (11) Each update is thus computed entirely from layer l’s own forward-pass activations, making the memory update both efficient and truly modular. By contrast, updating s−1(l)W^(l)_s-1 via a standard cross-entropy loss would require propagating gradients from the output logits back through all subsequent layers s−1(L),…,s−1(l+1)W^(L)_s-1,…,W^(l+1)_s-1 before reaching layer l. During inference, the query tokens are projected via qW_q and the memory is read without any weight update per Eq. (4), producing memory tokens supplied to the frozen backbone in Eq. (5). Table 1: Performance of ICL, RAG, and MoNe (ours) on S-NIAH, MK-NIAH, and Frequent Word Extraction across varying context lengths. Results are grouped by whether the input length falls within or beyond the model’s training context length (32K tokens), assessing generalization to out-of-training-distribution sequence lengths. Within model context length Beyond model context length Task Method 4K 8K 16K 32K 48K 64K 96K 128K S-NIAH ICL 0.95 0.98 0.94 0.94 0.64 0.60 0.42 0.28 RAG 0.94 0.93 0.91 0.93 0.93 0.93 0.97 0.89 MoNe (ours) 1.00 1.00 0.99 1.00 1.00 0.99 0.99 0.96 MK-NIAH ICL 0.89 0.84 0.83 0.93 0.41 0.13 0.11 0.00 RAG 0.92 0.91 0.88 0.79 0.80 0.66 0.74 0.71 MoNe (ours) 1.00 0.99 0.99 0.99 0.99 0.98 0.97 0.94 Frequent Word Extraction ICL 0.59 0.57 0.42 0.41 0.29 0.31 0.28 0.23 RAG 0.58 0.61 0.60 0.61 0.61 0.59 0.60 0.60 MoNe (ours) 1.00 1.00 1.00 1.00 0.99 0.99 0.96 0.96 4 Experiments In this section, we first detail the experimental setup, and follow up by presenting quantitative results and a detailed computation cost analysis. 4.1 Experimental Setup Datasets. We evaluate on three tasks drawn from the RULER benchmark (15), each designed to require precise retrieval or aggregation over the full context rather than surface-level pattern matching. We use S-NIAH (Single Needle-in-a-Haystack), MK-NIAH (Multi-Key Needle-in-a-Haystack), and FWE (Frequent Word Extraction) tasks. Evaluation Metric. S-NIAH and MK-NIAH use substring exact match (Sub-EM): whether the ground-truth value appears as a substring of the model’s output (1). For FWE, we measure variable recall: the fraction of the three target words that appear in the model’s output. All tasks are evaluated at context lengths of 4K, 8K, 16K, and 32K tokens (within the backbone’s native window) and 48K, 64K, 96K, and 128K tokens (beyond the native window). Further training and evaluation details are described in the Appendix A. Baselines. We use Qwen2.5-0.5B-Instruct (25) as the frozen backbone for all methods. In-Context Learning (ICL) feeds the full context directly as a prompt, up to the backbone’s native 32K context window; performance degrades sharply beyond this limit as the model must attend over increasingly many tokens. Retrieval-Augmented Generation (RAG) treats the long context as an external memory from which relevant information can be retrieved at query time—a natural baseline that avoids quadratic attention cost by first compressing the context into a retrievable index (20; 13; 18). We split each context into chunks of 128 tokens and use BGE-Large (3) to compute embeddings; the top-K chunks with highest cosine similarity to the query are concatenated as the prompt. We evaluate K∈1,4,8K∈\1,4,8\ and report the best score at each context length. 4.2 Evaluation Results Table 1 reports results of all three methods across three tasks and eight context lengths, covering both within and beyond the backbone’s native 32K window. The results reveal a consistent pattern: MoNe sustains near-perfect performance across all tasks, while ICL collapses once the context exceeds the native limit and RAG plateaus on tasks that require integrating information distributed across the full context. Within the Native Context Window (4K–32K). MoNe achieves near-perfect Sub-EM across all tasks and context lengths (0.99–1.00), substantially outperforming ICL (0.41–0.98) and RAG (0.58–0.93). ICL’s degradation even within the native window is most pronounced on Frequent Word Extraction, where attending to all tokens is necessary but the backbone’s limited capacity is overwhelmed; MoNe’s compressed representation is more robust. Beyond the Native Context Window (48K–128K). ICL degrades sharply once the context exceeds the 32K window: from 0.64 on S-NIAH at 48K to 0.28 at 128K, and from 0.41 to 0.00 on MK-NIAH. RAG is more stable but plateaus below 0.97 on S-NIAH and below 0.80 on MK-NIAH, as embedding-based retrieval struggles with multi-hop facts distributed across many passages. MoNe maintains 0.96, 0.94, and 0.96 Sub-EM on S-NIAH, MK-NIAH, and Frequent Word Extraction at 128K tokens, despite being trained exclusively on contexts up to 4K tokens. This generalization in context length with no additional training or positional interpolation is enabled by the segment-local RoPE encoding (§3.3), which keeps position indices bounded in [0,T)[0,T) regardless of total context length. Table 2: Computational cost comparison at different context lengths. Values in parentheses denote reduction rate relative to ICL (↓ larger is better). Method Peak GPU Usage Total FLOPs Context length = 32K ICL 2.48 GB 58.06 T MoNe (total) 1.41 GB (↓43% 43\%) 38.10 T (↓34% 34\%) ⌞ Test-time learning 1.41 GB 37.24 T ⌞ Inference 1.29 GB 0.86 T Context length = 128K ICL 7.07 GB 786.33 T MoNe (total) 1.41 GB (↓80% 80\%) 149.61 T (↓81% 81\%) ⌞ Test-time learning 1.41 GB 148.97 T ⌞ Inference 1.29 GB 0.64 T Figure 4: Ablation study on performance-efficiency trade-off: (Left) Layer selection. Attaching MoNe to all 24 decoder layers achieves the highest accuracy across all context lengths; restricting to the last 16 layers (last 23 23) yields modest degradation beyond 32K, while using only the last 8 layers (last 13 13) collapses beyond 32K. (Right) Segment size. T=512T=512 (default) generalizes from 4K training contexts to 128K evaluation contexts, whereas T=128T=128 and T=256T=256 degrade rapidly past 16K, yet offer only marginal efficiency gains over the default. 4.3 Computational Cost Analysis Table 2 compares the end-to-end computational cost of ICL and MoNe, including both test-time learning and inference. At 32K tokens, MoNe reduces peak GPU memory by 43% (1.41 GB vs. 2.48 GB) and total FLOPs by 34% (38.10 T vs. 58.06 T). The advantage becomes more pronounced at longer contexts. At 128K tokens, ICL requires 7.07 GB and 786.3 T FLOPs for inference alone, whereas MoNe requires only 1.41 GB and 149.61 T FLOPs in total. Importantly, MoNe’s peak GPU memory is independent of the context length N. During test-time learning, fast-weight updates are computed locally within fixed-size segments (Eq. (8)), resulting in a constant 1.41 GB. At inference time, queries attend only to the fixed-size memory tokens, requiring 1.29 GB of memory and 0.64 T FLOPs. 4.4 Performance-Efficiency Trade-off Figure 4 presents two ablations on MK-NIAH, examining how layer coverage and segment size each contribute to long-context generalization. Layer Selection. We ablate the effect of attaching MoNe to different subsets of the 24 decoder layers. Figure 4(a) shows that attaching MoNe to all 24 layers achieves Sub-EM scores of 1.00 / 0.98 / 0.94 at 4K / 64K / 128K. Restricting MoNe to the last 16 layers (the final 23 23) preserves short-context performance (1.00 at 4K) but degrades at longer contexts (0.88 at 64K, 0.70 at 128K). Restricting further to the last 8 layers (the final 13 13) leads to a substantial collapse, dropping to 0.15 at 64K and 0.02 at 128K. Each additional group of 8 layers adds only ≈12.6≈12.6M parameters (≈2.1%≈2.1\% of backbone size), and Fig. 4(b) shows that the FLOPs overhead is equally modest: at 128K, all-24L incurs 149.1T FLOPs versus 132.6T for last-16L and 116.1T for last-8L—the restricted configurations reduce FLOPs by only 11–22% while suffering a far greater loss in performance. Segment Size. We compare three segment sizes T∈128,256,512T∈\128,256,512\ across the same context lengths. Increasing the segment size consistently improves long-context performance: at 64K / 128K, T=512T=512 achieves 0.98 / 0.94, outperforming T=256T=256 (0.91 / 0.75) and T=128T=128 (0.85 / 0.53). Although all configurations are trained on contexts of up to 4K tokens, smaller segments degrade more sharply at longer test contexts. As shown in Fig. 4(d), smaller segments offer only marginal FLOPs savings; as shown in Fig. 4(c), this comes at a substantial cost to long-context performance. We therefore adopt T=512T=512, which delivers the strongest performance among the tested configurations. 5 Conclusion MoNe augments frozen, pretrained Transformers with long-context capabilities via modular neural memory, requiring no modifications to the backbone weights. Although trained on relatively short sequences, MoNe generalizes to context lengths far exceeding the native window. In our experiments, MoNe consistently outperforms ICL, which suffers from sharp performance degradation beyond the training limit. It also surpasses RAG on tasks requiring the synthesis of distributed information. These performance gains are achieved with high efficiency; at 128K tokens, MoNe reduces both FLOPs and peak GPU memory by approximately 80% compared to ICL. Crucially, its memory footprint remains constant regardless of context length. Finally, our ablation studies show that full layer coverage with a segment size of 512 optimizes the performance-efficiency trade-off. MoNe is thus highly effective for resource-constrained applications, such as on-device personalization (6; 24) and continual learning (7; 8). Limitations and Future Work. Current experiments validate MoNe’s core properties—long-context generalization, constant-memory inference, and modular integration—using a Qwen2.5-0.5B backbone and the controlled retrieval tasks in RULER. Extending these results to larger model scales and to naturalistic long-document workloads such as single/multi-document QA (35; 26) and real-world conversational histories (22; 29) is a natural next step: MoNe’s plug-in design requires no architectural changes to accommodate larger backbones, and the efficiency advantages of constant-memory inference are expected to become even more pronounced at scale, where KV-cache cost is otherwise extremely prohibitive. A further direction is to apply different LoRA-based approaches (16; 17) for each neural memory to improve efficiency and enable flexible management of multiple memory modules. References Asai et al. (2024) A. Asai, Z. Wu, Y. Wang, A. Sil, and H. Hajishirzi Self-rag: learning to retrieve, generate, and critique through self-reflection. In International conference on learning representations, Vol. 2024, p. 9112–9141. Cited by: §4.1. Behrouz et al. (2026) A. Behrouz, P. Zhong, and V. Mirrokni Titans: learning to memorize at test time. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, External Links: Link Cited by: §1, §2, §2. Chen et al. (2024a) J. Chen, S. Xiao, P. Zhang, K. Luo, D. Lian, and Z. Liu M3-embedding: multi-linguality, multi-functionality, multi-granularity text embeddings through self-knowledge distillation. In Findings of the Association for Computational Linguistics: ACL 2024, L. Ku, A. Martins, and V. Srikumar (Eds.), Bangkok, Thailand, p. 2318–2335. External Links: Link, Document Cited by: §4.1. Chen et al. (2024b) Y. Chen, S. Qian, H. Tang, X. Lai, Z. Liu, S. Han, and J. Jia Longlora: efficient fine-tuning of long-context large language models. In International Conference on Learning Representations, Vol. 2024, p. 8220–8238. Cited by: §2. Chhikara et al. (2025) P. Chhikara, D. Khant, S. Aryan, T. Singh, and D. Yadav Mem0: building production-ready ai agents with scalable long-term memory. arXiv preprint arXiv:2504.19413. Cited by: §2. Cho et al. (2024) W. Cho, S. Choi, D. Das, M. Reisser, T. Kim, S. Yun, and F. Porikli Hollowed net for on-device personalization of text-to-image diffusion models. Advances in Neural Information Processing Systems 37, p. 43058–43079. Cited by: §5. Cho et al. (2023) W. Cho, J. Park, and T. Kim Complementary domain adaptation and generalization for unsupervised continual domain shift learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 11442–11452. Cited by: §5. Dorovatas et al. (2026) V. Dorovatas, M. Schwerin, A. D. Bagdanov, L. Caccia, A. Carta, L. Charlin, B. Hammer, T. L. Hayes, T. Hess, C. Kanan, et al. Modular memory is the key to continual learning agents. arXiv preprint arXiv:2603.01761. Cited by: §5. Gao et al. (2025) T. Gao, A. Wettig, H. Yen, and D. Chen 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. Cited by: §2. Gu and Dao (2024) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, External Links: Link Cited by: §1, §2. Guo et al. (2025) D. Guo, D. Yang, H. Zhang, J. Song, P. Wang, Q. Zhu, R. Xu, R. Zhang, S. Ma, X. Bi, et al. Deepseek-r1: incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948. Cited by: §2. Gupta et al. (2024) A. Gupta, A. Shirgaonkar, A. d. L. Balaguer, B. Silva, D. Holstein, D. Li, J. Marsman, L. O. Nunes, M. Rouzbahman, M. Sharp, et al. RAG vs fine-tuning: pipelines, tradeoffs, and a case study on agriculture. arXiv preprint arXiv:2401.08406. Cited by: §2. Guu et al. (2020) K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang REALM: retrieval-augmented language model pre-training. In Proceedings of the 37th International Conference on Machine Learning, ICML’20. Cited by: §1, §2, §4.1. Hong et al. (2025) K. Hong, A. Troynikov, and J. Huber Context rot: how increasing input tokens impacts llm performance. URL https://research. trychroma. com/context-rot, retrieved October 20, p. 2025. Cited by: §2. Hsieh et al. (2024) C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, and B. Ginsburg RULER: what’s the real context size of your long-context language models?. In First Conference on Language Modeling, External Links: Link Cited by: §A.1, §2, §4.1. Hu et al. (2022) E. J. Hu, yelong shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen LoRA: low-rank adaptation of large language models. In International Conference on Learning Representations, External Links: Link Cited by: §3.2, §5. Hwang et al. (2025) J. Hwang, W. Cho, and T. Kim PiCa: parameter-efficient fine-tuning with column space projection. arXiv preprint arXiv:2505.20211. Cited by: §5. Izacard and Grave (2021) G. Izacard and E. Grave Leveraging passage retrieval with generative models for open domain question answering. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, P. Merlo, J. Tiedemann, and R. Tsarfaty (Eds.), Online, p. 874–880. External Links: Link, Document Cited by: §2, §4.1. Kontonis et al. (2026) V. Kontonis, Y. Zeng, S. Garg, L. Chen, H. Tang, Z. Wang, A. Awadallah, E. Horvitz, J. Langford, and D. Papailiopoulos MEMENTO: teaching llms to manage their own context. arXiv preprint arXiv:2604.09852. Cited by: §2. Lewis et al. (2020) P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W. Yih, T. Rocktäschel, S. Riedel, and D. Kiela Retrieval-augmented generation for knowledge-intensive nlp tasks. In Advances in Neural Information Processing Systems, H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (Eds.), Vol. 33, p. 9459–9474. External Links: Link Cited by: §1, §2, §4.1. Liu et al. (2024) N. F. Liu, K. Lin, J. Hewitt, A. Paranjape, M. Bevilacqua, F. Petroni, and P. Liang Lost in the middle: how language models use long contexts. Transactions of the Association for Computational Linguistics 12, p. 157–173. External Links: Link, Document Cited by: §2. Maharana et al. (2024) A. Maharana, D. Lee, S. Tulyakov, M. Bansal, F. Barbieri, and Y. Fang Evaluating very long-term conversational memory of llm agents. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), p. 13851–13870. Cited by: §5. Mysore et al. (2024) S. Mysore, Z. Lu, M. Wan, L. Yang, B. Sarrafzadeh, S. Menezes, T. Baghaee, E. B. Gonzalez, J. Neville, and T. Safavi Pearl: personalizing large language model writing assistants with generation-calibrated retrievers. In Proceedings of the 1st Workshop on Customizable NLP: Progress and Challenges in Customizing NLP for a Domain, Application, Group, or Individual (CustomNLP4U), S. Kumar, V. Balachandran, C. Y. Park, W. Shi, S. A. Hayati, Y. Tsvetkov, N. Smith, H. Hajishirzi, D. Kang, and D. Jurgens (Eds.), Miami, Florida, USA, p. 198–219. External Links: Link, Document Cited by: §2. Park et al. (2026) S. Park, J. Kim, H. Park, D. Das, S. Yun, M. Hayat, J. Choo, F. Porikli, and S. Choi Memory-efficient fine-tuning diffusion transformers via dynamic patch sampling and block skipping. arXiv preprint arXiv:2603.20755. Cited by: §5. Qwen et al. (2025) Qwen, :, A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, H. Lin, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Lin, K. Dang, K. Lu, K. Bao, K. Yang, L. Yu, M. Li, M. Xue, P. Zhang, Q. Zhu, R. Men, R. Lin, T. Li, T. Tang, T. Xia, X. Ren, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Wan, Y. Liu, Z. Cui, Z. Zhang, and Z. Qiu Qwen2.5 technical report. External Links: 2412.15115, Link Cited by: §4.1. Rajpurkar et al. (2016) P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang Squad: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 conference on empirical methods in natural language processing, p. 2383–2392. Cited by: §5. Salemi et al. (2024) A. Salemi, S. Mysore, M. Bendersky, and H. Zamani LaMP: when large language models meet personalization. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), L. Ku, A. Martins, and V. Srikumar (Eds.), Bangkok, Thailand, p. 7370–7392. External Links: Link, Document Cited by: §2. Schmidhuber (1992) J. Schmidhuber Learning to control fast-weight memories: an alternative to dynamic recurrent networks. Neural Comput. 4 (1), p. 131–139. External Links: ISSN 0899-7667, Link, Document Cited by: §2. Shaham et al. (2023) U. Shaham, M. Ivgi, A. Efrat, J. Berant, and O. Levy ZeroSCROLLS: a zero-shot benchmark for long text understanding. In Findings of the Association for Computational Linguistics: EMNLP 2023, p. 7977–7989. Cited by: §5. Shi et al. (2023) F. Shi, X. Chen, K. Misra, N. Scales, D. Dohan, E. Chi, N. Schärli, and D. Zhou Large language models can be easily distracted by irrelevant context. In Proceedings of the 40th International Conference on Machine Learning, ICML’23. Cited by: §2. Sun et al. (2025) Y. Sun, X. Li, K. Dalal, J. Xu, A. Vikram, G. Zhang, Y. Dubois, X. Chen, X. Wang, S. Koyejo, T. Hashimoto, and C. Guestrin Learning to (learn at test time): RNNs with expressive hidden states. In Forty-second International Conference on Machine Learning, External Links: Link Cited by: §1, §2. Vodrahalli et al. (2024) K. Vodrahalli, S. Ontanon, N. Tripuraneni, K. Xu, S. Jain, R. Shivanna, J. Hui, N. Dikkala, M. Kazemi, B. Fatemi, et al. Michelangelo: long context evaluations beyond haystacks via latent structure queries. arXiv preprint arXiv:2409.12640. Cited by: §2. Yang et al. (2025) C. Yang, N. Srebro, D. McAllester, and Z. Li Pencil: long thoughts with short memory. arXiv preprint arXiv:2503.14337. Cited by: §2. Yang et al. (2024) S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim Gated linear attention transformers with hardware-efficient training. In Proceedings of the 41st International Conference on Machine Learning, ICML’24. Cited by: §2. Yang et al. (2018) Z. Yang, P. Qi, S. Zhang, Y. Bengio, W. Cohen, R. Salakhutdinov, and C. D. Manning HotpotQA: a dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 conference on empirical methods in natural language processing, p. 2369–2380. Cited by: §5. Yen et al. (2025) H. Yen, T. Gao, M. Hou, K. Ding, D. Fleischer, P. Izsak, M. Wasserblat, and D. Chen Helmet: how to evaluate long-context language models effectively and thoroughly. ICLR. Cited by: §2. Zhang et al. (2026) T. Zhang, S. Bi, Y. Hong, K. Zhang, F. Luan, S. Yang, K. Sunkavalli, W. T. Freeman, and H. Tan Test-time training done right. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §A.2, §2, §3.2, §3.3. Zhou et al. (2025) Z. Zhou, A. Qu, Z. Wu, S. Kim, A. Prakash, D. Rus, J. Zhao, B. K. H. Low, and P. P. Liang Mem1: learning to synergize memory and reasoning for efficient long-horizon agents. arXiv preprint arXiv:2506.15841. Cited by: §2. Appendix A Experimental Details A.1 Datasets From the RULER benchmark (15), we use the following three tasks: • S-NIAH (Single Needle-in-a-Haystack) embeds a single key–value pair of adjective-noun compound words at a random depth within a passage of filler text; the model must recall the exact value given the key. We measure Sub-EM: whether the ground-truth value appears as a substring of the model’s output. • MK-NIAH (Multi-Key Needle-in-a-Haystack) plants four key–value pairs of adjective-noun compound words at scattered positions within Paul Graham essay excerpts; the model is queried for the value of one specific key while the remaining pairs serve as distractors. Sub-EM is reported on the queried value. • Frequent Word Extraction (FWE) fills the context with synthetic 6-character words sampled from a Zipfian distribution (α=2.0α=2.0), with the top-ranked word replaced by a noise token; the model must identify the top-3 most frequent non-noise words. We measure variable recall: the fraction of the three target words that appear in the model’s output. With a segment size of 512 tokens, we generate 30K samples for each N∈2,3,…,8N∈\2,3,…,8\, covering context lengths from 1K to 4K tokens, for offline training of LoRA adapters and meta-parameters ((l) η^(l), momentum projections). For evaluation, we generate 100 samples per N∈8,16,…,256N∈\8,16,…,256\, ranging from 4K to 128K context lengths. At test time, context segments are processed sequentially to update the neural memory; the model then answers using only the final memory state, with no direct access to the context. A.2 Implementation Details We use Qwen2.5-0.5B-Instruct as the frozen backbone, attaching a SwiGLU MLP to each decoder layer (H=4H=4 heads, dh=224d_h=224). Following 37, keys and queries undergo affine rescaling, SiLU, and per-token L2 normalization, values SiLU only, and the output is RMSNorm-normalized and SiLU-gated before injection into attention. Weights are channel-wise L2-renormalized after each update, and the gradient uses a learned per-token learning rate (softplus-activated) and data-dependent momentum decay (sigmoid-activated). The data-dependent decay βs(l)∈(0,1) _s^(l)∈(0,1) is predicted from each chunk’s hidden states, allowing the momentum to be selectively reset based on content. For LoRA adapters, rank 128 and α=128α=128 are used. We train for one epoch with AdamW (β1=0.9 _1=0.9, β2=0.95 _2=0.95, wd 0.10.1), batch 16, lr 10−310^-3/5×10−55×10^-5 (non-LoRA/LoRA), 200-step warmup and cosine decay to ηmin=10−5 _ =10^-5, gradient clip 1.0, bf16.