Skip to content

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

RequirementValue
GPUOne H100-class GPU (80 GB)
Time~1 GPU-day for stage 1; shardable across N processes
API keyOPENAI_API_KEY (stage 5 only)
Python3.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 reproducing-lightmem python=3.11 -y
conda activate reproducing-lightmem
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.

ModelPortStages
Qwen/Qwen3-30B-A3B-Instruct-250780001: summarise; 4: answer
Qwen/Qwen3-Embedding-0.6B80013: retrieve
Qwen/Qwen3-Embedding-4B80023: retrieve
Qwen/Qwen3-Embedding-8B80033: 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
#StageOutput
1Ingestoutputs/_memory/
2Corpus + qrelsoutputs/_qrels/
3Retrieve + evaluateoutputs/_retrieval/
4Answeroutputs/_answers/
5Judgemetrics.<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 --mem-base-url http://localhost:8000/v1

~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

RetrieverSpec entry
BM25bm25
SPLADE-v3sparse:naver/splade-v3
PromptReps-sparsesparse:promptreps:<Llama-3-8B>

Dense

RetrieverSpec entryServer
all-MiniLM-L6-v2dense:sentence-transformers/all-MiniLM-L6-v2
Qwen3-Embedding-0.6Bdense:Qwen/Qwen3-Embedding-0.6B:8001
Qwen3-Embedding-4Bdense:Qwen/Qwen3-Embedding-4B:8002
Qwen3-Embedding-8Bdense:Qwen/Qwen3-Embedding-8B:8003
PromptReps-densedense:promptreps:<Llama-3-8B>

Hybrid

RetrieverSpec entryServer
PromptReps-hybridhybrid:promptreps:<Llama-3-8B>
MiniLM + BM25fusion:sentence-transformers/all-MiniLM-L6-v2
Qwen3-0.6B + BM25fusion: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.

RunCorpus / configurationSpecStagesAnswer depthsUsed by
ALightMem (r=0.8, th=1024)default1 → 53,5,6,10,20RQ1, RQ2, RQ3
BRaw user turnsuser turns2 → 5 + full-context3,5,10RQ1, RQ3
CLightMem (th=768)r=0.4 / r=0.61 → 510RQ1

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.

QuestionRunsResult slice
RQ1: reproductionA + B + CTop-10 + baselines + costs
RQ2: retriever effectA11 retrievers at top-10 + oracle
RQ3: construction valueA + BDepth/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

MethodACC (%)Constr. tokens (k)Constr. calls
Full-context54.8
Naive RAG60.8
LightMem (r=0.4, th=768)62.3144.16192.56
LightMem (r=0.6, th=768)65.1135.43172.90
LightMem (r=0.8, th=1024)67.3146.42177.80

Our reproduced results

MethodACC (%)Constr. tokens (k)Constr. callsAnswer tokens (k)
Full-context56.50018.84
Naive RAG66.9001.03
LightMem (r=0.4, th=768)58.372.6864.690.72
LightMem (r=0.6, th=768)67.8106.90104.520.67
LightMem (r=0.8, th=1024)70.7119.88117.620.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.

FamilyRetrieverRecall@10ACC (%)
SparseBM250.39057.2
SparseSPLADE-v30.54271.2
SparsePromptReps-sparse0.45663.5
Denseall-MiniLM-L6-v2 (default)0.53070.7
DenseQwen3-Embedding-0.6B0.56272.5
DenseQwen3-Embedding-4B0.58775.5
DenseQwen3-Embedding-8B0.58372.5
DensePromptReps-dense0.51469.1
HybridPromptReps-hybrid0.51268.5
HybridMiniLM + BM250.50267.3
HybridQwen3-0.6B + BM250.50668.0
Oracle1.00075.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

Retrievertop-3top-5top-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.90.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

0 stars

Watchers

0 watching

Forks

Contributors

Languages