CaliMem: Search-Efficient Data Agent via Calibrated Memory

Hanze Li; Athan Z. Li; Feiyang Kang; Jiachen T. Wang; Yuhan Chen; Ruishan Liu; Ruoxi Jia

cs.AIOct 25, 2026

CaliMem: Search-Efficient Data Agent via Calibrated Memory

Authors: Hanze Li, Athan Z. Li, Feiyang Kang, Jiachen T. Wang, Yuhan Chen, Ruishan Liu, Ruoxi Jia

Organizations: Virginia Tech · University of Southern California · Princeton University

Abstract

Data curation is a first-order design choice in machine learning, yet tailoring training data to a downstream model and task remains costly in labor, expertise, and compute. Datasets are therefore often processed through fixed pipelines and reused across applications. We show that this assumption is fragile: across 55 dataset--model tasks spanning biological domains, the same curation recipe can help one model while hurting another. Effective curation must adapt to the downstream setting, motivating automated search over data recipes. We first build an agentic system that searches for curation policies at deployment, improving data utility via test-time curation (TTC). Although achieving performance gains, its search is driven largely by recent feedback and remains local. Adding persistent memory helps the agent retain evidence across iterations and reveals a deeper problem: it forms claims that extend beyond the conditions in which they were observed, sometimes abandoning an entire strategy after a single unsuccessful trial. To address this problem, we introduce CaliMem, a calibrated-memory agentic loop that requires the agent to predict the outcome of each proposed policy before executing it. Comparing this prediction with the observed result gives the agent a reality check on the claims behind its proposal. It uses the prediction error to retain supported claims, revise mistaken ones, and specify the conditions under which a conclusion holds. This makes the agent's beliefs testable, reduces overgeneralization, and supports more effective long-horizon exploration. Across biological and LLM post-training tasks, CaliMem consistently outperforms fixed curation baselines and published methods, and exceeds full-data training by up to 19.4%. Code is open-sourced.

Figures & tables

Figure 1: Curation depends on the downstream model, and calibrated memory guides the search for a suitable policy. (a) 12 standard recipes on a single-cell batch-integration dataset, ranked by the performance of two models trained on their half-size subsets. The recipe that ranks 1st for cVAE ranks 11th for MMD. (b) A basic agent proposes a policy, trains and evaluates the model, and responds to the latest score change. CaliMem records evidence and claims in memory, predicts the outcome of each proposed policy, and compares its prediction with the result. The comparison guides revisions to the claims used in subsequent iterations.

Iteration	Score	Status	What changed	Agent’s note
1	61.30	keep	Random 50% subset (baseline)	The score to beat.
4	63.63	keep	Give each perturbation at least 100 cells (all of its cells if it has fewer); split the rest by perturbation size	New best: small perturbations now keep all their cells, and large ones keep their share.
5	63.83	keep	Control cells: $\sim$ 4,900 $\to$ 2,000	New best: cutting controls frees cells for perturbations at little cost.
6	67.53	keep	Control cells: 2,000 $\to$ 1,000	Large gain: more perturbation cells matter more than more controls.
7	66.60	discard	Control cells: 1,000 $\to$ 500	Worse: 500 controls are too few, so 1,000 is the lower limit. Next, keep 1,000 controls and try a higher per-perturbation minimum.
8	65.32	discard	Control cells back to 1,000; per-perturbation minimum: 100 $\to$ 150	Worse: a higher minimum takes cells from large perturbations, so 100 is the right minimum.

Table 1: Excerpt from the trajectory of the basic agent on single-cell perturbation-response prediction (Norman19, GEARS); the complete log of all 15 iterations is in Table 10 . Scores are evaluation cosine similarity between predicted and observed log fold changes (

\times

100). A perturbation is the set of cells that received the same genetic perturbation, and control cells are unperturbed. Keep means the policy became the incumbent; discard means the agent reverted to the incumbent. Each row lists what the agent changed and a paraphrase of the note it wrote after seeing the score.

Figure 2: Without a prediction, one failure closes a whole direction; with a prediction, the failure is tied to its condition. Cell-type annotation ( pancreas_crosstech , SVM; evaluation macro-F1). Both agents see rebalancing fail on top of a confidence filter (0.860

\to

0.853). Top: the failure becomes a claim about rebalancing as a whole. Bottom: the gap between the predicted

+0.01

and the observed

-0.007

ties the failure to the filter, and rebalancing stays open. The top row follows the memory-only run; the bottom row adapts the memory + prediction run on the same task to the same move, and its best score is from that run.

