🤖 AI

MoNe: Modular Neural Memory for Efficient Long Context Inference

MoNe is a lightweight, modular neural memory system that attaches to frozen Transformers to enable efficient long-context inference by decoupling inference cost from context length, achieving constant query complexity and significantly reduced memory usage without retraining.

Original authors: Wonguk Cho, Kyubyung Chae, Tribhuvanesh Orekondy, Sunghyun Park, Hyoungwoo Park, Jeongho Kim, Arash Behboodi, Kyuwoong Hwang, Sungrack Yun

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wonguk Cho, Kyubyung Chae, Tribhuvanesh Orekondy, Sunghyun Park, Hyoungwoo Park, Jeongho Kim, Arash Behboodi, Kyuwoong Hwang, Sungrack Yun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence has reached a point where it can read entire books, analyze years of chat logs, or digest vast legal contracts in a single go. To do this, these systems rely on a mechanism that allows them to look back at everything they have just processed to understand the current question. However, this ability comes with a steep price. As the amount of text grows, the energy and computer memory required to process it do not just increase; they explode. Imagine trying to read a library by comparing every single word you are reading right now against every single word you have read since you started. The effort required grows so fast that it quickly becomes impossible for standard computers, especially smaller ones found in phones or laptops, to handle more than a few thousand words at a time. This limitation forces current systems to either forget the beginning of a long story or rely on external tools to find specific facts, often missing the bigger picture that connects scattered details.

A team of researchers at Qualcomm AI Research has developed a new approach called MoNe, which allows these artificial intelligence models to handle massive amounts of information without needing to be retrained or without overwhelming their hardware. Instead of forcing the computer to constantly re-read the entire history of a conversation or document every time it answers a question, MoNe acts like a specialized notebook that the model writes in as it reads. The system processes the long text in small, manageable chunks. As it reads each chunk, it updates this internal notebook, distilling the essential information into a compact form. Once the entire text has been read and the notebook is full, the model can answer questions using only the contents of that notebook. Crucially, the model never needs to look back at the original long text again. This design means that the cost of answering a question stays the same, no matter if the original document was a short email or a hundred-thousand-word novel.

The researchers tested this method on a standard set of challenges designed to see if a model could find a specific fact hidden deep within a large pile of text, a task often called finding a needle in a haystack. They also tested its ability to extract frequently mentioned words and to solve problems that required connecting multiple pieces of information spread across the text. When the text length stayed within the model's original training limits, the new method performed nearly perfectly, outperforming standard approaches that try to read everything at once. More impressively, when the researchers pushed the system to handle text lengths far beyond what the model was originally designed for—up to one hundred and twenty-eight thousand tokens—the new method continued to work with high accuracy. In contrast, the standard approach, which tries to read the whole text at once, failed completely as the text grew longer, dropping to near-zero accuracy. A common alternative method that tries to search for relevant snippets of text also struggled when the answer required piecing together many different, scattered facts.

The efficiency gains from this new method are substantial. At the longest text length tested, the researchers found that their approach reduced the amount of computer power needed by about eighty percent compared to the standard method. It also reduced the peak memory required by the same margin. This is a critical improvement because it means that powerful long-context reasoning could eventually run on devices with limited resources, rather than requiring massive, expensive servers. The system achieves this by keeping the memory footprint constant; the size of the internal notebook does not grow as the text gets longer. The researchers demonstrated that this system could be attached to existing, pre-trained models without changing their core structure, adding only a small amount of extra data storage. They also showed that the system could generalize to text lengths thirty-two times longer than what it was trained on, simply by processing the text in fixed-size segments and updating its internal state step by step.

While the current experiments focused on specific retrieval tasks using a particular model size, the results suggest a viable path forward for making artificial intelligence more capable of handling real-world, long-form information. The method does not require building entirely new types of computer brains from scratch, nor does it demand that the model be retrained on massive new datasets. Instead, it introduces a lightweight, modular addition that learns to compress information on the fly. This allows the system to maintain a clear, constant memory of the past while keeping the cost of thinking about the present low. As the demand for artificial intelligence that can reason over hours of conversation or thousands of pages of documentation continues to grow, this approach offers a practical way to keep those systems fast, efficient, and capable of understanding the full context of the world they are asked to navigate.

Technical Summary: MoNe – Modular Neural Memory for Efficient Long Context Inference

1. Problem Statement

Long-context processing is essential for applications such as personalized assistants, document QA, and agentic systems. However, the dominant approach, In-Context Learning (ICL), scales quadratically ( $O(N^2)$ ) in FLOPs and linearly ( $O(N)$ ) in KV-cache memory with context length $N$ . This makes long-context inference prohibitively expensive on resource-constrained hardware (e.g., mobile devices) and causes small models to struggle with reliable information extraction even within their native context windows.

Alternative approaches face significant limitations:

Retrieval-Augmented Generation (RAG): While it reduces attention costs, RAG relies on embedding-based retrieval of independent chunks. This fails when tasks require synthesizing multiple scattered, interdependent facts that are individually unremarkable.
Recurrent/Linear Models (e.g., Mamba, TTT, Titans): These achieve linear complexity but require training entirely new architectures from scratch, preventing seamless integration with existing frozen pretrained Transformers without full retraining.

