Agent Stack for ML Training Lifecycle: Generation, Validation, and Optimization

Jaideep Ray; Daniel Tang; Zejun Lin; Haoyue Tang; Lijuan Zhang; Tao Huang; Tommy Li; Aditya Jain

cs.SEOct 30, 2026

Agent Stack for ML Training Lifecycle: Generation, Validation, and Optimization

Authors: Jaideep Ray, Daniel Tang, Zejun Lin, Haoyue Tang, Lijuan Zhang, Tao Huang, Tommy Li, Aditya Jain,

Organizations: LinkedIn

Abstract

Production model development requires repeated architecture design, platform integration, debugging, training, and evaluation. We present an agent stack deployed at LinkedIn with two independently usable agent loops: AutoTorch for model generation and migration, and AutoResearch for model optimization. AutoTorch generates PyTorch models and training pipelines for TensorFlow migration, XGBoost replacement, and new-model development. Executable validation hooks check platform API compliance, trainability, and conversion fidelity on sampled data, guiding iterative repair before engineer review and full-scale distributed training. AutoResearch explores architecture and training-configuration changes through expert-steered or autonomous experimentation. An external orchestrator manages parallel trials, experiment history, candidate selection, and wall-clock budgets across agent sessions. A shared runtime preserves execution state independently of agent reasoning and allocates GPUs separately from agent execution. Across ranking, content understanding, and recommendation workloads, AutoTorch completed generation and full-scale training in 92 of 115 workflow runs (80%). Of these, 77 (83.7%) produced models meeting post-training quality requirements. In a separate cohort, AutoResearch improved the primary task metric on held-out data for 28 of 100 production pipelines. Engineering teams estimate that the stack reduced elapsed end-to-end migration time from approximately eight weeks to two weeks, a 4x speedup.

Figures & tables

Panel of Figure 1 — Figure 1: Model migration as a composition of independently usable agent loops. AutoTorch generates a training model bundle; AutoResearch can subsequently optimize the accepted model. Each loop also supports independent invocation.

Panel of Figure 2 — Figure 2: AutoTorch’s two stages. Stage 1 clones a TF model or builds a new architecture, with sampled checks guiding repair. After engineer review, Stage 2 can use a warm start for a fast task-metric check before final cold-start training.

Panel of Figure 3 — Figure 3: AutoResearch’s synchronous hill climb. A round begins from a fixed incumbent, parallel trials run in separate worktrees, and selection occurs after all results return.

Panel of Figure 4 — Figure 4: Durable agent sandboxes and independently allocated accelerators. Training data is available only through controlled paths and is not placed in LLM context.

Panel of Figure 5 — Figure 5: Skills propose architectures, and executable checks return failure evidence for repair. New-model generation uses separate scoring after full training.

Panel of Figure 6 — Figure 6: Synchronous AutoResearch rounds. Each round waits for all trials, then retains or replaces the incumbent. All rounds share one elapsed-time budget.

Panel of Figure 7 — Figure 7: Hyperparameter and architecture search share the same evaluation loop and round-end selection policy.

1. Deployment outcomes		Result	Scope
1.1	AutoTorch generation runs	115	Workflow executions; not necessarily distinct models
1.2	Operational completion	92 of 115 (80%)	Generation and full training; denominator: all attempted runs
1.3	Post-training quality pass	77 of 92 (83.7%)	Evaluated model artifacts from completed runs; TF clones also require a conversion score $\geq 90$
1.4	Offline primary-metric difference	Typically within 1–5%	Relative to reference after full training; initialization modes not disaggregated; not a deployment gate
1.5	AutoResearch improvements	28 of 100 (28%)	Improved original baseline on final held-out evaluation
2. Measured AutoTorch execution		p50 / p90	Range

Table 1: Production outcomes and runtime. AutoTorch latencies refer to the indicated intervals; end-to-end time includes engineer review. The two denominator cohorts differ.

