🤖 machine learning

CoFEH: LLM-driven Feature Engineering Empowered by Collaborative Bayesian Hyperparameter Optimization

This paper introduces CoFEH, a collaborative framework that synergistically interleaves LLM-driven feature engineering with Bayesian hyperparameter optimization through a mutual conditioning mechanism and dynamic step selection to overcome the limitations of rigid search spaces and greedy workflows in traditional AutoML.

Original authors: Beicheng Xu, Keyao Ding, Wei Liu, Yupeng Lu, Bin Cui

Published 2026-05-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Beicheng Xu, Keyao Ding, Wei Liu, Yupeng Lu, Bin Cui

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Building a Better Machine Learning Machine

Imagine you are trying to build the ultimate race car. You have two main jobs:

Preparing the Fuel (Feature Engineering): You need to refine the raw fuel (data) so the engine can burn it efficiently. This involves filtering out impurities, mixing in special additives, or changing the fuel's chemical structure.
Tuning the Engine (Hyperparameter Optimization): You need to adjust the engine's settings (spark plug timing, fuel injection rate) to get the most speed out of that fuel.

The Old Way:
Traditionally, engineers did these jobs separately and in a rigid order.

First, they would mix the fuel using a fixed recipe book (a limited set of pre-defined chemical reactions). They couldn't invent new additives; they could only mix what was already in the book.
Once the fuel was "done," they would lock the recipe and start tuning the engine.
The Flaw: This is like tuning a Ferrari engine with low-grade gasoline. If the fuel is bad, no amount of engine tuning will make the car fast. Conversely, if you have amazing fuel but a poorly tuned engine, you still lose the race. The old methods didn't realize that the best fuel depends on the best engine settings, and vice versa.

The New Solution: CoFEH

The authors of this paper created CoFEH (Collaborative Feature Engineering and Hyperparameter Optimization). Think of CoFEH as a super-smart, two-person pit crew working together in real-time.

1. The "Creative Chef" (The LLM)

Instead of using a rigid recipe book, CoFEH uses a Large Language Model (LLM) as a "Creative Chef."

What it does: This Chef doesn't just mix ingredients from a list. It reads the data, understands the "flavor" of the problem, and invents new ingredients and cooking techniques on the fly.
The "Tree of Thought" Analogy: Imagine the Chef is exploring a massive, infinite forest of possible recipes. Instead of walking a single path, the Chef uses a "Tree of Thought" strategy. They branch out, try a few paths, realize one leads to a dead end, backtrack, and try a completely different direction. This allows them to find truly unique and powerful data transformations that rigid systems would never think of.

2. The "Precision Mechanic" (Bayesian Optimization)

While the Chef is inventing new fuel, a "Precision Mechanic" (using a method called Bayesian Optimization) is tuning the engine.

The Old Problem: Usually, the Chef makes the fuel, hands it to the Mechanic, and says, "Here, tune this." The Mechanic tries to tune it, but if the fuel is weird, the Mechanic might give up or tune it poorly.
The CoFEH Fix: The Chef and Mechanic are now talking to each other constantly.
- The Mechanic tells the Chef: "Hey, this new fuel additive you made works great with this specific engine setting, but fails with that one."
- The Chef tells the Mechanic: "Okay, I'll tweak the recipe to match your engine settings better."
- They work in a loop, refining the fuel and the engine settings simultaneously. This is called Mutual Conditioning.

3. The "Team Captain" (Dynamic Selector)

The team has a limited amount of time (budget) to work on the car. Who should they focus on right now?

The Old Way: They would spend 50% of the time on fuel and 50% on the engine, no matter what.
The CoFEH Fix: A "Team Captain" watches the progress.
- If the car is struggling because the fuel is terrible, the Captain shouts, "Chef, keep cooking! Ignore the engine for a moment!"
- If the fuel is perfect but the engine is sputtering, the Captain yells, "Mechanic, fine-tune the engine! The Chef can take a break."
- This ensures they spend their time exactly where it matters most for that specific race.

