Evolving Demonstration Optimization for Chain-of-Thought Feature Transformation

TL;DR

This paper introduces a dynamic, reinforcement learning-based framework for optimizing feature transformations in data-centric AI, leveraging evolving demonstrations to enhance the diversity, validity, and alignment of transformations with downstream tasks.

Contribution

It presents a novel closed-loop approach that evolves demonstration contexts for LLM-driven feature transformation, improving over static methods and enhancing performance and stability.

Findings

Outperforms classical and LLM-based baselines on tabular benchmarks.

More stable and diverse transformation generation compared to one-shot methods.

Framework generalizes across different LLMs and downstream evaluators.

Abstract

Feature Transformation (FT) is a core data-centric AI task that improves feature space quality to advance downstream predictive performance. However, discovering effective transformations remains challenging due to the large space of feature-operator combinations. Existing solutions rely on discrete search or latent generation, but they are frequently limited by sample inefficiency, invalid candidates, and redundant generations with limited coverage. Large Language Models (LLMs) offer strong priors for producing valid transformations, but current LLM-based FT methods typically rely on static demonstrations, resulting in limited diversity, redundant outputs, and weak alignment with downstream objectives. We propose a framework that optimizes context data for LLM-driven FT by evolving trajectory-level experiences in a closed loop. Starting from high-performing feature transportation sequences explored by reinforcement learning, we construct and continuously update an experience library of downstream task-verified transformation trajectories, and use a diversity-aware selector to form contexts along with a chain-of-thought and guide transformed feature generation toward higher performance. Experiments on diverse tabular benchmarks show that our method outperforms classical and LLM-based baselines and is more stable than one-shot generation. The framework generalizes across API-based and open-source LLMs and remains robust across downstream evaluators.