MLZero: A Multi-Agent System for End-to-end Machine Learning Automation

Ying Nian Wu (University of California, Los Angeles) · George Karypis (University of Minnesota, Minneapolis) · Haoyang Fang (AWS) · Boran Han (AWS) · Nick Erickson (Amazon Web Services) · Xiyuan Zhang (AWS AI) · Su Zhou (Carnegie Mellon University) · Anirudh Dagar (AWS) · Jiani Zhang (Google) · Caner Turkmen (Amazon Web Services) · Tony Hu (AWS AI) · Huzefa Rangwala (George Mason University) · Yuyang (Bernie) Wang (AWS AI)
automlcognitive perception modulecompact llmepisodic memoryiterative code generationmle-bench litemulti-agent frameworkmultimodal automl agent benchmarkmultimodal dataperceptual contextperformance evaluationsemantic memorysolution qualitysuccess rate

Existing AutoML systems have advanced the automation of machine learning (ML); however, they still require substantial manual configuration and expert input, particularly when handling multimodal data. We introduce MLZero, a novel multi-agent framework powered by Large Language Models (LLMs) that enables end-to-end ML automation across diverse data modalities with minimal human intervention. A cognitive perception module is first employed, transforming raw multimodal inputs into perceptual context that effectively guides the subsequent workflow. To address key limitations of LLMs, such as hallucinated code generation and outdated API knowledge, we enhance the iterative code generation process with semantic and episodic memory. MLZero demonstrates superior performance on MLE-Bench Lite, outperforming all competitors in both success rate and solution quality, securing six gold medals. Furthermore, when evaluated on our Multimodal AutoML Agent Benchmark, which includes 25 more challenging tasks spanning diverse data modalities, MLZero outperforms the competing methods by a large margin with a success rate of 0.92 (+263.6\%) and an average rank of 2.28. Our approach maintains its robust effectiveness even with a compact 8B LLM, outperforming full-size systems from existing solutions.