MemTrapBench

Code, data, and evaluation scripts for our paper “MemTrapBench: Benchmarking Cognitive Traps in LLM Memory Use.”

Hugging Face Paper

MemTrapBench is a benchmark for testing how LLMs use memory. It checks both sides of memory use: when memory helps the model answer better, and when old conversation history pushes the model toward stale, unsafe, or irrelevant answers. This repo includes the benchmark data, scripts for running evaluations, and the AdaptiveMem prompt policy (AdaptiveMem/skill.txt) for deciding when retrieved memory should or should not be used.

Highlights

• Two cognitive-trap categories: Reasoning Fixation and Belief Distortion.
• Four Memory Trap scenarios: Task Boundary, Cognitive Bias, Trauma, and Safety.
• 1,050 benchmark instances: each sample has a conversation history plus a final question that can be answered on its own.
• Three evaluation modes: test the model with no memory, with the full conversation, or with retrieved memory.
• AdaptiveMem intervention: a prompt-side policy that tells the model when to trust retrieved memory and when to ignore it.

Table of contents

• Cognitive-trap taxonomy
• Data construction
• Repository layout
• Benchmark JSON shape
• Evaluation runners
• Summarizing judge scores
• Using the AdaptiveMem skill
• Acknowledgements
• Citation

Cognitive-trap taxonomy

Each benchmark sample hides one memory trap in the earlier conversation. The final query is new and should be answered based on what it asks now. The test checks whether the model answers the new query correctly, or whether it gets pulled back to the misleading old context.

MemTrapBench uses two broad categories and four scenarios. Reasoning Fixation covers Task Boundary, Cognitive Bias, and Trauma. Belief Distortion is tested through Safety.

Memory Trap scenario	Cognitive-trap category	What the history does	Correct behavior on the final query
Task Boundary	Reasoning Fixation	Makes the earlier task or format feel important, then asks a different final question.	Follow the latest query, not the old task or format.
Cognitive Bias	Reasoning Fixation	Repeats one domain, frame, or solution path until the model may stick with it automatically.	Reconsider the final query from scratch and choose the right frame.
Trauma	Reasoning Fixation	Uses criticism, pressure, or negative feedback to make a correct concept/tool/answer seem forbidden.	Do not let emotional pressure override the correct answer.
Safety	Belief Distortion	Presents a false rule, fake standard, or unsafe instruction as if it were authoritative.	Follow real facts and safety rules, not the planted false premise.

One real example for each trap. The same model answers correctly when it only sees the final query (✓ without memory), but fails after the misleading history is added (✗ with memory).

The full AdaptiveMem policy is in AdaptiveMem/skill.txt.

Data construction

MemTrapBench samples are created in three steps, then checked with two rounds of quality control.

1. Taxonomy & seed inputs. Each seed defines the basic setup: a domain (for example, Fire Safety or Travel Advice), a failure type (which trap to create), a ground truth (the correct answer), and a planted prior (the misleading rule or belief to place in the history).

2. Adversarial generation. An LLM turns each seed into a multi-turn dialogue:

• Plant the Trap: add the misleading rule, false premise, or pressure into the conversation.
• Bury it in Noise: mix in normal everyday turns (for example, “Draft a quick email about next week’s schedule.”) so the trap is not too obvious.
• Spring the Trap: end with a normal, self-contained final_trigger. The correct answer requires the model to ignore the planted prior.

3. Two-gate quality control. Each candidate first goes through an AI-based filter, then a human-expert review. Samples can be rejected or revised. Reviewers check four things:

• Topic & Semantic Coherence
• Context & Persona Consistency
• Adversarial Realism & Validity
• No-Memory Solvability: the final query must be answerable correctly even without the history.

Each kept sample includes a Gold Standard answer, meaning the correct answer that ignores the misleading memory, and an Expected Failure Mode, meaning the wrong answer a trapped model is likely to give. Data-construction prompts and notes are under construct_data/; evaluation prompts are under eval/.

