README.md

metadata

license: apache-2.0
base_model: Qwen/Qwen3.5-4B
tags:
  - memory
  - long-context
  - retrieval
  - qwen3.5

MemOnce

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 (~0.4B parameters); 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)

Method

Long-horizon interaction forces a language model to keep reusing its past. A Transformer does this by accumulating a per-layer KV cache for every token, so memory grows with both context length and depth. Linear-time models avoid that growth but compress history into a bounded state that loses fine-grained detail.

MemOnce separates what is remembered from how it is attended to, with two coupled designs:

1. Memory encoder-decoder. Instead of caching KV for the whole history, MemOnce stores only a trailing 1/r subset of final-layer hidden states as compact memory blocks, and decodes the selected blocks back into layer-specific KV caches on demand. The persistent cache loses the depth factor entirely and is compressed along the sequence dimension.
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.

The three operations 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 — each 512-token segment (or dialogue turn) is compressed once into a 64-token memory block by the memory LoRA adapter (r = 128).
• Index — each block is described by a learned 1024-d projection (mem_to_emb) and retrieved by a single cosine pass, top-K = 64.
• Read — the selected blocks are decoded into per-layer KV caches; generation runs with the adapter disabled.

Key conclusions

• The active state stops growing with context. What stays on the GPU is a fixed memory budget, not a KV cache proportional to the number of tokens seen.
• Depth no longer multiplies the cache. Storing only trailing final-layer states means the persistent footprint is independent of the number of layers.
• Retrieval happens once, not per layer. A single routing decision is reused across all layers and decoding steps instead of being recomputed for each of them.
• General capability is preserved. The backbone stays frozen during generation and history enters strictly as decoded KV caches, so long-term memory is added without catastrophic forgetting.
• One architecture, two settings. MemOnce works as a foundation model trained from scratch, and as a post-training adaptation of a frozen Qwen3.5 backbone — this checkpoint is the latter.

Usage

You also need the frozen backbone:

hf download Qwen/Qwen3.5-4B --local-dir Qwen3.5-4B
hf download ClareNie/MemOnce --local-dir memonce-ckpt

cd test/post-training
python inference.py --base-model ../Qwen3.5-4B --checkpoint ../memonce-ckpt --self-test
python inference.py --base-model ../Qwen3.5-4B --checkpoint ../memonce-ckpt --TEXT

--self-test verifies that every trained tensor in the checkpoint is consumed by the loader before any generation happens. 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.

Files

File	Size	Contents
model.safetensors	264 MB	Memory LoRA adapter (r = 128), per-layer KV readout projections, mem_to_emb
config.json	665 B	Memory hyper-parameters: 64 memory tokens, top-K = 64, 1024-d descriptors

Checksums:

f5588a7527bc748f686e26f1a76854ca88bd0061f9733505b8cd23703ebfc927  model.safetensors
8a363b36ed02450a93d22efd16b42f9c15d4b304b40e3be84ec73f5c479cb7d8  config.json

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}
}