Figure 3: Main results on the biology domains. Each entry is the mean, over the dataset–model pairs of the domain, of the best headline score reached within the 15-iteration budget, scaled to 0–100. Random is the seeded subset every run starts from, Full data trains on the entire pool without curation, and the fixed-pipeline rows are standard curation recipes of each domain: the mean

\pm

sd over the 12 recipes, and the best recipe chosen separately for every subtask. Detailed results can be found in Appendix D .

Figure 4: Main results on LLM post-training tasks. Each entry is the best mean accuracy (%) over the setup’s benchmark suite reached within the 10-iteration budget. Random, LESS ( Xia et al., 2024 ) , DELIFT ( Agarwal et al., 2025 ) , and the agents all select a subset of the same size from the same pool (the selection ratio is given under each backbone): the dashed line and band mark Random as the mean

\pm

sd over seeds, LESS and DELIFT are human-designed selection methods, and Entire dataset trains on the entire pool without curation. Points and segments left of the dashed line fall below Random. Detailed results can be found in Appendix D .

Figure 5: Best score so far per iteration under the basic agent and CaliMem (Claude Opus 4.8). Left two panels: single-cell batch integration and cell-type annotation, averaged over the 15 dataset–model tasks of each domain, 15 iterations. Right two panels: code instruction tuning (Qwen2.5-Coder-1.5B) and instruction tuning (Llama-2-7B on the LESS mixture), 10 iterations. In every panel CaliMem passes the basic agent’s final score within the first 5 iterations.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Domain	Dataset	Pool	Budget	Models	Headline metric	Iter.
Protein variant effect	blat_ecolx	2,894	1,447	ESM-2 ridge, ESM-2 MLP, AA-property ridge	Spearman	15
	cp2c9_human	3,736	1,868
	tpmt_human	2,188	1,094
	pten_human	4,387	2,194
Non-coding variant effect	clinvar_splice	11,163	5,582	DNABERT logistic, 4-mer logistic	AUPRC	15
	mpra_hepg2	10,978	5,489

Appendix

Table 2: The benchmark grid, pairing each dataset with the models listed for it. Pool is the number of training rows from which the agent selects, and Budget is the number of rows that each submission must contain. Iter. is the number of iterations per run.

Figure 7: Selecting 50% data using standard, human-designed data processing pipelines for biology tasks. Biology tasks typically apply standard, human-designed processing pipelines on the datasets, and apply the sanitized datasets for training different ML models in downstream applications. Yet, different ML models and training pipelines carry different assumption for the data. The fixed data cleaning protocol may yield substantially different outcomes in different use cases. As depicted in the figure, in every case, there are more processing pipelines underperforming training on unprocessed data.

Recipe	Kind	Rule	Default	Practice / source
random	–	uniform subset, floors reserved	–	seeded baseline; random training subsets in training-size studies ( Gelman et al., 2021 ; Hsu et al., 2022 )
drop_dead_q20 , _q35	F	drop variants below the $q$ -quantile of the assay score	$q=0.20,\,0.35$	removal of low-fitness “holes” ( Wittmann et al., 2021 )
cap_per_position_k10 , _k6	A	at most $k$ variants per residue position	$k=10,\,6$	position breadth over depth ( Sruthi and Prakash, 2020 )
position_outlier_2sd	F	drop variants $>2$ sd from their position’s median	2 sd	proxy for measurement-error and read-count filtering ( Rubin et al., 2017 ; Faure et al., 2020 )
flip_low_bin	F	keep only variants in ProteinGym’s non-functional bin	bin $=0$	FLIP low-versus-high split ( Dallago et al., 2021 )
flip_below_median	F	keep only variants below the median score	median	median analog of the same split ( Dallago et al., 2021 )

Appendix

Table 3: Protein variant-effect recipes (four assays, three frozen regressors). Kind: F filter, A allocation, R ranking.

Recipe	Kind	Rule	Default	Practice / source
random	–	uniform subset, floors reserved	–	seeded baseline
all_positives	R	keep every positive, fill with random negatives	–	majority-class undersampling ( Whalen et al., 2022 )
balanced_1to1 , _1to2	A	target positive:negative ratio	1:1, 1:2	class rebalancing ( Kircher et al., 2014 ; Whalen et al., 2022 )
gc_match_0.02 , _0.05	F	at most as many negatives as positives per GC bin	bin width 0.02, 0.05	GC-matched null sets ( Ghandi et al., 2014 ; Lee et al., 2015 )
dedup_32mer , _24mer	F	drop windows sharing any exact $k$ -mer with a kept window	$k=32,\,24$	redundancy / LD pruning ( Purcell et al., 2007 ; Whalen et al., 2022 )
consequence_stratified	A	equal quota per consequence class	–	consequence annotations ( Landrum et al., 2018 ; Marin et al., 2024 )

