Towards Reliable ML Feature Engineering via Planning in Constrained-Topology of LLM Agents

TL;DR

The paper tackles reliability and adaptability challenges in automated ML feature engineering by introducing a planner-guided, constrained-topology multi-agent framework that orchestrates repo-level code generation. An LLM-powered planner manages a graph-based environment of specialized actors, leveraging downstream feedback and human-in-the-loop interventions to produce reliable, maintainable code. Empirical results on a novel PySpark-based benchmark show strong gains over manual and unplanned workflows, with a mean pass@3 of $0.833$ compared to $0.600$ and $0.333$ for baselines, respectively, and real-world impact of reducing feature engineering cycles from $3$ weeks to $1$ day. The work demonstrates the value of structured, planner-driven orchestration for scalable, production-ready ML feature pipelines with robust error recovery.

Abstract

Recent advances in code generation models have unlocked unprecedented opportunities for automating feature engineering, yet their adoption in real-world ML teams remains constrained by critical challenges: (i) the scarcity of datasets capturing the iterative and complex coding processes of production-level feature engineering, (ii) limited integration and personalization of widely used coding agents, such as CoPilot and Devin, with a team's unique tools, codebases, workflows, and practices, and (iii) suboptimal human-AI collaboration due to poorly timed or insufficient feedback. We address these challenges with a planner-guided, constrained-topology multi-agent framework that generates code for repositories in a multi-step fashion. The LLM-powered planner leverages a team's environment, represented as a graph, to orchestrate calls to available agents, generate context-aware prompts, and use downstream failures to retroactively correct upstream artifacts. It can request human intervention at critical steps, ensuring generated code is reliable, maintainable, and aligned with team expectations. On a novel in-house dataset, our approach achieves 38% and 150% improvement in the evaluation metric over manually crafted and unplanned workflows respectively. In practice, when building features for recommendation models serving over 120 million users, our approach has delivered real-world impact by reducing feature engineering cycles from three weeks to a single day.

Figures (2)

• Figure 1: Overview of our approach, illustrating interactions between the planner, actors, and executor. Actors and their dependencies match the actual the development environment $E$ used in experiments.
• Figure 2: Actor-level success and failure counts across all tasks, showing which actors had higher failures during execution.