Case	Model / use case	Agent action	Primary task metric
AutoTorch — generation and migration
AT-1	Two-tower video retrieval	Delivered PyTorch model and training pipeline based on source model	Recall@200 matched source model
AT-2	XGBoost $\rightarrow$ TabM; job-posting classification	Changed loss to per-example binary cross-entropy	AUC matched
AT-3	TensorFlow Ranking-based CTR model	Ported to PyTorch with weight transfer; fixed data and evaluation mismatches	AUC matched
AT-4	TensorFlow multi-task model (200M embedding parameters)	Ported to PyTorch; repaired preprocessing, sparse gradients, and task-output handling	Embedding outputs matched (output-parity check)
AT-5	TensorFlow graph neural network	Ported to a PyTorch GNN; restored node encoder and concatenation combiner	Downstream classifier AUC matched

Table 2: Selected agent-loop cases (AT: AutoTorch; AR: AutoResearch). These examples are not randomly sampled.

Panel of Figure 8 — Figure 8: Harness-specific execution control. AutoTorch uses hooks to gate each generate–evaluate–repair iteration. AutoResearch maintains persistent state and round-end selection externally to agent sessions.

Automated Machine Learning

Automated machine learning (AutoML) builds and tunes data-to-model pipelines with minimal manual effort, balancing predictive quality, compute, and deployment constraints.

See all 99 papers

Think it, Run it: Autonomous ML pipeline generation via self-healing multi-agent AI
Apr 29, 2026Adela Bara, Gabriela Dobrita, Simona-Vasilica OpreaAutomated Machine Learning Multi-Agent LLMs
Bucharest University of Economic Studies, Department of Economic Informatics and Cybernetics, no.6 Piața Romană, Bucharest, 010374, Romania
AgentOpt v0.1 Technical Report: Client-Side Optimization for LLM-Based Agent
Apr 7, 2026Wenyue Hua, Sripad Karne, Qian Xie +4LLM Agent Workflow Optimization LLM Agent Efficiency
Microsoft Research, AI Frontiers · Columbia University · Cornell University
AutoTrainess: Teaching Language Models to Improve Language Models Autonomously
Jun 30, 2026Zhaojian Yu, Penghao Yin, Shuzheng Gao +3Tool-Augmented Language Model Agents LLM Tool Use
1Tsinghua University · 2The Chinese University of Hong Kong · 3Simple Agent Lab

Apr 29, 2026J/K move · Enter open · S save

Think it, Run it: Autonomous ML pipeline generation via self-healing multi-agent AI

Adela Bara, Gabriela Dobrita, Simona-Vasilica Oprea

Bucharest University of Economic Studies, Department of Economic Informatics and Cybernetics, no.6 Piața Romană, Bucharest, 010374, Romania

Automated Machine Learning Multi-Agent LLMs

Open paper AskPDF arXiv

The purpose of our paper is to develop a unified multi-agent architecture that automates end-to-end machine learning (ML) pipeline generation from datasets and natural-language (NL) goals, improving efficiency, robustness and explainability. A five-agent system is proposed to handle profiling, intent parsing, microservice recommendation, Directed Acyclic Graph (DAG) construction and execution. It integrates code-grounded Retrieval-Augmented Generation (RAG) for microservice understanding, an explainable hybrid recommender combining multiple criteria, a self-healing mechanism using Large Language Model (LLM)-based error interpretation and adaptive learning from execution history. The approach is evaluated on 150 ML tasks across diverse scenarios. The system achieves an 84.7% end-to-end pipeline success rate, outperforming baseline methods. It demonstrates improved robustness through self-healing and reduces workflow development time compared to manual construction. The study introduces a novel integration of code-grounded RAG, explainable recommendation, self-healing execution and adaptive learning within a single architecture, showing that tightly coupled intelligent components can outperform isolated solutions.

Abstract

Automated Machine Learning

Explore similar work

Think it, Run it: Autonomous ML pipeline generation via self-healing multi-agent AI

AgentOpt v0.1 Technical Report: Client-Side Optimization for LLM-Based Agent

AutoTrainess: Teaching Language Models to Improve Language Models Autonomously