Why is this a Big Deal?

The paper tested this system on 28 different datasets (think of these as 28 different types of race tracks with different rules).

The Result: CoFEH beat all the other methods, including the old rigid systems and other AI systems that use LLMs.
The Secret Sauce:
1. Freedom: It isn't stuck using a fixed list of operations. It can invent new ones.
2. Collaboration: It doesn't do the fuel and engine jobs separately. It does them together, so they help each other.
3. Adaptability: It knows when to focus on data and when to focus on the model, saving time and energy.

In Summary

Imagine you are trying to solve a puzzle.

Old AI: Tries to force puzzle pieces into a box using a hammer (rigid rules).
Other New AI: Uses a smart robot to find pieces, but the robot only looks for pieces one at a time, ignoring how they fit together.
CoFEH: Uses a team of experts. One expert (the LLM) is a master puzzle-solver who can invent new pieces if needed. The other expert (the Optimizer) knows exactly how to arrange them. They stand side-by-side, constantly saying, "If we move this piece here, it fits better with that one," until the puzzle is perfectly solved.

The paper claims that by letting these two experts collaborate and share information in real-time, they can solve complex data problems much better and faster than any previous method.

Technical Summary: CoFEH

Problem Statement

Feature Engineering (FE) is a cornerstone of machine learning pipelines but remains a bottleneck for traditional automated methods. Existing approaches suffer from three fundamental limitations:

Semantic Agnosticism: They rely on rigid, predefined search spaces and trial-and-error without leveraging domain priors.
Rigid Topology: They enforce fixed pipeline sequences (e.g., generation followed by selection), preventing dynamic interleaving or flexible composition of heterogeneous operations.
Restricted Operation Sets: They are confined to closed libraries of built-in algorithms or basic arithmetic, unable to discover task-specific operations beyond predefined primitives.

Furthermore, while Large Language Models (LLMs) offer a promising alternative for generating semantically grounded, unbounded operators, current LLM-based FE methods typically focus on isolated subtasks (like feature generation) and fail to construct free-form pipelines. Crucially, these methods rarely couple FE with Hyperparameter Optimization (HPO) of the downstream model. This leads to a "greedy" sequential workflow (FE-then-HPO) that ignores the strong dependency between feature representation and model capacity, often trapping the system in sub-optimal local minima.

Methodology: CoFEH Framework

The authors propose CoFEH (Collaborative Feature Engineering and Hyperparameter Optimization), a framework designed to interleave LLM-driven FE and Bayesian HPO for robust end-to-end AutoML. The system comprises three core components:

1. LLM-Driven Tree-of-Thought (ToT) FE Optimizer

To address the challenge of Free-form FE, CoFEH reformulates pipeline construction as a sequential decision-making problem solved via Monte Carlo Tree Search (MCTS).

State Representation: The dataset at any transformation stage is treated as a state node in a search tree.
Steerable Reasoning: The LLM generates new FE operations based on a structured prompt containing dataset metadata, the ancestor pipeline, and a dynamic directive ( $d$ $d$ ):
- Initialization: Generates high-quality initial operations for the root node.
- Exploration: Proposes novel operations to probe unknown regions of the search space.
- Exploitation: Refines the current state by reusing and adapting high-performing historical operations.
Global Memory: A Pareto-based retrieval mechanism archives successful operations (reasoning chains, feature subsets, scores, and relative gains) to guide the LLM during the Exploitation phase.

2. Mutual Conditioning Mechanism

To solve the Collaborative FE–HPO challenge, CoFEH establishes a bidirectional information flow between the LLM (FE) and Bayesian Optimization (BO) (HPO), avoiding isolated tuning.

