MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement
Open accessAbstract
Agents based on large language models (LLMs) for machine learning engineering (MLE) can automatically implement ML models via code generation. However, existing approaches to build such agents often rely heavily on inherent LLM knowledge and employ coarse exploration strategies that modify the entire code structure at once. This limits their ability to select effective task-specific models and perform deep exploration within specific components, such as experimenting extensively with feature engineering options. To overcome these, we propose MLE-STAR, a novel approach to build MLE agents. MLE-STAR first leverages external knowledge by using a search engine to retrieve effective models from the web, forming an initial solution, then iteratively refines it by exploring various strategies targeting specific ML components. This exploration is guided by ablation studies analyzing the impact of individual code blocks. Furthermore, we introduce a novel ensembling method using an effective strategy suggested by MLE-STAR. Our experimental results show that MLE-STAR achieves medals in 64% of the Kaggle competitions on the MLE-bench Lite, significantly outperforming the best alternative.
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Self-check
Percentages, sums and reported differences re-computed from the paper’s own numbers.
The self-check needs the paper’s full text.
Cites
Still being checked.
Research works printed but not linked: 48
- T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems , 2020.
- J. S. Chan, N. Chowdhury, O. Jaffe, J. Aung, D. Sherburn, E. Mays, G. Starace, K. Liu, L. Maksin, T. Patwardhan, et al. Mle-bench: Evaluating machine learning agents on machine learning engineering. International Conference on Learning Representations , 2025.
- T. Chen and C. Guestrin. Xgboost: A scalable tree boosting system. ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , 2016.
- A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference on Learning Representations , 2021.
- T. Elsken, J. H. Metzen, and F. Hutter. Neural architecture search: A survey. Journal of Machine Learning Research , 2019.
- N. Erickson, J. Mueller, A. Shirkov, H. Zhang, P. Larroy, M. Li, and A. Smola. Autogluon-tabular: Robust and accurate automl for structured data. arXiv preprint arXiv:2003.06505 , 2020. · arXiv
- L. Fan, F. Zhang, H. Fan, and C. Zhang. Brief review of image denoising techniques. Visual computing for industry, biomedicine, and art , 2019.
- W. Fan, E. Zhong, J. Peng, O. Verscheure, K. Zhang, J. Ren, R. Yan, and Q. Yang. Generalized and heuristic-free feature construction for improved accuracy. SIAM International Conference on Data Mining , 2010.
- M. Feurer, K. Eggensperger, S. Falkner, M. Lindauer, and F. Hutter. Auto-sklearn 2.0: Hands-free automl via meta-learning. Journal of Machine Learning Research , 2022.
- S. Guo, C. Deng, Y. Wen, H. Chen, Y. Chang, and J. Wang. DS-agent: Automated data science by empowering large language models with case-based reasoning. International Conference on Machine Learning , 2024.
- K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. IEEE Conference on Computer Vision and Pattern Recognition , 2016.
- N. Hollmann, S. Müller, and F. Hutter. Large language models for automated data science: Introducing caafe for context-aware automated feature engineering. Advances in Neural Information Processing Systems , 2023.
- D. Holzmüller, L. Grinsztajn, and I. Steinwart. Better by default: Strong pre-tuned mlps and boosted trees on tabular data. Advances in Neural Information Processing Systems , 2024.
- S. Hong, Y. Lin, B. Liu, B. Liu, B. Wu, C. Zhang, C. Wei, D. Li, J. Chen, J. Zhang, et al. Data interpreter: An llm agent for data science. arXiv preprint arXiv:2402.18679 , 2024. · arXiv
- F. Horn, R. Pack, and M. Rieger. The autofeat python library for automated feature engineering and selection. Joint European Conference on Machine Learning and Knowledge Discovery in Databases , 2019.
- X. Hu, Z. Zhao, S. Wei, Z. Chai, Q. Ma, G. Wang, X. Wang, J. Su, J. Xu, M. Zhu, et al. Infiagent-dabench: Evaluating agents on data analysis tasks. arXiv preprint arXiv:2401.05507 , 2024. · arXiv
- Q. Huang, J. Vora, P. Liang, and J. Leskovec. Mlagentbench: Evaluating language agents on machine learning experimentation. International Conference on Machine Learning , 2024a.
- Y. Huang, J. Luo, Y. Yu, Y. Zhang, F. Lei, Y. Wei, S. He, L. Huang, X. Liu, J. Zhao, et al. Dacode: Agent data science code generation benchmark for large language models. arXiv preprint arXiv:2410.07331 , 2024b. · arXiv
- Y. Ichihara, Y. Jinnai, T. Morimura, K. Abe, K. Ariu, M. Sakamoto, and E. Uchibe. Evaluation of best-of-n sampling strategies for language model alignment. Transactions on Machine Learning Research , 2025.
- N. Jain, K. Han, A. Gu, W.-D. Li, F. Yan, T. Zhang, S. Wang, A. Solar-Lezama, K. Sen, and I. Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code. International Conference on Learning Representations , 2025.
- Z. Jiang, D. Schmidt, D. Srikanth, D. Xu, I. Kaplan, D. Jacenko, and Y. Wu. Aide: Ai-driven exploration in the space of code. arXiv preprint arXiv:2502.13138 , 2025. · arXiv
- C. E. Jimenez, J. Yang, A. Wettig, S. Yao, K. Pei, O. Press, and K. Narasimhan. Swe-bench: Can language models resolve real-world github issues? International Conference on Learning Representations , 2024.
- H. Jin, Q. Song, and X. Hu. Auto-keras: An efficient neural architecture search system. ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , 2019.
- L. Jing, Z. Huang, X. Wang, W. Yao, W. Yu, K. Ma, H. Zhang, X. Du, and D. Yu. Dsbench: How far are data science agents to becoming data science experts? International Conference on Learning Representations , 2025.
- J. M. Kanter and K. Veeramachaneni. Deep feature synthesis: Towards automating data science endeavors. IEEE International Conference on Data Science and Advanced Analytics , 2015.
- J. L. Kolodner. An introduction to case-based reasoning. Artificial intelligence review , 1992.
- L. Kotthoff, C. Thornton, H. H. Hoos, F. Hutter, and K. Leyton-Brown. Auto-weka 2.0: Automatic model selection and hyperparameter optimization in weka. Journal of Machine Learning Research , 2017.
- E. LeDell and S. Poirier. H2O AutoML: Scalable automatic machine learning. ICML Workshop on AutoML , 2020.
- L. Li, H. Wang, L. Zha, Q. Huang, S. Wu, G. Chen, and J. Zhao. Learning a data-driven policy network for pre-training automated feature engineering. International Conference on Learning Representations , 2023.
- Z. Li, Q. Zang, D. Ma, J. Guo, T. Zheng, M. Liu, X. Niu, Y. Wang, J. Yang, J. Liu, et al. Autokaggle: A multi-agent framework for autonomous data science competitions. arXiv preprint arXiv:2410.20424 , 2024. · arXiv
- J. Nam, K. Kim, S. Oh, J. Tack, J. Kim, and J. Shin. Optimized feature generation for tabular data via llms with decision tree reasoning. Advances in Neural Information Processing Systems , 2024.
- R. S. Olson and J. H. Moore. Tpot: A tree-based pipeline optimization tool for automating machine learning. ICML Workshop on AutoML , 2016.
- F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, et al. Scikit-learn: Machine learning in python. Journal of Machine Learning Research , 2011.
- H. Pham, M. Guan, B. Zoph, Q. Le, and J. Dean. Efficient neural architecture search via parameters sharing. International Conference on Machine Learning , 2018.
- L. Prokhorenkova, G. Gusev, A. Vorobev, A. V. Dorogush, and A. Gulin. Catboost: unbiased boosting with categorical features. Advances in Neural Information Processing Systems , 2018.
- E. Real, A. Aggarwal, Y. Huang, and Q. V. Le. Regularized evolution for image classifier architecture search. AAAI Conference on Artificial Intelligence , 2019.
- S. Schmidgall, Y. Su, Z. Wang, X. Sun, J. Wu, X. Yu, J. Liu, Z. Liu, and E. Barsoum. Agent laboratory: Using llm agents as research assistants. arXiv preprint arXiv:2501.04227 , 2025. · arXiv
- Y. Shen, K. Song, X. Tan, D. Li, W. Lu, and Y. Zhuang. Hugginggpt: Solving ai tasks with chatgpt and its friends in hugging face. Advances in Neural Information Processing Systems , 2023.
- M. Tan and Q. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. International Conference on Machine Learning , 2019.
- G. Team, P. Georgiev, V. I. Lei, R. Burnell, L. Bai, A. Gulati, G. Tanzer, D. Vincent, Z. Pan, S. Wang, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.05530 , 2024. · arXiv
- H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023. · arXiv
- G. Wang, Y. Xie, Y. Jiang, A. Mandlekar, C. Xiao, Y. Zhu, L. Fan, and A. Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv: Arxiv-2305.16291 , 2023.
- X. Wang, B. Li, Y. Song, F. F. Xu, X. Tang, M. Zhuge, J. Pan, Y. Song, B. Li, J. Singh, et al. Openhands: An open platform for ai software developers as generalist agents. International Conference on Learning Representations , 2024.
- I. Watson and F. Marir. Case-based reasoning: A review. The knowledge engineering review , 1994.
- S. Yao, J. Zhao, D. Yu, N. Du, I. Shafran, K. Narasimhan, and Y. Cao. React: Synergizing reasoning and acting in language models. International Conference on Learning Representations , 2023.
- Z. You, Y. Zhang, D. Xu, Y. Lou, Y. Yan, W. Wang, H. Zhang, and Y. Huang. Datawiseagent: A notebookcentric llm agent framework for automated data science. arXiv preprint arXiv:2503.07044 , 2025. · arXiv
- T. Zhang, Z. A. Zhang, Z. Fan, H. Luo, F. Liu, Q. Liu, W. Cao, and L. Jian. Openfe: Automated feature generation with expert-level performance. International Conference on Machine Learning , 2023.
- B. Zoph and Q. V. Le. Neural architecture search with reinforcement learning. International Conference on Learning Representations , 2017.
Cited by
Sorted by citations
MLEvolve: A Self-Evolving Framework for Automated Machine Learning Algorithm Discovery
Du, S., Yan, X., Shi, J., Cao, Z., Feng, S., Liang, Z., Sun, B., Peng, T., Zhou, Y., Li, X., et al. (2026) · arXiv · doi:10.48550/arxiv.2606.06473
PACE: Unleashing the Power of Code Embeddings to Boost AutoML Agents
Zhao, G., Liang, H., Tang, H., Ma, Y., Liu, J., Du, Z., Zheng, Y., Xiao, C., & Hao, J. (2026) · ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD) · doi:10.1145/3770855.3817742
- Authors
- Nam, J., Yoon, J., Chen, J., Shin, J., Arık, S. Ö., & Pfister, T.
- Affiliations
- Korea Advanced Institute of Science and Technology (KR)
- Google (United States) (US)
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2025
- Pages
- 129290-129310
- ISSN
- 1049-5258 (ISSN-L)
- arXiv
- 2506.15692
- OpenAlex
- W4417536788
- Cited by
- 1
- References
- 0
- Citations by year
- 2026: 1
- FWCI
- 1.14
- Type
- conference paper
- Language
- en
- Open access
- green
- arxiv.org
- Locations
- neural information processing systems (publishedVersion)
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University) (cc-by)
- Topics
- Physical Sciences › Computer Science › Information Systems › Software Engineering Research
- Physical Sciences › Computer Science › Artificial Intelligence › Machine Learning and Data Classification
- Physical Sciences › Computer Science › Artificial Intelligence › Artificial Intelligence Applications
- Keywords
- large language models, code generation, search engine retrieval, ablation study, ensembling, feature engineering, Kaggle competition
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
