MemOnce

A Scalable Memory Architecture for Large Language Models

This repository hosts the post-training checkpoint of MemOnce on a frozen Qwen3.5-4B backbone. It contains only the trained memory modules (~71M parameters, 264 MB): the LoRA write adapter, the per-layer KV readout projections and the mem_to_emb index projection. The backbone itself is not included.

• Project page: https://clare-nie.github.io/MemOnce/
• Paper: MemOnce: A Scalable Memory Architecture for Large Language Models (2026)
• Code, training scripts and evaluation: linked from the project page.

Method

MemOnce is a foundation architecture for long-horizon continual language modeling. A Transformer reuses its past by accumulating a per-layer KV cache for every token, so memory grows with both context length and depth; linear/recurrent models avoid that growth but compress history into a bounded state that loses fine-grained detail.

The principle is simple: cache a single compact global memory, and retrieve it only once. Two coupled designs implement it:

1. Memory encoder-decoder. Only a trailing 1/r subset of final-layer hidden states is cached as compact memory blocks, and the selected blocks are decoded back into layer-specific KV caches on demand. The persistent cache therefore loses the depth factor entirely and is compressed along time (L × r depth-time compression).
2. Memory sparse attention. For each new input, the top-K relevant blocks are retrieved once, and that routing decision is shared by every attention layer and every decoding step. The persistent bank may grow and be offloaded to CPU or SSD, while the GPU-resident active state stays constant.

Write / index / read are factorized as

G_t = g_θ(x_{1:t-1})      # write:  encode history once
M_t = i_θ(x_t, G_t)       # index:  describe and retrieve blocks
C_t = u_θ(M_t)            # read:   decode selected blocks into KV

Concretely, in this checkpoint:

• Write — turns longer than 4,096 tokens are split into 4,096-token segments, and each segment is written once as a 64-token memory block by the memory LoRA adapter (r = 128).
• Index — each block is described by a learned 1024-d mem_to_emb projection and selected by a single cosine pass with top-K = 64.
• Read — the selected blocks are decoded into per-layer KV caches; generation runs with the adapter disabled, so history enters strictly as KV caches.

Highlights

• 🧠 Scalable Architecture: Caches and retrieves memory only once, simultaneously achieving accurate recall, linear-time complexity, and a constant active GPU memory budget.
• 🧱 Foundation Performance: Consistently outperforms Transformer, linear GDN, and YOCO across model sizes, benefiting from depth-time cache recurrence and compression.
• 🎯 Native Memory: The first native memory model to surpass full-context baselines, beating prior methods like M+ and Metis without catastrophic forgetting.
• ⚡ Inference Scalability: At 1M context, slashes active GPU cache by 1,313×, accelerates prefill by 6.9×, and boosts decoding throughput by 3.5×.

Detailed results and figures are on the project page.

Usage

Qwen3.5 post-training requires transformers >= 5.0 (see requirements/posttraining.txt in the code repository). Download the frozen backbone and this checkpoint, then run the inference script from the code repository:

pip install -U huggingface_hub
hf download Qwen/Qwen3.5-4B --local-dir Qwen3.5-4B
hf download ClareNie/MemOnce --local-dir model

cd eval/post-training
export BASE=../../Qwen3.5-4B
export CKPT=../../model

# 1. Self-test (verifies every trained tensor is strictly consumed)
python inference.py --base-model "$BASE" --checkpoint "$CKPT" --self-test

# 2. Built-in multi-hop demo
python inference.py --base-model "$BASE" --checkpoint "$CKPT" --TEXT

--self-test verifies that all trained tensors are consumed by the loader before any generation. For your own data, pass --input my_data.jsonl --output generations.jsonl; each line is a messages list whose preceding turns are written into memory and whose last user turn is the query. Other useful flags: --max-new-tokens, --retrieve-blocks, --attn-implementation [sdpa|flash_attention_2|eager].

Files

Citation

@article{nie2026memonce,
  title   = {MemOnce: A Scalable Memory Architecture for Large Language Models},
  author  = {Nie, Chang and Fu, Chaoyou and Liu, Songlin and Hu, Yiqing and Shan, Caifeng},
  journal = {arXiv preprint},
  year    = {2026}
}