Appendix

Table 4: Non-coding variant-effect recipes (three datasets, two frozen classifiers).

Recipe	Kind	Rule	Default	Practice / source
random	–	uniform subset, floors reserved	–	seeded baseline
qc_fixed_mito20 , _mito5	F	genes $\geq 200$ , counts $\geq 500$ , mitochondrial fraction below threshold	20%, 5%	tutorial thresholds ( Wolf et al., 2018 ; Hao et al., 2021 ; Luecken and Theis, 2019 )
qc_mad5 , qc_mad3	F	per-batch MAD outliers on log counts and log genes (both tails) and on mitochondrial fraction (upper tail)	5, 3 MADs	MAD-based QC ( Lun et al., 2016 ; Germain et al., 2020 ; Heumos et al., 2023 )
doublet_scrublet	F	remove doublets called by Scrublet within each batch	default threshold	Scrublet ( Wolock et al., 2019 ) ; scDblFinder ( Germain et al., 2022 )
class_cap_500 , _1000	A	at most $N$ cells per class, remainder water-filled	$N=500,\,1000$	down-weighting abundant types ( Domínguez Conde et al., 2022 )
class_balanced	A	equal quota per class	–	type-balanced sampling ( Domínguez Conde et al., 2022 ) ; errors on small populations ( Abdelaal et al., 2019 )

Appendix

Table 5: Cell-type annotation recipes (three datasets, five frozen classifiers).

Recipe	Kind	Rule	Default	Practice / source
random	–	uniform subset, floors reserved	–	seeded baseline
batch_equal	A	equal quota per batch	–	unequal batch sizes ( Johnson et al., 2007 ; Korsunsky et al., 2019 ; Tran et al., 2020 )
batch_sqrt	A	quota proportional to $\sqrt{\text{batch size}}$	–	square-root allocation
qc_fixed_mito20	F	genes $\geq 200$ , counts $\geq 500$ , mitochondrial fraction $<20\%$	20%	threshold-based QC ( Luecken and Theis, 2019 )
qc_mad5	F	per-batch outliers beyond 5 MADs	5 MADs	MAD-based QC ( Lun et al., 2016 ; Heumos et al., 2023 )
doublet_scrublet	F	remove Scrublet-called doublets per batch	default threshold	Wolock et al. (2019)

Appendix

Table 6: Batch integration recipes

Component	Value
Initial model	OLMo-2-0425-1B (base)
Training hardware	1 NVIDIA A100 GPU
Training epochs	1
Trainable modules	All parameters (full fine-tuning)
Per-device batch size	1
Gradient accumulation	16

Appendix

Table 7: Fine-tuning settings for OLMo-2-1B on the Tülu 3 SFT mixture.

Component	Value
Initial model	Llama-2-7B (base)
Training hardware	1 NVIDIA A100 GPU
Training epochs	1
Trainable modules	All parameters (full fine-tuning)
Per-device batch size	2
Gradient accumulation	8

Appendix

Table 8: Fine-tuning settings for Llama-2-7B on the LESS mixture.

Component	Value
Initial model	Qwen2.5-Coder-1.5B-Instruct
Training hardware	1 NVIDIA A100 GPU
Training epochs	1
Trainable modules	All parameters (full fine-tuning)
Per-device batch size	4
Gradient accumulation	4

Appendix

Table 9: Fine-tuning settings for Qwen2.5-Coder-1.5B-Instruct on OpenCodeInstruct.

Iteration	Score	Status	What changed from the previous iteration	Agent’s note (paraphrased)
1	61.30	keep	Random 50% subset (baseline)	The score to beat.
2	57.58	discard	All $\sim$ 9,700 control cells; the same number of cells for every perturbation; within each perturbation, the cells with the highest counts	Worse than baseline. At least one of the three changes hurt, most likely picking the highest-count cells, which may be doublets.
3	60.83	discard	Controls back to $\sim$ 4,900; random cells within each perturbation again	Still slightly below baseline. Picking the highest-count cells was the problem; the same number per perturbation is also slightly worse than splitting by size.
4	63.63	keep	Give each perturbation at least 100 cells (all of its cells if it has fewer); split the rest by perturbation size	New best: small perturbations now keep all their cells, and large ones keep their share.
5	63.83	keep	Control cells: $\sim$ 4,900 $\to$ 2,000	New best: cutting controls frees cells for perturbations at little cost.
6	67.53	keep	Control cells: 2,000 $\to$ 1,000	Large gain: more perturbation cells matter more than more controls.