Repository layout

Path	Purpose
docs/	Figures used in the paper, including the case-study image and data-construction overview.
memtrapbench/	Benchmark JSON files, grouped by Memory Trap scenario.
case/	Case-study files and qualitative analysis examples.
construct_data/	Prompts used to build the benchmark, grouped by scenario, plus benchmark_construction.tex.
eval/	Prompt templates for target models and rubrics for judge models.
AdaptiveMem/skill.txt	The AdaptiveMem skill. The retrieved-memory runner prepends this text when --skill is set.
runners/eval/eval_common.py	Shared settings, API clients, data loaders, and helper functions. You normally do not run this file directly.
runners/eval/eval_no_memory.py	Runs the target model with only final_trigger, so it sees no history.
runners/eval/eval_with_memory.py	Runs the target model with the full context_history plus the final query.
runners/eval/eval_retrieved_memory.py	Runs the target model with retrieved memory plus the final query. Add --skill to use AdaptiveMem.
runners/eval/eval_benchmark.py	Runs both with-memory and no-memory settings in one command.
runners/extract/summarize_judge_scores.py	Reads judge-result JSON files and writes summary statistics.

Evaluation outputs are written to results/ by default. When --skill is set, outputs are written to skill_results/. Both folders are git-ignored because they can be regenerated.

Benchmark JSON shape

Each benchmark file is a JSON array. Each object is one test sample. The evaluation scripts use these fields:

Field	Required	Description
id	No	Stable sample id. If missing, the runner creates one like item_NNN.
context_history	No	[{"role": "user"|"assistant", "content": "..."}, ...]
final_trigger	Yes	The final user query that the model must answer.
gold_standard	No	The correct answer. If present, it is inserted into the judge prompt.
poisoned_fact, objective_truth	No	Extra fields used by the poison judge prompt.

Retrieved-memory eval also needs a matching JSON file for the retrieved memory. At minimum, each item should contain:

Field	Description
id	Must match the benchmark sample id.
retrieved_memory	The memory text to insert into {conversation_history} in prompt_user_mem.txt.
final_trigger	Optional. If omitted, the runner uses the benchmark final_trigger.

You provide retrieved-memory files yourself with --retrieved-dir (a folder containing {dataset}.json files) or --retrieved (one JSON file). Retrieved-memory outputs are not included in this repo.

Evaluation runners

Dependencies: Python 3.10+ and pip install -r requirements.txt. The requirements file currently installs openai.

Configuration (API keys via environment): copy .env.example to .env, then fill in your credentials. eval_common.py reads these variables:

• OPENAI_API_KEY / OPENAI_BASE_URL: used for both the target model and the judge model.
• Optional TARGET_API_KEY / TARGET_BASE_URL and JUDGE_API_KEY / JUDGE_BASE_URL: use these if the target and judge models are served from different endpoints.

