Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Reproducing LightMem: Naive RAG Is Just as Good for Memory Management

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.

Workflow: LightMem builds memory entries with an LLM while Naive RAG uses raw conversations directly; both are then retrieved over with matched retrievers, depths and token budgets.

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.

Setup

What you need

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.

1. Environment

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.txt

LightMem 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.

2. Dataset

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

3. Judge key

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.

4. Tokeniser cache

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')"

5. Verify

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.

Model endpoints

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 fp8

Each 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.

Pipeline

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

Stage 1 — Ingest

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,3

Construction cost lands in outputs/_memory/**/ingest_usage.json.

Stage 2 — Corpus + qrels

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.

Stage 3 — Retrieve + eval

Runs every retriever in the spec's retrievers: list, then scores each run against the qrels.

SPEC=$SPEC bash experiments/run_all_retrieval.sh

Knobs: 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).

Stage 4 — Answer

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 only

Knobs: 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).

Stage 5 — Judge

Discovers every predictions.jsonl under this experiment's answers root and scores it.

SPEC=$SPEC bash experiments/run_all_judge.sh

Knobs: 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).

Reproduce

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.

Notes

  • Question set. Following the original LightMem study we exclude the 56 single-session-assistant questions (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 via EXCLUDE_QUESTION_TYPES; keep it in sync between run_all_answers.sh and run_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.

Appendix: Results judged by GPT-4o-mini

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.

RQ1 — Reproduction

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

RQ2 — Retriever effect

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

RQ3 — LightMem versus Naive RAG

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.

About

Code for the paper “Reproducing LightMem: Naive RAG Is Just as Good for Memory Management”.

Resources

Stars

Watchers

Forks

Contributors

Languages