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

```bash
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

```bibtex
@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}
}
```