English | 中文

• Introduction
• Latest News
• Installation
• Quick Start
• Reproduce Results
• Training
• Citation

MemSifter is an LLM memory retrieval offloading system based on outcome-driven proxy reasoning. Given a large pool of personal conversation sessions (the "haystack"), MemSifter efficiently identifies the sessions most relevant to a user query and passes them as context to a downstream chat LLM—without burdening the LLM itself with retrieval.

The system follows a three-stage pipeline:

Session Embedding  →  Session Ranking (MemSifter)  →  Chat LLM
   (bge-m3)           (generative reranker)          (any LLM)

1. Session Embedding — a dense embedding model (bge-m3) performs a coarse similarity pre-filter across all sessions.
2. Session Ranking — MemSifter, a lightweight generative model trained with DAPO reinforcement learning, performs fine-grained reranking of the pre-filtered candidates.
3. Chat LLM — the top-ranked sessions are assembled into a context window and passed to any OpenAI-compatible chat model to produce the final answer.

• [03/03/2026] Our paper is available on arXiv.
• [20/02/2026] Code, models, and data are released.

Requirements: Python 3.8+, two CUDA-capable GPUs (for Quick Start single-sample inference).

git clone https://github.com/plageon/MemSifter.git
cd MemSifter
pip install -r requirements.txt

Download the required models into a local models/ directory:

# Example using huggingface-cli
huggingface-cli download BAAI/bge-m3 --local-dir models/bge-m3
huggingface-cli download zstanjj/MemSifter-4B-Thinking \
    --local-dir models/zstanjj/MemSifter-4B-Thinking

For the Quick Start single-sample inference (no Ray), install only the required packages:

pip install torch sentence-transformers vllm openai pyyaml loguru numpy pandas

The toolkit in memsifter/toolkit.py provides three classes for single-sample inference with no Ray dependency. It assumes you have two GPUs: the embedding model runs on cuda:0 and the MemSifter ranker runs on cuda:1.

import json
from memsifter.toolkit import SessionEmbedder, SessionRanker, LLMChat

# Load one sample and unpack all fields
with open("data/test_memory.json") as f:
    entry = json.load(f)[0]

question             = entry["question"]
haystack_sessions    = entry["haystack_sessions"]
haystack_dates       = entry["haystack_dates"]
haystack_session_ids = entry["haystack_session_ids"]
answer_session_ids   = entry["answer_session_ids"]

# Initialise models (loaded once, reusable)
embedder = SessionEmbedder(model_path="models/bge-m3", device="cuda:0")
ranker   = SessionRanker(
    model_path="models/zstanjj/MemSifter-4B-Thinking",
    device="cuda:1",
)
chat = LLMChat(api_key="YOUR_KEY", base_url="YOUR_BASE_URL", model_name="YOUR_MODEL")

# Stage 1 — embedding pre-filter
top_sessions = embedder.get_top_sessions(
    question=question, sessions=haystack_sessions, dates=haystack_dates, top_k=20
)

# Stage 2 — generative reranking
ranked_sessions = ranker.rerank(
    question=question, pre_ranked_sessions=top_sessions, top_k=5
)

# Stage 3 — LLM answer
predicted_answer = chat.answer(question=question, ranked_sessions=ranked_sessions)

print("Question:", question)
print("Answer:  ", predicted_answer)

Each entry in data/test_memory.json has the following fields:

SessionEmbedder(model_path, device="cuda:0")

• get_top_sessions(question, sessions, dates=None, top_k=20) → List[(idx, session_turns, date, score)]

SessionRanker(model_path, device="cuda:1")

• rerank(question, pre_ranked_sessions, top_k=5) → List[(idx, session_turns, date)]

LLMChat(api_key, base_url, model_name)

• answer(question, ranked_sessions) → str

This section covers batch inference across all benchmark datasets using the released MemSifter checkpoint. The batch pipeline uses Ray for distributed multi-GPU inference.

https://huggingface.co/datasets/zstanjj/MemSifter-Datasets

Start a Ray cluster before running the scripts:

ray start --head

Set the required environment variables:

export API_KEY="YOUR_LLM_API_KEY"
export BASE_URL="YOUR_LLM_BASE_URL"
export CUDA_VISIBLE_DEVICES=0,1

Computes bge-m3 embeddings for all sessions in the benchmark datasets and stores similarity scores.

cd scripts/infer
./session_embedding.sh

Key variables (edit inside the script or export before running):

Runs the MemSifter generative ranker over the embedding-pre-filtered candidates.

cd scripts/infer
./session_ranking.sh

Key variables:

Passes the ranked sessions to a chat LLM and collects the generated answers.

cd scripts/infer
./chat_infer.sh

Key variables:

Stop Ray after inference is complete:

ray stop

This section describes how to train your own MemSifter ranker on custom data.

Run the embedding and ranking steps on your own datasets first (see Reproduce Results), then prepare the DAPO training data:

cd scripts/train
./prepare_rl_data.sh

Key variables:

Outputs:

• {OUTPUT_DIR}/rl_train_data/{VERSION}/train_*.parquet
• {OUTPUT_DIR}/rl_train_data/{VERSION}/test.parquet

cd scripts/train
./qwen3_4b_task_reward.sh

Key variables:

The training uses task reward mode (Marginal Utility Reward + Rank-Sensitive Reward) via the DAPO algorithm.

Convert VERL checkpoints to HuggingFace format:

cd scripts/train
./collect_verl_ckpt.sh

Merge multiple checkpoint steps by weight averaging (optional but recommended):

export CKPT_STEPS="20 30 40"
export MODEL_NAME="MemSifter-Qwen3-4B-Task-Reward"
./merge_ckpts.sh

The merged model is saved to {MODEL_DIR}/{MODEL_NAME}-merged.

If you use MemSifter in your research, please cite:

@misc{memsifter,
      title={MemSifter: Offloading LLM Memory Retrieval via Outcome-Driven Proxy Reasoning}, 
      author={Jiejun Tan and Zhicheng Dou and Liancheng Zhang and Yuyang Hu and Yiruo Cheng and Ji-Rong Wen},
      year={2026},
      eprint={2603.03379},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2603.03379}, 
}