MemOnce / README.md

Use README highlights, drop checksum blocks

af6e1e2 verified about 10 hours ago

5.21 kB

---
license: apache-2.0
base_model: Qwen/Qwen3.5-4B
tags:
  - memory
  - long-context
  - retrieval
  - qwen3.5
---
# 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](https://clare-nie.github.io/MemOnce/).

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

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

| File | Size | Contents |
|---|---|---|
| `model.safetensors` | 264 MB | 100 trained tensors: LoRA write adapter (`r = 128`), per-layer KV readout projections, `mem_to_emb` |
| `config.json` | 554 B | Memory hyper-parameters: 64 memory tokens, top-`K = 64`, 1024-d descriptors, 4,096-token target length |
| `index.json` | 25 KB | Per-tensor manifest for `model.safetensors` (names, dtypes, shapes) |

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

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

	# 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](https://clare-nie.github.io/MemOnce/).

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

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

	\| File \| Size \| Contents \|
	\|---\|---\|---\|
	\| `model.safetensors` \| 264 MB \| 100 trained tensors: LoRA write adapter (`r = 128`), per-layer KV readout projections, `mem_to_emb` \|
	\| `config.json` \| 554 B \| Memory hyper-parameters: 64 memory tokens, top-`K = 64`, 1024-d descriptors, 4,096-token target length \|
	\| `index.json` \| 25 KB \| Per-tensor manifest for `model.safetensors` (names, dtypes, shapes) \|

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