Source-linked AI summary
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
TL;DR
Long-context inference is costly because full-context prompting scales quadratically and requires memory that grows with context length. MoNe adds a test-time-learned neural memory to a frozen Transformer, reducing compute and peak GPU memory by approximately 80% at 128K tokens while sustaining strong performance beyond the native window.
Problem
Full-context prompting makes long-context inference expensive, with O(N^2) computation and O(N) KV-cache memory that can be prohibitive on resource-constrained devices.
Method
MoNe attaches layer-localized fast-weight neural memories to a frozen Transformer, sequentially encoding fixed-size context segments and generating query-conditioned memory tokens without rereading context.
Results
Approximately 80% lower FLOPs and peak GPU memory at 128K tokens accompany near-perfect performance across evaluated tasks, while ICL collapses beyond the native window.
Takeaways & Limitations
MoNe supports efficient long-context inference beyond the backbone’s native window with constant memory footprint and 6.4% parameter overhead.
Takeaways & Limitations
Experiments use a Qwen2.5-0.5B backbone and controlled RULER retrieval tasks, leaving larger models and naturalistic long-document workloads for future validation.
Abstract
from arXiv · showhide
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.
1. Introduction
MoNe is a lightweight neural-memory module for frozen pretrained Transformers that enables long-context reasoning without backbone retraining. Its sequential test-time learning and memory-only inference reduce dependence on context length while preserving strong performance far beyond the native context window.
- Motivation: Full-context prompting scales quadratically with context length, making long-context processing prohibitive on resource-constrained hardware and for complex reasoning.The introduction identifies practical applications requiring tens of thousands of tokens, including personalized assistants, document QA, and agentic systems.
- Method: MoNe attaches to any frozen pretrained Transformer, adding only 6.4% parameter overhead without modifying backbone weights.The module uses layer-localized gradient updates to adapt fast-weight neural memory networks during test-time learning.
- Method: MoNe sequentially processes fixed-size context segments in O(N) total FLOPs, then generates memory keys and values from query tokens without rereading context.At inference, the query attends only to generated memory tokens rather than the original context.
- Efficiency: At 128K tokens, MoNe reduces total FLOPs and peak GPU memory by approximately 80% compared with in-context learning.The fast weights occupy a constant footprint regardless of context length.
- Long-context generalization: Training on contexts up to 4K tokens enables MoNe to generalize to 128K, a 32× extrapolation, while ICL collapses beyond the native context window.The reported benchmarks include S-NIAH, MK-NIAH, and Frequent Word Extraction.
2. Related Work
Prior work addresses long-context challenges through retrieval and test-time neural memory, but each approach has important limitations in reasoning coverage or hardware efficiency. RAG retrieves independently chunked segments, while TTT methods trade full attention for fast-weight memory yet often underutilize hardware.
- Long Context Problem: Self-attention scales quadratically with sequence length, making long-context reasoning challenging for LLMs.Long-context settings include retrieval-augmented generation and extended mathematical reasoning traces.
- Long Context Problem: LLMs often fail to exploit long contexts robustly, missing buried facts and degrading as input length increases.This degradation is described as context rot.
- Retrieval-Augmented Generation (RAG): RAG extends LLMs beyond their context window by retrieving relevant documents or fragments at inference time.Personalization methods retrieve relevant portions of user behavior history to augment prompts.
- Retrieval-Augmented Generation (RAG): Embedding-based retrieval is ill-suited to tasks requiring synthesis of multiple fragmented, interdependent facts across the full context.Its independently chunked segments can fail to preserve the relationships needed for long-context reasoning.
- Test-Time Learning with Neural Memory: Test-time neural memory updates fast weights during inference to compress past context, enabling sub-quadratic sequence modeling as an alternative to self-attention.Existing TTT methods often use below 5% peak FLOPs because of small mini-batch update sizes; LaCT reaches up to 70% GPU utilization with 2K–1M-token updates.
3. MoNe: Modular Neural Memory for Pretrained Transformer Attention
MoNe augments each frozen decoder layer with a fast-weight neural memory that learns associations from context segments locally at test time. It then reads the cached memory from query tokens without rereading context, while keeping the memory KV footprint fixed and supporting arbitrarily long contexts through segment-local positions.
- Layer-local test-time learning: Each segment updates only its layer-local fast weights using activations from that layer, without propagating gradients through other layers.This makes the update modular and avoids the cross-layer backpropagation required by standard cross-entropy updates.
- Memory architecture: MoNe attaches a SwiGLU fast-weight neural memory to every frozen decoder layer, with online-updated parameters for each fast-weight head.The nonlinear memory increases stored associative capacity within the same parameter budget.
- Query-time memory retrieval: At inference, query projections read the final cached fast-weight state without updates, producing memory tokens that are combined with standard attention entries.The memory uses frozen backbone Wk and Wv projections, with only small low-rank adapters added to the frozen weights.
- Efficiency: The memory KV pair contains a fixed T entries per layer regardless of N, keeping the KV-cache footprint constant during test-time learning and inference.This addresses the standard attention cost of O(N^2) FLOPs and an O(N) KV cache for long contexts.
- Position handling: Segment-local RoPE assigns context positions modulo T, while query positions remain within the same range, enabling generalization to arbitrarily long contexts without positional interpolation.LoRA adapters and meta-parameters are trained offline and remain frozen during deployment.
4. Experiments
Experiments on three RULER retrieval and aggregation tasks show that MoNe sustains near-perfect performance from 4K to 128K tokens, while ICL collapses beyond the native 32K window and RAG plateaus on distributed-information tasks. MoNe also reduces long-context computation and memory, with layer coverage and segment size controlling the accuracy–efficiency trade-off.
- Overall Results: MoNe sustains near-perfect performance across S-NIAH, MK-NIAH, and Frequent Word Extraction at all eight evaluated context lengths.The evaluation spans 4K–32K tokens within the backbone’s native window and 48K–128K tokens beyond it.
- Computational Cost: 43% lower peak GPU memory and 34% lower total FLOPs are achieved by MoNe than ICL at 32K tokens.At 128K, ICL requires 7.07 GB and 786.3 T FLOPs for inference alone, whereas MoNe requires 1.41 GB and 149.61 T FLOPs in total.
- Computational Cost: 1.41 GB of peak GPU memory is maintained independently of context length because fast-weight updates occur locally within fixed-size segments.Inference uses fixed-size memory tokens, requiring 1.29 GB of memory and 0.64 T FLOPs.
- Layer Selection: All 24 decoder layers achieve Sub-EM scores of 1.00 / 0.98 / 0.94 at 4K / 64K / 128K, while restricting MoNe to the last 8 layers drops performance to 0.15 at 64K and 0.02 at 128K.Each additional group of 8 layers adds ≈12.6M parameters, while restricted configurations reduce FLOPs by only 11–22% at 128K.
- Segment Size: 0.98 / 0.94 at 64K / 128K is achieved with T=512, outperforming T=256 at 0.91 / 0.75 and T=128 at 0.85 / 0.53.Smaller segments degrade more sharply at longer test contexts while offering only marginal FLOPs savings.
5. Conclusion
MoNe adds modular neural memory to frozen Transformers, enabling long-context generalization without modifying backbone weights. Experiments show consistent advantages over ICL and RAG, while current validation remains limited to a small backbone and controlled RULER retrieval tasks.
- Contributions: MoNe extends frozen pretrained Transformers to context lengths far beyond the native window without modifying backbone weights.Its modular neural memory provides the long-context capability while preserving the pretrained backbone.
- Empirical advantages: MoNe consistently outperforms ICL, which experiences sharp performance degradation beyond its training limit.It also surpasses RAG on tasks requiring synthesis of distributed information.
- Limitations and Future Work: Current experiments validate long-context generalization, constant-memory inference, and modular integration using a Qwen2.5-0.5B backbone and controlled RULER retrieval tasks.Future work should test larger model scales and naturalistic workloads, including single- and multi-document QA and real-world conversational histories.
A. Experimental Details · A.1. Datasets
The evaluation uses three RULER tasks covering single-key retrieval, distractor-resistant multi-key retrieval, and frequent-word extraction. Training spans 1K–4K-token contexts, while evaluation extends from 4K to 128K tokens using sequential memory updates and final-memory-only answering.
- A.1. Datasets: Three tasks from RULER evaluate retrieval and extraction under long contexts.The tasks are S-NIAH, MK-NIAH, and Frequent Word Extraction.
- A.1. Datasets: S-NIAH embeds one adjective-noun key–value pair at random depth, requiring exact-value recall measured by Sub-EM.Sub-EM tests whether the ground-truth value appears as a substring of the output.
- A.1. Datasets: MK-NIAH plants four key–value pairs in Paul Graham essay excerpts and queries one key amid distractors, reporting Sub-EM.The remaining three pairs act as distractors during retrieval.
- A.1. Datasets: FWE uses synthetic 6-character Zipfian words with α=2.0, replacing the top-ranked word with noise and measuring variable recall over three targets.The model must identify the top-3 most frequent non-noise words.
- A.1. Datasets: 512-token segments support offline training on 30K samples for each N ∈{2, 3, . . . , 8}, covering 1K–4K-token contexts.Training targets LoRA adapters and meta-parameters η(l) and momentum projections.
- A.1. Datasets: 100 evaluation samples per N ∈{8, 16, . . . , 256} cover 4K–128K-token contexts, with segments processed sequentially to update memory before final-memory-only answering.The evaluation model answers using only the final memory state.
A.2. Implementation Details
MoNe is implemented by attaching normalized, gated SwiGLU memories to every layer of a frozen Qwen2.5-0.5B-Instruct backbone. Training uses learned decay and learning-rate mechanisms, optional LoRA adapters, and AdamW optimization with bf16 computation.
- Architecture: MoNe attaches a SwiGLU MLP to each decoder layer of frozen Qwen2.5-0.5B-Instruct, using H=4 heads and d_h=224.Keys and queries use affine rescaling, SiLU, and per-token L2 normalization; values use SiLU, while outputs are RMSNorm-normalized and SiLU-gated.
- Fast-weight updates: The implementation predicts data-dependent decay β(l) and per-token learning rates, selectively resetting momentum from chunk hidden states and renormalizing weights after updates.The learning rate is softplus-activated, and channel-wise L2 renormalization follows each update.
- Optimization: LoRA adapters use rank 128 and α=128, while training runs for one epoch with AdamW, batch 16, learning rate 10−3/5×10−5, 200-step warmup, and cosine decay to ηmin=10−5.AdamW uses β1=0.9, β2=0.95, weight decay 0.1, gradient clipping 1.0, and bf16 computation.