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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.
IterationScoreStatusWhat changedAgent’s note
161.30keepRandom 50% subset (baseline)The score to beat.
463.63keepGive each perturbation at least 100 cells (all of its cells if it has fewer); split the rest by perturbation sizeNew best: small perturbations now keep all their cells, and large ones keep their share.
563.83keepControl cells: ∼\sim 4,900 →\to 2,000New best: cutting controls frees cells for perturbations at little cost.
667.53keepControl cells: 2,000 →\to 1,000Large gain: more perturbation cells matter more than more controls.
766.60discardControl cells: 1,000 →\to 500Worse: 500 controls are too few, so 1,000 is the lower limit. Next, keep 1,000 controls and try a higher per-perturbation minimum.
865.32discardControl cells back to 1,000; per-perturbation minimum: 100 →\to 150Worse: 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+0.01 and the observed −0.007-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

DomainDatasetPoolBudgetModelsHeadline metricIter.
Protein variant effectblat_ecolx2,8941,447ESM-2 ridge, ESM-2 MLP, AA-property ridgeSpearman15
cp2c9_human3,7361,868
tpmt_human2,1881,094
pten_human4,3872,194
Non-coding variant effectclinvar_splice11,1635,582DNABERT logistic, 4-mer logisticAUPRC15
mpra_hepg210,9785,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.
RecipeKindRuleDefaultPractice / 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 , _q35Fdrop variants below the qq -quantile of the assay scoreq=0.20, 0.35q=0.20,\,0.35removal of low-fitness “holes” ( Wittmann et al., 2021 )
cap_per_position_k10 , _k6Aat most kk variants per residue positionk=10, 6k=10,\,6position breadth over depth ( Sruthi and Prakash, 2020 )
position_outlier_2sdFdrop variants >2>2 sd from their position’s median2 sdproxy for measurement-error and read-count filtering ( Rubin et al., 2017 ; Faure et al., 2020 )
flip_low_binFkeep only variants in ProteinGym’s non-functional binbin =0=0FLIP low-versus-high split ( Dallago et al., 2021 )
flip_below_medianFkeep only variants below the median scoremedianmedian 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.
RecipeKindRuleDefaultPractice / source
random–uniform subset, floors reserved–seeded baseline
all_positivesRkeep every positive, fill with random negatives–majority-class undersampling ( Whalen et al., 2022 )
balanced_1to1 , _1to2Atarget positive:negative ratio1:1, 1:2class rebalancing ( Kircher et al., 2014 ; Whalen et al., 2022 )
gc_match_0.02 , _0.05Fat most as many negatives as positives per GC binbin width 0.02, 0.05GC-matched null sets ( Ghandi et al., 2014 ; Lee et al., 2015 )
dedup_32mer , _24merFdrop windows sharing any exact kk -mer with a kept windowk=32, 24k=32,\,24redundancy / LD pruning ( Purcell et al., 2007 ; Whalen et al., 2022 )
consequence_stratifiedAequal 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).
RecipeKindRuleDefaultPractice / source
random–uniform subset, floors reserved–seeded baseline
qc_fixed_mito20 , _mito5Fgenes ≥200\geq 200 , counts ≥500\geq 500 , mitochondrial fraction below threshold20%, 5%tutorial thresholds ( Wolf et al., 2018 ; Hao et al., 2021 ; Luecken and Theis, 2019 )
qc_mad5 , qc_mad3Fper-batch MAD outliers on log counts and log genes (both tails) and on mitochondrial fraction (upper tail)5, 3 MADsMAD-based QC ( Lun et al., 2016 ; Germain et al., 2020 ; Heumos et al., 2023 )
doublet_scrubletFremove doublets called by Scrublet within each batchdefault thresholdScrublet ( Wolock et al., 2019 ) ; scDblFinder ( Germain et al., 2022 )
class_cap_500 , _1000Aat most NN cells per class, remainder water-filledN=500, 1000N=500,\,1000down-weighting abundant types ( Domínguez Conde et al., 2022 )
class_balancedAequal 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).
RecipeKindRuleDefaultPractice / source
random–uniform subset, floors reserved–seeded baseline
batch_equalAequal quota per batch–unequal batch sizes ( Johnson et al., 2007 ; Korsunsky et al., 2019 ; Tran et al., 2020 )
batch_sqrtAquota proportional to batch size\sqrt{\text{batch size}}–square-root allocation
qc_fixed_mito20Fgenes ≥200\geq 200 , counts ≥500\geq 500 , mitochondrial fraction <20%<20\%20%threshold-based QC ( Luecken and Theis, 2019 )
qc_mad5Fper-batch outliers beyond 5 MADs5 MADsMAD-based QC ( Lun et al., 2016 ; Heumos et al., 2023 )
doublet_scrubletFremove Scrublet-called doublets per batchdefault thresholdWolock et al. (2019)
Appendix
Table 6: Batch integration recipes
ComponentValue
Initial modelOLMo-2-0425-1B (base)
Training hardware1 NVIDIA A100 GPU
Training epochs1
Trainable modulesAll parameters (full fine-tuning)
Per-device batch size1
Gradient accumulation16
Appendix
Table 7: Fine-tuning settings for OLMo-2-1B on the Tülu 3 SFT mixture.
ComponentValue
Initial modelLlama-2-7B (base)
Training hardware1 NVIDIA A100 GPU
Training epochs1
Trainable modulesAll parameters (full fine-tuning)
Per-device batch size2
Gradient accumulation8
Appendix
Table 8: Fine-tuning settings for Llama-2-7B on the LESS mixture.
ComponentValue
Initial modelQwen2.5-Coder-1.5B-Instruct
Training hardware1 NVIDIA A100 GPU
Training epochs1
Trainable modulesAll parameters (full fine-tuning)
Per-device batch size4
Gradient accumulation4
Appendix
Table 9: Fine-tuning settings for Qwen2.5-Coder-1.5B-Instruct on OpenCodeInstruct.
IterationScoreStatusWhat changed from the previous iterationAgent’s note (paraphrased)
161.30keepRandom 50% subset (baseline)The score to beat.
257.58discardAll ∼\sim 9,700 control cells; the same number of cells for every perturbation; within each perturbation, the cells with the highest countsWorse than baseline. At least one of the three changes hurt, most likely picking the highest-count cells, which may be doublets.
360.83discardControls back to ∼\sim 4,900; random cells within each perturbation againStill slightly below baseline. Picking the highest-count cells was the problem; the same number per perturbation is also slightly worse than splitting by size.
463.63keepGive each perturbation at least 100 cells (all of its cells if it has fewer); split the rest by perturbation sizeNew best: small perturbations now keep all their cells, and large ones keep their share.
563.83keepControl cells: ∼\sim 4,900 →\to 2,000New best: cutting controls frees cells for perturbations at little cost.
667.53keepControl cells: 2,000 →\to 1,000Large 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.
IterationScoreStatusWhat changed from the previous iterationAgent’s note (paraphrased)
184.92keepRandom 50% subset (baseline)The score to beat. One cell type has a recall of only 0.52 and drags the average down.
284.42discardThe same number of cells for every cell type, capping the large onesWorse: the weakest cell type improves, but the large ones lose cells. The SVM already weights classes by frequency, so balancing the counts is redundant.
389.84keepNo count balancing; keep the 4,082 cells whose labels a cross-validated logistic regression predicts most confidentlyLarge gain: dropping ambiguous or mislabeled cells sharpens the linear boundary, and the weakest cell type’s recall rises from 0.52 to 0.79.
488.35discardScore cells with a classifier trained on the other 10x protocol instead of cross-validationWorse: a scorer trained on one protocol is less reliable than cross-validation over the whole pool.
588.96discardBack 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.
689.17discardAll 2,000 genes again; score cells with the evaluation model itself (the SVM) instead of logistic regressionWorse: 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.
ModelAgent 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.6013.3864.6568.7946.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 .
ModelAgent 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.5242.44 ±\pm 0.2857.72 ±\pm 0.24
\Block [l]1-2Entire dataset-30.3545.7062.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 kk (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.
Figure 10: Forecasting error of a CaliMem fitted on the first kk 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.
More in this topic

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

Explore similar work

CardsList
Jul 7, 2026

CurateEvo: Data-Curation Evolving for Agentic Post-Training

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.

Dingzirui Wang, Xuanliang Zhang, Keyan Xu +2

Harbin Institue of Technology

Training Data CurationLLM Agent TrainingLLM Post-Training
AskPDF
Jun 2, 2026

Can Generalist Agents Automate Data Curation?

Curating training data is among the most consequential yet labor-intensive parts of modern AI development: practitioners iteratively propose, implement, evaluate, and revise data policies against noisy benchmark feedback. We ask whether generalist coding agents can automate this data-curation loop. We introduce Curation-Bench, an agent-centric benchmark that fixes the model, training recipe, and evaluation suite while giving agents command-line access to inspect data, implement policies, submit them to a fixed training/evaluation pipeline, and revise. In a vision-language instruction-tuning instantiation, out-of-the-box agents reach strong published data-selection baselines within ten iterations. However, trajectory analysis reveals a persistent execution-research gap: agents mainly tune local policy variants rather than explore new policy families, even when given strategy guides and paper references. Scaffolds requiring each iteration to cite, instantiate, and adapt a prior method shift agents toward method-guided exploration. The scaffolded agent autonomously composes -- without human design input -- a data-selection policy that outperforms strong published baselines at one-tenth their data budget. Overall, current agents can run the curation loop, but reliable data research requires scaffolded method adaptation, not open-ended prompting alone. Code and benchmark are open-sourced.

Feiyang Kang, Hanze Li, Adam Nguyen +5

1Virginia Tech · University of Illinois Urbana-Champaign · University of Wisconsin-Madison +1

AI Coding AgentsTraining Data CurationData SelectionTraining Data SelectionAgent Benchmarks
AskPDF
Sep 23, 2026

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

Agentic memory systems reuse past experience to improve future performance, yet most existing designs curate memory at write time: once a task is completed, its trajectory is distilled into a fixed artifact, such as a reflection, workflow, skill, or reasoning strategy, that is later retrieved by similarity. This forces the system to decide what is worth remembering before the future query is known, irreversibly discarding information and producing a query-independent summary that must serve many possible downstream tasks. Learning such a write-time curator is also difficult because the value of a storage decision may only become apparent when a relevant query arrives, potentially many tasks later, creating a long-horizon credit-assignment problem. We instead retain raw trajectories and defer curation until read time, when the current task is known. Given the retrieved traces and the new task, a memory curator synthesizes a compact, task-adaptive payload tailored to the immediate need. Because this payload is consumed on the same task, the curator can be trained directly from immediate task success, avoiding delayed utility signals and the need to artificially group related tasks. Across ALFWorld, WebShop, and τ2τ^2-bench, our Just-in-Time Memory (JitMem) consistently outperforms no-memory agents as well as heuristic and learned write-time memory methods, improving over the strongest baseline by 16.2, 16.3, and 3.9 absolute success-rate points, respectively. Notably, even an untrained curator is already competitive with or surpasses these baselines, showing that task-adaptive read-time curation itself is a major source of the gain; training the curator further compounds the improvement.

Yefan Zhou, Yang Li, Zeyu Leo Liu +2

Salesforce AI Research

LLM Agent MemoryAgent Memory ManagementLLM MemoryMemory-Augmented Language Model Agents
AskPDF