Code to reproduce our paper: a reproduction of LightMem on LongMemEval-S, compared against Naive RAG — retrieval directly over raw user turns, with no memory-construction stage.
LightMem transforms raw dialogue into constructed memory entries before retrieval; Naive RAG retrieves the raw turns directly. We compare the two representations across retrievers, retrieval depths, answering-token budgets, and oracle conditions.
| Requirement | Value |
|---|---|
| GPU | One H100-class GPU (80 GB) |
| Time | ~1 GPU-day for stage 1; shardable across N processes |
| API key | OPENAI_API_KEY (stage 5 only) |
| Python | 3.11 |
The model endpoints are needed one stage at a time, so one GPU is sufficient. Everything except
the gpt-5.5 judge is self-hosted; see Model endpoints.
Conda provides the interpreter; every package comes from pip, so the pins in requirements.txt are what you actually get.
conda create -n lightmem python=3.11 -y
conda activate lightmem
pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txtLightMem needs no install. It is vendored at baselines/lightmem/, and
bootstrap.py puts baselines/ on sys.path before any stage runs —
so lightmem.* (upstream code) and baselines.lightmem.* (this repo's adapters) both import
with nothing pip-installed. Ignore baselines/lightmem/setup.py; installing it has no effect,
since the vendored copy precedes site-packages on the path.
Get LongMemEval-S from xiaowu0162/LongMemEval and put it at exactly this path — benchmarks/adapters/init.py registers it relative to the repo root, with no environment override:
data/LongMemEval/longmemeval_s_cleaned.json
Put your judge key in a .env at the repo root (OPENAI_API_KEY=...). It is used only for
the LLM judge; the generator, memory summariser and embedding models are all self-hosted.
Pre-compression needs tiktoken's BPE files, downloaded on first use. On a cluster whose compute nodes have no internet, warm the cache on a login node:
python -c "import tiktoken; tiktoken.encoding_for_model('gpt-3.5-turbo')"python -c "import scripts; from baselines.lightmem.longmemeval import LightMemLongMemEval; print('ok')"That exercises the bootstrap, the vendored LightMem and the adapter in one import. If it prints
ok, everything but the model endpoints is ready. Run the stages from the repo root — the .sh
drivers resolve outputs/ and logs/ relative to the working directory.
The pipeline expects these OpenAI-compatible endpoints to be already running. We serve them with vLLM (v0.22.1) on H100s; any compatible server works.
| Model | Port | Stages |
|---|---|---|
Qwen/Qwen3-30B-A3B-Instruct-2507 |
8000 | 1: summarise; 4: answer |
Qwen/Qwen3-Embedding-0.6B |
8001 | 3: retrieve |
Qwen/Qwen3-Embedding-4B |
8002 | 3: retrieve |
Qwen/Qwen3-Embedding-8B |
8003 | 3: retrieve |
Each endpoint uses http://localhost:<port>/v1 as its base URL.
No stage needs more than one of them. Stage 1 and stage 4 use only the generator, stage 3
only the embedders its spec's dense:/fusion: entries name, stage 2 no endpoint at all, and
stage 5 only the OpenAI API. On a single GPU, run the stages in order and restart the server
between them.
The generator is what we served with, sized for the full-context baseline:
vllm serve Qwen/Qwen3-30B-A3B-Instruct-2507 \
--host 0.0.0.0 --port 8000 \
--max-model-len 65536 \
--gpu-memory-utilization 0.8 \
--kv-cache-dtype fp8Each embedding server is the same command with its own model and port.
Ports are what the specs' dense:MODEL@http://localhost:PORT/v1 entries point at — change
either side to match your setup. The generator URL is overridable with GEN_BASE_URL, and it
needs a context window of at least 64k tokens to fit the full-context baseline.
The other retrievers load in-process and need no server — see the table in Stage 3 for which is which.
Five stages, each reading the previous one's output tree. Every experiment's coordinates live
in one YAML spec under experiments/ — Python stages read it via --exp, the
.sh drivers via SPEC=. Set it once:
SPEC=experiments/longmemeval_lightmem_has_answer.yaml| # | Stage | Output |
|---|---|---|
| 1 | Ingest | outputs/_memory/ |
| 2 | Corpus + qrels | outputs/_qrels/ |
| 3 | Retrieve + evaluate | outputs/_retrieval/ |
| 4 | Answer | outputs/_answers/ |
| 5 | Judge | metrics.<judge>.json beside each answer set |
Builds the LightMem memory store. LightMem specs only — Naive RAG, Full-context and Oracle have no construction stage and skip straight to stage 2.
python -m scripts.pipeline.ingest --exp $SPEC~1 GPU-day. Resumable and shardable — an already-saved sample is skipped, and N processes can cover disjoint subsets writing to the same memory dir:
python -m scripts.pipeline.ingest --exp $SPEC --num-shards 4 --shard-id 0 # ... and 1,2,3Construction cost lands in outputs/_memory/**/ingest_usage.json.
Writes the retrieval corpus and its relevance judgements. Which builder depends on the representation:
python -m scripts.pipeline.qrels_has_answer --exp $SPEC # LightMem memory entries
python -m scripts.pipeline.qrels_turns --dataset longmemeval_s # raw user turns (Naive RAG etc.)A LightMem entry counts as relevant if it was constructed or updated from a turn
LongMemEval-S marks has_answer.
Runs every retriever in the spec's retrievers: list, then scores each run against the qrels.
SPEC=$SPEC bash experiments/run_all_retrieval.shKnobs: K_VALUES=3,5,10,20 (recall/NDCG cutoffs), PARALLEL=1 (dense endpoints concurrently),
EVAL_ONLY=1 (re-score an existing run.tsv, no servers needed), EVAL=0 (retrieve only).
The paper's 11 retrievers, grouped by family below. To keep the tables narrow, <Llama-3-8B>
stands for meta-llama/Meta-Llama-3-8B-Instruct, and endpoint URLs are represented by their
ports. The experiment specs contain the full, directly executable entries.
Sparse
| Retriever | Spec entry |
|---|---|
| BM25 | bm25 |
| SPLADE-v3 | sparse:naver/splade-v3 |
| PromptReps-sparse | sparse:promptreps:<Llama-3-8B> |
Dense
| Retriever | Spec entry | Server |
|---|---|---|
| all-MiniLM-L6-v2 | dense:sentence-transformers/all-MiniLM-L6-v2 |
— |
| Qwen3-Embedding-0.6B | dense:Qwen/Qwen3-Embedding-0.6B |
:8001 |
| Qwen3-Embedding-4B | dense:Qwen/Qwen3-Embedding-4B |
:8002 |
| Qwen3-Embedding-8B | dense:Qwen/Qwen3-Embedding-8B |
:8003 |
| PromptReps-dense | dense:promptreps:<Llama-3-8B> |
— |
Hybrid
| Retriever | Spec entry | Server |
|---|---|---|
| PromptReps-hybrid | hybrid:promptreps:<Llama-3-8B> |
— |
| MiniLM + BM25 | fusion:sentence-transformers/all-MiniLM-L6-v2 |
— |
| Qwen3-0.6B + BM25 | fusion:Qwen/Qwen3-Embedding-0.6B |
:8001 |
fusion: is BM25 plus the named dense model, 0.5/0.5 min-max score fusion. The three
PromptReps rows come from one encode pass — the hybrid: entry writes all three run.tsv
files, so the dense:/sparse: entries exist only so stage 4 also answers for them.
Gated models. SPLADE-v3 and the three PromptReps rows are the only retrievers that need
Hugging Face access: accept the terms for naver/splade-v3 and
meta-llama/Meta-Llama-3-8B-Instruct on the Hub, then log in on the run machine
(hf auth login, or set HF_TOKEN).
Generates an answer per question from each retriever's top-k, plus the oracle upper bound.
SPEC=$SPEC TOP_K_VALUES=10 bash experiments/run_all_answers.sh
SPEC=$SPEC FULL_ONLY=1 bash experiments/run_all_answers.sh # Full-context baseline onlyKnobs: TOP_K_VALUES (comma-separated depths to sweep), INCLUDE_ORACLE=0 (skip the oracle),
FULL_ONLY=1 (prompt with the entire corpus and nothing else — needs the generator's large
context window, see Model endpoints). Generation is pinned in the driver
(greedy, seed 42, 2000 output tokens).
Discovers every predictions.jsonl under this experiment's answers root and scores it.
SPEC=$SPEC bash experiments/run_all_judge.shKnobs: JUDGE_MODEL=gpt-4o-mini (the original study's judge — metrics files coexist per judge,
so re-scoring the same predictions needs no re-generation), BATCH=1 (OpenAI Batch API —
cheaper, no rate limit), FORCE=1 (re-judge sets this judge already scored).
Three runs cover the whole paper. They are organised by the artifact they build rather than by RQ, because the artifacts are shared: one LightMem memory store serves RQ1, RQ2 and RQ3, so nothing here is run twice.
| Run | Corpus / configuration | Spec | Stages | Answer depths | Used by |
|---|---|---|---|---|---|
| A | LightMem (r=0.8, th=1024) |
default | 1 → 5 | 3,5,6,10,20 |
RQ1, RQ2, RQ3 |
| B | Raw user turns | user turns | 2 → 5 + full-context | 3,5,10 |
RQ1, RQ3 |
| C | LightMem (th=768) |
r=0.4 / r=0.6 | 1 → 5 | 10 |
RQ1 |
Run A's spec sweeps all 11 retrievers (see Stage 3) over one fixed
memory store — that sweep is RQ2. Its TOP_K_VALUES covers RQ3's
matched-depth and matched-token-budget grids at once: top-6/10/20 on LightMem pair with
top-3/5/10 on raw turns (run B). The oracle condition is produced automatically in stage 4.
Run C is reported at top-10 with all-MiniLM-L6-v2 only, so a single depth is enough; its
construction cost comes from each store's ingest_usage.json and needs no answering at all.
mem_variant suffixes every stage's output dir, so the three configs never collide — the two
specs can be looped over back to back.
| Question | Runs | Result slice |
|---|---|---|
| RQ1: reproduction | A + B + C | Top-10 + baselines + costs |
| RQ2: retriever effect | A | 11 retrievers at top-10 + oracle |
| RQ3: construction value | A + B | Depth/token matches + oracles |
Answer scores are in metrics.*.json; construction costs are in
outputs/_memory/**/ingest_usage.json. For RQ3, compare LightMem top-6/10/20 with Naive RAG
top-3/5/10.
- Question set. Following the original LightMem study we exclude the 56
single-session-assistantquestions (LightMem constructs memories from user turns only), leaving 444 questions — 422 of which carry gold evidence and are scored for recall. The drivers apply this viaEXCLUDE_QUESTION_TYPES; keep it in sync betweenrun_all_answers.shandrun_all_judge.sh. - Not generated here. The significance tests (paired t-test on Recall@10, McNemar's
exact test on accuracy, Bonferroni-corrected) and the paper's tables and figures were
computed from the
outputs/trees separately.
The paper reports results judged by gpt-5.5, whereas the original LightMem evaluation used
gpt-4o-mini. For a like-for-like comparison with the original judge, we re-graded the same
generated answers with gpt-4o-mini; retrieval scores, generated answers, and token/call counts
are unchanged. All accuracies below use the same 444-question evaluation set described in
Notes.
Construction and answering token counts are in thousands. Construction costs include both online summarisation and offline OP-update; the original paper did not report answering tokens.
Original LightMem results
| Method | ACC (%) | Constr. tokens (k) | Constr. calls |
|---|---|---|---|
| Full-context | 54.8 | — | — |
| Naive RAG | 60.8 | — | — |
LightMem (r=0.4, th=768) |
62.3 | 144.16 | 192.56 |
LightMem (r=0.6, th=768) |
65.1 | 135.43 | 172.90 |
LightMem (r=0.8, th=1024) |
67.3 | 146.42 | 177.80 |
Our reproduced results
| Method | ACC (%) | Constr. tokens (k) | Constr. calls | Answer tokens (k) |
|---|---|---|---|---|
| Full-context | 56.5 | 0 | 0 | 18.84 |
| Naive RAG | 66.9 | 0 | 0 | 1.03 |
LightMem (r=0.4, th=768) |
58.3 | 72.68 | 64.69 | 0.72 |
LightMem (r=0.6, th=768) |
67.8 | 106.90 | 104.52 | 0.67 |
LightMem (r=0.8, th=1024) |
70.7 | 119.88 | 117.62 | 0.66 |
Overall top-10 results for LightMem. Recall@10 is judge-independent; only answer accuracy is
re-graded. The default LightMem retriever is all-MiniLM-L6-v2.
| Family | Retriever | Recall@10 | ACC (%) |
|---|---|---|---|
| Sparse | BM25 | 0.390 | 57.2 |
| Sparse | SPLADE-v3 | 0.542 | 71.2 |
| Sparse | PromptReps-sparse | 0.456 | 63.5 |
| Dense | all-MiniLM-L6-v2 (default) | 0.530 | 70.7 |
| Dense | Qwen3-Embedding-0.6B | 0.562 | 72.5 |
| Dense | Qwen3-Embedding-4B | 0.587 | 75.5 |
| Dense | Qwen3-Embedding-8B | 0.583 | 72.5 |
| Dense | PromptReps-dense | 0.514 | 69.1 |
| Hybrid | PromptReps-hybrid | 0.512 | 68.5 |
| Hybrid | MiniLM + BM25 | 0.502 | 67.3 |
| Hybrid | Qwen3-0.6B + BM25 | 0.506 | 68.0 |
| — | Oracle | 1.000 | 75.5 |
Each cell is the answer-accuracy difference in percentage points,
LightMem - Naive RAG; positive values favour LightMem.
Matched retrieval depth
| Retriever | top-3 | top-5 | top-10 |
|---|---|---|---|
| BM25 | -7.0 | -10.6 | -10.4 |
| SPLADE-v3 | -9.5 | -5.6 | -6.5 |
| PromptReps-sparse | -8.3 | -8.6 | -6.1 |
| all-MiniLM-L6-v2 (default) | +9.2 | +5.0 | +3.8 |
| Qwen3-Embedding-0.6B | -4.7 | -6.5 | -6.1 |
| Qwen3-Embedding-4B | +1.4 | -1.1 | +0.7 |
| Qwen3-Embedding-8B | -1.6 | -2.3 | -4.5 |
| PromptReps-dense | +3.4 | +0.7 | +2.0 |
| PromptReps-hybrid | -2.5 | -6.3 | -6.8 |
| MiniLM + BM25 | -2.0 | -9.0 | -5.4 |
| Qwen3-0.6B + BM25 | -7.9 | -11.9 | -7.9 |
Matched answering-token budget
| Retriever | ~330 tokens | ~500 tokens | ~935 tokens |
|---|---|---|---|
| BM25 | -2.5 | -5.4 | -2.0 |
| SPLADE-v3 | +1.6 | +0.7 | -3.2 |
| PromptReps-sparse | -1.4 | -0.5 | -1.4 |
| all-MiniLM-L6-v2 (default) | +18.0 | +11.0 | +6.8 |
| Qwen3-Embedding-0.6B | +2.5 | -0.9 | -2.7 |
| Qwen3-Embedding-4B | +9.2 | +5.6 | +2.7 |
| Qwen3-Embedding-8B | +5.2 | +1.4 | -3.6 |
| PromptReps-dense | +13.1 | +12.8 | +4.3 |
| PromptReps-hybrid | +5.9 | 0.0 | -4.5 |
| MiniLM + BM25 | +6.1 | +0.2 | -3.6 |
| Qwen3-0.6B + BM25 | -3.8 | -2.9 | -7.2 |
Under oracle retrieval, Naive RAG reaches 87.8% accuracy and LightMem reaches 75.5%. LightMem's fixed construction cost is approximately 119,884 tokens and 117 LLM calls per ingested sample, compared with zero construction cost for Naive RAG.