HPO Conditioned on FE: The BO surrogate model uses meta-features ( $\phi(s)$ ) extracted from the dataset state to characterize the FE pipeline. This allows the BO to model the joint influence of the pipeline and hyperparameters, recommending the optimal configuration ( $\lambda^*$ ) for a specific dataset state ( $s^*$ ).
FE Conditioned on HPO: The MCTS selection and playout phases are informed by HPO results. The "best performance" metric for a node is determined by the optimal hyperparameters found for that specific dataset state, ensuring the LLM evaluates feature quality under optimal model configurations rather than default settings.

3. Dynamic Optimizer Selector

To address Task-adaptive scheduling, CoFEH models the allocation of the optimization budget between FE and HPO as a Multi-Armed Bandit (MAB) problem using a Predictor Upper Confidence Bound (PUCB) policy.

The selector dynamically decides whether to execute an FE step or an HPO step at each iteration based on empirical success rates ( $Q(a)$ ) and time-varying prior weights ( $\omega_a$ ).
Prior Weight Scheduling: The system starts with a bias toward FE (to ensure data is "model-ready") and linearly decays this bias toward a 50/50 equilibrium as the budget is exhausted. This guarantees a balanced budget distribution in neutral scenarios while allowing task-specific empirical rewards to steer the allocation in practice.

Key Contributions

Free-Form Pipeline Optimization: An LLM-driven ToT optimizer that explores unconstrained pipeline topologies and operators, enabling truly free-form feature engineering that composes preprocessing, transformation, generation, and selection dynamically.
Mutual Conditioning Mechanism: A novel bidirectional coupling where FE and HPO inform each other's decisions, preventing the "greedy myopia" of sequential optimization and enabling the discovery of global optima in the joint space.
Dynamic Resource Allocation: A PUCB-based selector that adaptively allocates computational budget between FE and HPO based on the specific utility of each for the given task.
Empirical Validation: Extensive experiments demonstrating that CoFEH outperforms both traditional AutoML systems and existing LLM-based baselines.

Experimental Results

The authors evaluated CoFEH on 28 public datasets (19 classification, 9 regression) across three downstream model scenarios: XGBoost, MLP, and CASH (Combined Algorithm Selection and Hyperparameter Optimization).

Standalone FE: CoFEH achieved the best average rank (1.82) among all baselines, significantly outperforming traditional methods (Mindware, OpenFE) and other LLM-based approaches (OCTree, ELLM, LFG).
Joint FE+HPO: In the end-to-end setting, CoFEH demonstrated a 7.03% relative error reduction compared to its own standalone FE performance, outperforming all baselines. It achieved an average rank of 1.75, establishing a new state-of-the-art.
Generalization: The framework showed robust generalization across different model types (MLP, CASH), maintaining top performance and achieving significant error reductions (e.g., +45.1% improvement in CASH scenarios).
Cost Efficiency: CoFEH achieved the best cost-performance trade-off, requiring only ~$0.07 per task with the Gemini-2.0-flash backbone, while traditional search-heavy methods (OpenFE, Mindware) required significantly more runtime.
Ablation Studies: Removing the mutual conditioning mechanism or the dynamic selector significantly degraded performance, confirming the necessity of both coupled information management and adaptive resource scheduling.

Significance and Claims

The paper claims that CoFEH bridges the gap between the semantic reasoning capabilities of LLMs and the numerical efficiency of Bayesian Optimization. By moving away from rigid, sequential workflows to an interleaved, collaborative optimization strategy, CoFEH overcomes the "capability-integration paradox" where LLMs excel at FE but lag in FE-HPO coupling.

The authors assert that CoFEH:

Unleashes the full potential of LLMs to create human-expert-like, free-form pipelines.
Solves the problem of sub-optimal local minima caused by decoupled FE and HPO.
Provides a scalable, extensible foundation for future AutoML work that can be adapted to image and text data via prompt modifications.

The work is presented as a significant step toward holistic ML pipeline optimization, where feature discovery and model tuning are treated as a unified, mutually reinforcing process rather than isolated stages.