Thought-Retriever: Don’t Just Retrieve Raw Data, Retrieve Thoughts for Memory-Augmented Agentic Systems

<p align="center">
<a href="https://ulab-uiuc.github.io/Thought-Retriever/">🌐 Project Page</a> |
<a href="https://arxiv.org/pdf/2604.12231">📜 arXiv</a> |
<a href="https://x.com/TmlrPub/status/2044326160935395600">📮 Twitter Post</a>
<p>

[2026.04.13] 🌟 Thought-Retriever is accepted for TMLR 2026.

📌Preliminary

Environment Setup

# create a new environment
conda create -n thought-retriever python=3.10
conda activate thought-retriever

# install pytorch. Modify the command to match your CUDA version
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 \
  --extra-index-url https://download.pytorch.org/whl/cu113

# install related libraries
pip install -r requirements.txt

Dataset Preparation

QMSum
WCEP
Booksum
GovReport
SQuALITY

Save the downloaded files in the ./data/[DATASET_NAME] folder.

Generate Pre-queries

You can generate pre-queries using the following code:

python pre_query_generator.py \
    --root DATA_PATH \
    --api_base YOUR_API_BASE \
    --api_key YOUR_API_KEY \
    --dataset_type related_multi \
    --chunk_size 500 \
    --feat_cross \
    --with_doct5query \
    --device 0

⭐Experiments

Run thought-retriever and Metric Evaluation

You can follow this example:

python run_experiment.py \
  --root DATA_PATH \
  --api_base YOUR_API_BASE \
  --api_key YOUR_API_KEY \
  --llm_model YOUR_MODEL \
  --cuda 0 \
  --chunk_num 8 \
  --recall_coe 5 \
  --sim_thre 0.85

Evaluation using LLM-as-a-Judge

Evaluate the answers generated by thought-retriever through this:

python llm_evaluation.py \
  --api_base YOUR_API_BASE \
  --api_key YOUR_API_KEY \
  --llm_model YOUR_MODEL \
  --seed 42 \
  --tau 0

Citation

@article{feng2026thought,
  title={Thought-Retriever: Don't Just Retrieve Raw Data, Retrieve Thoughts for Memory-Augmented Agentic Systems},
  author={Feng, Tao and Han, Pengrui and Lin, Guanyu and Liu, Ge and You, Jiaxuan},
  journal={arXiv preprint arXiv:2604.12231},
  year={2026}
}

README