DAGPipe: Differentiable DAG Learning for Automated Data Preparation

Abstract

Automated data preparation is a critical bottleneck in machine learning on tabular data. Existing methods largely search over a fixed linear sequence of operations, even tools that branch by feature type use hand-specified, rule-based branches that cannot adapt operation composition to individual features. This limits their effectiveness on heterogeneous data. We propose DAGPipe, a novel framework that jointly learns both the topology and parameters of data preparation pipelines in a fully differentiable manner. Our key insight is to formulate automated data preparation as a structure learning problem over Directed Acyclic Graphs (DAGs), where nodes represent preprocessing operations and edges define data flow. DAGPipe couples a NOTEARS-based continuous relaxation of the DAG topology, a feature-aware routing mechanism that specializes operations to each feature's learned statistical profile, and an end-to-end training algorithm with convergence guarantees. Across 12 benchmark datasets with varying heterogeneity levels show that DAGPipe attains the best average rank among state-of-the-art preprocessing methods and exceeds the average baseline on every dataset. The advantage tends to grow with dataset heterogeneity, with larger margins toward the high-heterogeneity datasets.