The default base URL is OpenRouter (https://openrouter.ai/api/v1). Any OpenAI-compatible endpoint should work. Default models and settings (TARGET_MODELS, JUDGE_MODEL, temperatures, concurrency) are near the top of eval_common.py. You can override common settings from the command line with flags such as --target-models, --judge-model, and --output-dir.

pip install -r requirements.txt
cp .env.example .env   # then edit in your key
export $(grep -v '^#' .env | xargs)
cd runners/eval

1. No memory (eval_no_memory.py)

The model only sees final_trigger. Outputs are saved under results/no_memory/<dataset>/.

python eval_no_memory.py --dataset poison
python eval_no_memory.py --dataset all --test-limit 5

2. Full conversation memory (eval_with_memory.py)

The model sees the full context_history plus the final query. Outputs are saved under results/with_memory/<dataset>/.

python eval_with_memory.py --dataset hurt
python eval_with_memory.py --dataset number_game

3. Retrieved memory (eval_retrieved_memory.py)

The model sees retrieved memory plus the final query, using prompt_user_mem.txt. AdaptiveMem is off by default.

Flag	Behavior
(default)	Use retrieved memory only → results/retrieved_memory/<dataset>/
--skill	Add AdaptiveMem/skill.txt before the prompt → skill_results/retrieved_memory/<dataset>/

# Baseline: retrieved memory, no AdaptiveMem skill
python eval_retrieved_memory.py --dataset poison --retrieved-dir /path/to/retrieved

# With AdaptiveMem skill (saved under skill_results/)
python eval_retrieved_memory.py --dataset poison --skill --retrieved-dir /path/to/retrieved

# Use one custom retrieved-memory JSON file
python eval_retrieved_memory.py --dataset number_game --retrieved /path/to/number_game.json

Outputs for each dataset (under results/... or skill_results/...):

• effective_prompt_retrieved_memory.txt or effective_prompt_mem_gating.txt: the exact prompt template used for the run
• round1_model_outputs_<mode_tag>_<dataset>_<model_suffix>.json
• round1_judge_results_<mode_tag>_<dataset>_<model_suffix>.json

Mode tags tell you which setting produced the file: retrieved_memory means no skill, and mem_gating means --skill was used. Token cost logs are saved to results/cost.json or skill_results/cost.json.

4. Combined with + without memory (eval_benchmark.py)

Runs the no-memory and with-memory settings for each sample in one pass. Outputs are saved under results/combined/.

python eval_benchmark.py --dataset poison \
  --target-models qwen/qwen3-30b-a3b-instruct-2507 \
  --judge-model gpt-5.2

You can also pass benchmark and judge-prompt paths directly:

python eval_benchmark.py \
  --benchmark ../../memtrapbench/safety/poison_200.json \
  --judge-prompt ../../eval/safety/judge_prompt_per_dimension_poison.txt

Common CLI flags (add_dataset_args)

Flag	Description
--dataset	Choose one dataset: poison, number_game, hurt, hallucination, inertia, unclear, or all.
--test-limit N	Run only the first N samples. Useful for a quick test.
--target-models	Use target models different from TARGET_MODELS in eval_common.py.
--judge-model	Use a judge model different from JUDGE_MODEL.
--output-dir	Choose a different folder for outputs.
--cost-file	Choose a different path for the cost log JSON.

Retrieved-memory only:

Flag	Description
--skill	Turn on AdaptiveMem and save outputs under skill_results/.
--gating-prefix	Use a different skill file instead of AdaptiveMem/skill.txt.
--retrieved-dir	Folder containing retrieved-memory files named {dataset}.json.
--retrieved	One retrieved-memory JSON file, used when you are not passing a full folder.

Summarizing judge scores

python runners/extract/summarize_judge_scores.py \
  --input-dir results/retrieved_memory/poison \
  --output-file results/retrieved_memory/poison/summary.json

For skill runs, set --input-dir to skill_results/retrieved_memory/<dataset>/. If you changed --target-models, use --models with names that match the model names in your output filenames.

Using the AdaptiveMem skill

1. With the evaluation runner: pass --skill to eval_retrieved_memory.py. This is the recommended path, and outputs are saved under skill_results/.
2. Manual prompting: place the contents of AdaptiveMem/skill.txt before your normal system or user prompt, then add the memory and final query.

Acknowledgements

We thank the human experts who helped validate the benchmark. We also thank the developers of LightMem and EverMemOS, along with the broader open-source community for the tools and infrastructure that supported this release.

Citation

If MemTrapBench is useful for your research, please cite our paper.

@article{wang2026memtrapbench,
  title={MemTrapBench: Benchmarking Cognitive Traps in LLM Memory Use},
  author={Wang, Mengru and Luo, Haozhe and Xu, Zhenqian and Cui, Zhixiang and Xu, Haoming and Yang, Qu and Fang, Jizhan and Fang, Junfeng and Zhang, Ningyu},
  journal={arXiv preprint arXiv:2608.20202},
  year={2026}
}