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

TL;DR

This paper presents a planner-guided multi-agent framework that improves ML feature engineering by integrating environment-aware planning, human feedback, and failure correction, significantly enhancing efficiency and reliability.

Contribution

It introduces a novel constrained-topology multi-agent system guided by a planner that orchestrates code generation for feature engineering, addressing key real-world challenges.

Findings

Achieved 38% and 150% improvements on a novel dataset.

Reduced feature engineering cycles from three weeks to one day.

Enhanced reliability and maintainability of generated code.

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.