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.
Figures & tables
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.
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.
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 32 ) yields modest degradation beyond 32K, while using only the last 8 layers (last 31 ) collapses beyond 32K. ( Right ) Segment size. T=512 (default) generalizes from 4K training contexts to 128K evaluation contexts, whereas T=128 and T=256 degrade rapidly past 16K, yet offer only marginal efficiency gains over the default.