2. Methodology: MoNe

MoNe (Modular Neural Memory) is a lightweight plugin module designed to attach to any frozen pretrained Transformer, enabling efficient long-context reasoning without modifying backbone weights. The architecture operates in two distinct phases:

2.1. Architecture

MoNe attaches a fast-weight neural memory network to each decoder layer $l$ of the backbone.

Memory Module: Each layer utilizes a SwiGLU MLP with fast-weight parameters ( $W_{in}, W_{gate}, W_{out}$ ). Unlike linear maps, this nonlinear network increases the expressive capacity to encode richer associative structures.
Integration: During inference, query tokens are projected via the frozen backbone's $W_q$ and read from the cached fast-weight state. The resulting memory tokens are projected into key-value space using the frozen $W_k$ and $W_v$ and concatenated with standard attention entries.
Parameter Efficiency: The module adds only a small number of parameters (6.4% overhead) via low-rank adapters (LoRA) to the frozen weights.

2.2. Two-Phase Operation

Phase 1: Test-Time Learning (Preprocessing)

The context $C$ of $N$ tokens is partitioned into fixed-size segments of $T=512$ tokens.
Segments are processed sequentially. For each segment, the fast weights at each layer are updated via layer-localized gradient updates.
Associative Memory Loss: The memory module is optimized to maximize the inner product between the memory output and target values (derived from frozen $W_k$ and $W_v$ ). This "imprints" the association $k \to v$ into the fast weights without requiring prediction error or backpropagation through the entire backbone.
Positional Encoding: Segment-local RoPE is used, assigning local positions $p_{local} = j \mod T$ . This keeps position indices bounded in $[0, T)$ regardless of total context length, enabling generalization to arbitrary lengths without interpolation.

Phase 2: Inference

The query tokens attend only to the memory tokens generated from the updated fast-weight state.
No context tokens are re-read, and no context tokens are attended to directly.
Complexity: This decouples inference cost from context length, achieving $O(N)$ preprocessing and $O(1)$ query cost with peak GPU memory that remains constant regardless of $N$ .

3. Key Contributions

Modular Architecture: A plugin that integrates with any frozen pretrained Transformer, adding only 6.4% parameter overhead without backbone modification.
Test-Time Learning Procedure: A method using layer-localized gradient updates to adapt fast-weight neural memory networks. This enables $O(N)$ preprocessing and $O(1)$ query cost, contrasting with the $O(N^2)$ cost of ICL. It natively supports incremental context extension and multi-query reuse.
Generalization: Demonstrated ability to generalize from training contexts (up to 4K tokens) to evaluation contexts of 128K tokens (a 32× extrapolation) without additional training or positional interpolation.

4. Experimental Results

Experiments were conducted on the RULER benchmark (S-NIAH, MK-NIAH, Frequent Word Extraction) using a Qwen2.5-0.5B-Instruct backbone.

Performance:
- Within Native Window (4K–32K): MoNe achieved near-perfect performance (0.99–1.00 Sub-EM), significantly outperforming ICL (0.41–0.98) and RAG (0.58–0.93).
- Beyond Native Window (48K–128K): While ICL performance collapsed (e.g., dropping to 0.28 on S-NIAH at 128K) and RAG plateaued (struggling with multi-hop facts), MoNe maintained strong accuracy (0.96 on S-NIAH, 0.94 on MK-NIAH, 0.96 on Frequent Word Extraction at 128K).
Efficiency:
- At 128K tokens, MoNe reduced total FLOPs by approximately 81% and peak GPU memory by 80% compared to ICL.
- Peak GPU memory for MoNe remained constant at 1.41 GB regardless of context length, whereas ICL required 7.07 GB at 128K.
Ablation Studies:
- Layer Coverage: Attaching MoNe to all 24 decoder layers yielded the best results. Restricting to the last 8 layers caused performance to collapse beyond 32K.
- Segment Size: A segment size of $T=512$ offered the best trade-off, generalizing effectively to 128K, whereas smaller segments ( $T=128, 256$ ) degraded rapidly past 16K.

5. Significance and Claims

The paper claims that MoNe offers a practical solution for resource-constrained long-context inference by decoupling computational cost from context length. Its significance lies in:

Efficiency: Achieving constant memory footprint and linear preprocessing, making long-context reasoning feasible on edge devices where KV-cache costs are otherwise prohibitive.
Compatibility: Enabling long-context capabilities for existing frozen models without the need for expensive full fine-tuning or architectural redesign.
Robustness: Outperforming both standard ICL (which degrades with length) and RAG (which struggles with distributed reasoning) on tasks requiring the synthesis of information across long contexts.

The authors note that while current validation uses a 0.5B backbone and controlled retrieval tasks, the plug-in design is intended to scale to larger models and naturalistic workloads (e.g., multi-document QA, conversational histories) where the efficiency advantages of constant-memory inference are expected to become even more pronounced.