Yiwei Cai
Author directory2026
Optimizing Reasoning LLMs for Automated Feature Engineering via Multi-Stage Direct Preference Optimization
Mingzhi Hu | Menghai Pan | Yuzhong Chen | Junpeng Wang | Yingtong Dou | Vineeth Rakesh | Hoang Anh Duy Le | Uday Singh Saini | Xin Dai | Xiran Fan | Yuqing Liu | Yiwei Cai
Findings of the Association for Computational Linguistics: IJCNLP-AACL 2026
Mingzhi Hu | Menghai Pan | Yuzhong Chen | Junpeng Wang | Yingtong Dou | Vineeth Rakesh | Hoang Anh Duy Le | Uday Singh Saini | Xin Dai | Xiran Fan | Yuqing Liu | Yiwei Cai
Findings of the Association for Computational Linguistics: IJCNLP-AACL 2026
Feature engineering is central to tabular machine learning, but useful transformations must be both statistically predictive and semantically valid with respect to feature names, task descriptions, and observed data patterns. Although large language models can exploit such context, prompt-based feature engineering remains unreliable: generated transformations may be redundant, schema-invalid, non-executable, or weakly aligned with downstream performance. We propose MSDPO, a two-stage Direct Preference Optimization framework for automated feature engineering. Proposal-DPO aligns the model toward useful, diverse, and non-redundant feature proposals, while Function-DPO aligns it toward executable feature-construction code using execution and validation feedback. After alignment, the frozen model directly generates transformation functions for unseen tabular tasks. Across 20 classification datasets, MSDPO achieves the best overall mean rank among strong classical and LLM-based feature-engineering baselines. Under a budget-matched comparison inside the same feature-engineering harness, MSDPO also attains a better mean rank than Claude Sonnet 4.6, despite using an 8B open-weight backbone. Moreover, preferences constructed from classification signals transfer to four regression tasks. Ablation results show that the two DPO stages are complementary, improving both feature design and executable implementation.