Appendix

Table 10: Complete trajectory of the basic agent on single-cell perturbation-response prediction (Norman19, GEARS), all 15 iterations; Table 1 shows an excerpt. Scores are validation cosine similarity between predicted and observed log fold changes (

\times

100). A perturbation is the set of cells that received the same genetic perturbation, and control cells are unperturbed. Each row lists what the agent changed from the previous iteration and a paraphrase of the note it wrote after seeing the score. At iteration 10 the agent intended to keep the 1,000 control cells with the highest counts, but a column-name bug left the policy identical to iteration 6.

Iteration	Score	Status	What changed from the previous iteration	Agent’s note (paraphrased)
1	84.92	keep	Random 50% subset (baseline)	The score to beat. One cell type has a recall of only 0.52 and drags the average down.
2	84.42	discard	The same number of cells for every cell type, capping the large ones	Worse: the weakest cell type improves, but the large ones lose cells. The SVM already weights classes by frequency, so balancing the counts is redundant.
3	89.84	keep	No count balancing; keep the 4,082 cells whose labels a cross-validated logistic regression predicts most confidently	Large gain: dropping ambiguous or mislabeled cells sharpens the linear boundary, and the weakest cell type’s recall rises from 0.52 to 0.79.
4	88.35	discard	Score cells with a classifier trained on the other 10x protocol instead of cross-validation	Worse: a scorer trained on one protocol is less reliable than cross-validation over the whole pool.
5	88.96	discard	Back to the iteration-3 cells; drop the sparsest genes (2,000 $\to$ 1,684 genes)	Worse: removing genes hurts; the SVM needs the full 2,000-gene panel.
6	89.17	discard	All 2,000 genes again; score cells with the evaluation model itself (the SVM) instead of logistic regression	Worse: the SVM scorer removes too many B cells and megakaryocytes; the gentler logistic-regression scorer stays best.

Appendix

Table 11: Complete trajectory of CaliMem (Claude Opus 4.8) on single-cell cell type annotation (PBMC cross-protocol, SVM), all 15 iterations. Scores are validation macro-F1 (

\times

100), the F1 score averaged over five cell types: cytotoxic T cells, CD4 + T cells, CD14 + monocytes, B cells, and megakaryocytes. The training pool has 8,165 cells from two 10x protocols, and the validation cells come from a third protocol. Each iteration keeps 4,082 cells and can also choose the genes the model uses; the default is a panel of 2,000 highly variable genes. The SVM already weights each class inversely to its frequency. Each row lists what the agent changed from the previous iteration and a paraphrase of the note it wrote after seeing the score. Iterations 5 and 10 (bold) change the gene panel in opposite directions, and the agent’s memory joins the two into one claim: the 2,000-gene panel is a sweet spot.

Model	Agent variant	\Block Protein
(Spearman)	\Block Non-coding
(AUPRC)	\Block Cell type
(macro-F1)	\Block Batch
(scIB overall)	\Block Perturbation
(cosine)
\Block [l]1-2Full data (baseline)	-	35.60	13.38	64.65	68.79	46.04

Appendix

Table 12: Main results on the biology domains. Each entry is the mean, over the dataset–model pairs of the domain, of the best headline score reached within the 15-iteration budget, scaled to 0–100. Random is the seeded subset every run starts from, Full data trains on the entire pool without curation, and the fixed-pipeline rows are standard curation recipes of each domain: the mean

\pm

sd over the 12 recipes, and the best recipe chosen separately for every subtask. Percentages (%) are relative change against Full data; entries more than 5% above Full data are in green and entries more than 5% below it in red . Within each block, the best entry in each column is in bold .

Model	Agent variant	\Block OLMo-2-1B
(select 1.1%)	\Block Llama-2-7B
(select 3.7%)	\Block Qwen2.5-Coder-1.5B
(select 0.2%)
\Block [l]1-2Random (baseline)	-	31.00 $\pm$ 0.52	42.44 $\pm$ 0.28	57.72 $\pm$ 0.24
\Block [l]1-2Entire dataset	-	30.35	45.70	62.03
\Block [l]1-2LESS ( Xia et al., 2024 )	-	31.81 (+2.6%)	45.05 (+6.1%)	58.38 (+1.1%)

Appendix

Table 13: Main results on LLM post-training tasks. "Random", "LESS", "DELIFT", and "Agent variant" all select training data of the same size from the same pool. "Random" is a randomly downsampled subset, LESS ( Xia et al., 2024 ) and DELIFT ( Agarwal et al., 2025 ) are human-designed selection methods, whereas "Entire dataset" trains on the entire pool without curation. Each agent entry is the best mean accuracy over the setup’s benchmark suite reached within the 10-iteration budget. Percentages (%) are relative change against Random; entries more than 5% above Full data are in green and entries more than 5% below it in red .

Figure 8: Selecting 50% of the training data for machine-learning models in biology tasks. In biology, datasets are typically cleaned with standard, human-designed processing pipelines, and the processed data are then used to train different downstream models, which often make different assumptions about their data. A fixed cleaning protocol can therefore have very different effects across use cases: in every task, the pipelines that underperform training on unprocessed data (a random 50% subset) outnumber those that outperform it. Agentic data curation instead optimizes the training data directly for the target model pipeline, giving large performance gains through test-time scaling and outperforming fixed pipelines. CaliMem further strengthens the data agent and consistently outperforms the default agent.

Figure 9: The memory becomes more predictive as evidence accumulates. Absolute error when a fresh agent uses the CaliMem memory frozen after iteration $k$ (rows) to predict a later, unseen iteration (columns). Results are averaged over single-cell cell-type annotation tasks from Claude Sonnet 5 runs; lighter cells indicate smaller error. A fixed memory becomes less accurate as it extrapolates further ahead, whereas predictions of the same target iteration improve when the memory is fitted on more preceding experiments.

Panel of Figure 10 — Figure 10: Forecasting error of a CaliMem fitted on the first $k$ iterations (rows) when a fresh agent uses it to predict the score of every later iteration (columns), for the three domains not shown in the main text. Each entry is the absolute error in the domain’s headline metric, averaged over all tasks of the domain; lighter is smaller error.

Training Data Curation

Training data curation improves the quality, diversity, organization, and reliability of datasets used to train machine-learning models.

See all 156 papers

CurateEvo: Data-Curation Evolving for Agentic Post-Training
Jul 7, 2026Dingzirui Wang, Xuanliang Zhang, Keyan Xu +2Training Data Curation LLM Agent Training
Harbin Institue of Technology
Can Generalist Agents Automate Data Curation?
Jun 2, 2026Feiyang Kang, Hanze Li, Adam Nguyen +5AI Coding Agents Training Data Curation
1Virginia Tech · University of Illinois Urbana-Champaign · University of Wisconsin-Madison +1
Just-in-Time Memory: Learning to Curate Task-Adaptive Memory for LLM Agents
Sep 23, 2026Yefan Zhou, Yang Li, Zeyu Leo Liu +2LLM Agent Memory Agent Memory Management
Salesforce AI Research

Jul 7, 2026J/K move · Enter open · S save

CurateEvo: Data-Curation Evolving for Agentic Post-Training

Dingzirui Wang, Xuanliang Zhang, Keyan Xu, Qingfu Zhu+1

Harbin Institue of Technology

Training Data Curation LLM Agent Training LLM Post-Training

Open paper AskPDF arXiv

Large language model (LLM) agents require post-training methods that can improve long-horizon decision making from environment feedback. However, existing agentic post-training pipelines often treat data curation as a fixed preprocessing step, focusing mainly on data augmentation while neglecting filtering, refinement, and adaptation to downstream failures. We propose CurateEvo, a failure-driven dynamic evolution framework for agentic post-training data curation. CurateEvo represents the curation strategy as executable code and iteratively rewrites it using failed trajectories from a held-out development set. At each epoch, the evolved strategy transforms a fixed raw corpus into supervised fine-tuning data, reinforcement learning data, and an inference-time memory bank. The evolution process first improves effectiveness by diagnosing recurring failure modes and augmenting, filtering, or refining data accordingly, and then improves efficiency by pruning redundant or low-utility training turns under a cost-aware objective. Experiments on ACEBench-Agent, BFCL-V4, and τ^2-Bench under both labeled and wild-data settings show that CurateEvo consistently outperforms prior curation methods, improving average scores by 3.2 and 2.7 points, respectively. Further analyses demonstrate that CurateEvo is compatible with different post-training recipes and substantially reduces curation overhead.

Abstract

Training Data Curation

Explore similar work

CurateEvo: Data-Curation Evolving for Agentic Post-Training

Can Generalist Agents Automate Data Curation?

Just-in-Time Memory: Learning to Curate Task-Adaptive Memory for LLM Agents