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

AI summary

Overview

Research area: Automated Machine Learning (AutoML), specifically the coupling of LLM-driven feature engineering with Bayesian hyperparameter optimization for tabular/supervised learning pipelines.

Technical level: Advanced. The paper assumes familiarity with Bayesian optimization, Monte Carlo Tree Search, Upper Confidence Bound bandit policies, and LLM prompting/reasoning techniques.

Scope: The paper proposes CoFEH, a framework that interleaves LLM-based free-form feature engineering with Bayesian hyperparameter optimization through a mutual conditioning mechanism and an adaptive budget selector, and evaluates it against traditional and LLM-based baselines on public datasets. Authors are from Peking University (School of CS & Key Lab of High Confidence Software Technologies, MOE; and Beijing Key Laboratory of Software and Hardware Cooperative Artificial Intelligence Systems). The paper is published at KDD '26 (Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2, August 09–13, 2026, Jeju Island, Republic of Korea), DOI 10.1145/3770855.3817664, under a CC BY 4.0 license. Source code is listed at https://doi.org/10.5281/zenodo.20323800 and an extended version at https://arxiv.org/pdf/2602.09851.

What This Paper Is About

Feature engineering (FE) is the step that turns raw data into the inputs a machine learning model sees, and hyperparameter optimization (HPO) is the step that tunes the model itself. Traditional AutoML systems automate both but only inside rigid, pre-defined search spaces, while newer LLM-based methods can invent richer feature operations but usually handle only one sub-task (most often feature generation) and are bolted onto HPO in a greedy "FE-then-HPO" sequence that ignores how strongly features and model settings depend on each other. CoFEH's goal is to build free-form, LLM-generated feature pipelines and jointly optimize them with Bayesian HPO in an interleaved, mutually informed way.

Key Contributions

1. Free-form FE search (addresses challenge C1). An LLM-driven Tree-of-Thought FE pipeline optimizer, implemented via Monte Carlo Tree Search, that explores unconstrained pipeline topologies and operators rather than a fixed template or closed operation library.

2. Mutual conditioning (addresses challenge C2). A bidirectional information-flow mechanism between the LLM and the Bayesian optimizer: the BO surrogate is trained on dataset-state meta-features paired with configurations, and the MCTS FE search uses HPO-discovered performance ceilings and inherits the best configuration when expanding new feature nodes.

3. Dynamic optimizer selector (addresses challenge C3). A PUCB (Predictor Upper Confidence Bound)-based Multi-Armed Bandit policy that adaptively allocates the optimization budget between FE and HPO, with a linear prior-weight schedule and a "Budget Equilibrium" theorem guaranteeing balanced allocation under a neutral reward signal.

4. Empirical evaluation. Experiments on 28 public datasets comparing CoFEH against traditional and LLM-based baselines in both standalone FE and joint FE+HPO settings.

Main Findings

• Two prior conclusions motivate the design: The authors conclude (Conclusion #1) that FE is intrinsically semantic and therefore LLMs are its natural architect, but that this power requires truly free-form pipeline topologies; and (Conclusion #2) that BO remains the gold standard for HPO, surpassing what LLMs currently achieve in hyperparameter spaces.

• The "capability-integration paradox" is identified as the central gap. LLMs excel at FE exploration but lag traditional methods in FE–HPO coupling, because LLM and BO optimizers are heterogeneous and operate over differently represented spaces. Existing LLM-based FE methods such as CAAFE and ELLM-FT therefore fall into a greedy sequential routine: optimize FE under a fixed learner, then tune hyperparameters on the frozen features. The paper argues this causes greedy myopia, discarding features that would only pay off with specific hyperparameter tuning.

• Traditional AutoML suffers three stated limitations. Semantic agnosticism, rigid pipeline topology (e.g., MindWare's strict four-stage workflow; OpenFE's generation-selection pipeline), and a restricted operation set limited to a closed library of predefined algorithms or primitives such as +, −, ×.

• CoFEH's MCTS uses a specific, repurposed reward. Backpropagation tracks visit count N_s, cumulative reward R_s, and subtree-best performance v_max_s. The reward is a binary indicator r = I(v_new > v_max_s0) of whether a new global optimum was achieved, and the UCT exploitation term is Q(s') = R_s'/N_s' + normalized v_max_s'. Expansion is steered by three directives — Initialization, Exploration, Exploitation — with the root treated as non-fully expanded until Init has run five times, and other nodes until Exploration and Exploitation have each run twice.

• A global memory with Pareto-based retrieval drives exploitation. Every historical operation is archived as a tuple (reasoning chain, required feature set, performance score, relative improvement Δv = v − v_base). Candidates are first filtered by functional compatibility with the current dataset state, then a Pareto frontier is built in the (v, Δv) objective space to select "elite experiences" that span both breakthrough transformations and marginal refinements.

• BO is adapted to a dynamic, non-continuous search space. Instead of static-data HPO, the surrogate is trained on pairs ([φ(s_i), λ_i,j], v_i,j) over all nodes in the MCTS tree, using meta-features φ(s) to map discrete tree nodes into a representative feature space. The surrogate is a Random Forest, chosen for categorical-variable support and computational efficiency, with an acquisition function such as Expected Improvement. Candidate configurations are proposed via local perturbation around each node's historical configurations plus global random sampling across Λ.

• Budget allocation is governed by a stated equilibrium theorem. Theorem 3.1 (Budget Equilibrium) states that under a neutral reward signal (Q(a) = const), if the initial bias satisfies 0.5 ≤ p₁ < (M+1.5)/(M+3), the linear scheduling of ω_FE and ω_HPO guarantees N_FE(M), N_HPO(M) ∈ {⌊M/2⌋, ⌈M/2⌉}. The prior weights start biased toward FE (learnability warmup, to make data "model-ready" before HPO) and converge to 0.5 at m = M.

• Reported benchmark setup. Table 1 compares Mindware, OpenFE, OCTree, ELLM, LFG, and CoFEH in standalone FE and joint FE+HPO scenarios, reporting test error (%) for classification and test MSE for regression across 3 runs as Mean ± Std, with best and second-best results bolded/underlined and an Imp. (%) column quantifying relative error reduction of joint FE+HPO versus standalone FE. The paper's stated headline result is that CoFEH outperforms both traditional and LLM-based baselines in both standalone FE and joint FE+HPO settings across 28 public datasets.

• Example from the visible portion of the results table (electricity, classification): standalone FE test error of 9.08 ± 0.13 (Mindware), 8.47 ± 0.07 (OpenFE), 8.54 ± 0.10 (OCTree), 7.71 ± 0.21 (ELLM), 8.84 ± 0.24 (LFG), and 7.31 ± 0.35 (CoFEH); under joint FE+HPO, 8.28 ± 0.17 (Mindware), 7.93 ± 0.07 (OpenFE), 7.76 ± 0.31 (OCTree), 6.79 ± 0.48 (ELLM), 7.65 ± 0.26 (LFG), and 7.15 ± 0.14 (CoFEH). This illustrates that CoFEH is not uniformly best on every individual dataset and column.

• Not reported in the provided content: The full results table is truncated in the supplied text, so the aggregate Avg. Rank values per scenario, the Imp. (%) numbers, and the remaining regression/classification rows are not available here. The dataset names visible before truncation include rl, electricity, compass, wine, house_16H, Magic, higgs, jannis, credit, eye_movements, kddCup09, road-safety, bank-marketing, phoneme, covtype, and california — fewer than the stated 28 datasets. No explicit limitations section appears in the provided content.

Methodology in Plain English

The authors treat a feature engineering pipeline as a sequence of decisions rather than a fixed recipe. The dataset at any stage is a "state," and each feature operation (cleaning, encoding, transformation, generation, selection) is an "action" that moves the data to a new state. An LLM proposes these actions as executable code, guided by metadata about the dataset, the sequence of operations applied so far, a memory of previously successful operations, and a directive telling it whether to be creative (Exploration), to refine using past successes (Exploitation), or to seed the search (Initialization).

The search itself is run with Monte Carlo Tree Search, so the system can branch, backtrack, and revisit promising paths instead of committing to one linear chain. Successful operations are stored in a global memory and retrieved for reuse only when their required features exist in the current dataset, and only if they sit on a Pareto frontier balancing absolute performance against relative improvement.

In parallel, a Bayesian optimizer tunes the downstream model. Because the feature space has no continuous representation, the authors describe each dataset state by "meta-features" and feed pairs of (meta-features, configuration) into a Random Forest surrogate that predicts how well any (dataset state, configuration) combination will perform. An acquisition function then picks the next combination to try, simultaneously choosing the model configuration and the dataset state most likely to benefit from it.

The two halves are kept in the loop with each other rather than run one after the other. When the Bayesian optimizer finds a better score at some tree node, that new performance ceiling propagates back up to the node's ancestors so the MCTS prefers that branch in future selection. When the MCTS expands a new feature node, the new dataset inherits the best-known configuration from its parent, so it is evaluated at its best capacity rather than with defaults.

Finally, a bandit-style controller decides at each step whether to spend budget on feature engineering or on hyperparameter optimization. It balances the empirical success rate of each action against an exploration bonus, and it starts biased toward FE and gradually equalizes the two, so the dataset becomes "model-ready" before aggressive tuning begins. The theorem bounds this schedule so that under neutral rewards the budget is split evenly.

Why This Matters

Impact on research. The paper reframes an important tension in AutoML: LLMs are strong at proposing semantically meaningful feature transformations, but Bayesian optimization is stronger at hyperparameter search, and existing systems typically reconcile these by running them sequentially. CoFEH's mutual conditioning and PUCB-based budget allocation offer a concrete alternative, and its Budget Equilibrium theorem gives a formal handle on how to schedule effort between two heterogeneous optimizers. If the reported results hold broadly, this shifts the field away from "LLM generates features, then tune the model" toward genuinely interleaved joint optimization — a direction relevant to both AutoML researchers and the growing body of work on LLMs for tabular data.

Real-world applications:

• Credit and risk scoring: Tabular datasets with strong domain semantics (the paper's credit-related benchmarks) benefit from domain-aware transformations that a closed operation library cannot express, combined with careful tuning of the model's regularization.

• Marketing and customer analytics: Campaign-response and demographic datasets (the bank-marketing benchmark) often contain categorical and interaction-heavy features where LLM-proposed encodings and cross-features matter more than raw model capacity.

• Energy and utility load modeling: Time- and sensor-derived tabular data (the electricity benchmark) where feature construction and model tuning interact strongly.

• General AutoML platforms and MLOps tooling: CoFEH's recipe of interleaving a generative feature proposer with a Bayesian tuner and a budget controller is directly implementable inside an existing AutoML service.

Industry relevance. Any organization running tabular prediction pipelines — banking, insurance, retail, healthcare analytics, ad tech — currently pays either for data scientists to hand-craft features or for automated tools that search a restricted space. A framework that is both free-form and able to jointly tune the model addresses a real cost/performance bottleneck, provided the LLM inference cost and reproducibility concerns are manageable.

Future Directions

• Cost and efficiency of LLM calls. The framework invokes an LLM repeatedly inside an MCTS loop. The paper reports accuracy but the provided content does not report wall-clock time, token cost, or a comparison of LLM usage against baselines. Quantifying and reducing that overhead is a natural next step.

• Generalization beyond the reported dataset suite. The stated benchmark is 28 public datasets, and the visible portion of the results table is dominated by tabular classification and regression benchmarks. Extending to other data modalities, domain-specific tabular problems, or datasets with heavy class imbalance remains open.

• Theoretical analysis of the full loop. Theorem 3.1 covers budget balance only under a neutral reward signal (Q(a) = const). Whether the interleaved FE–HPO process converges, and under what conditions, is not established in the provided content.

• Reduced reliance on the reward definition and directives. The MCTS reward is a binary "new global optimum" indicator, and expansion is controlled by three hand-specified directives with fixed expansion thresholds (five Init executions at the root; two Exploration and two Exploitation executions elsewhere). Whether these are the right choices, and how sensitive results are to them, is not reported in the supplied content.

• Robustness and leakage checks. Because an LLM proposes code that transforms data, a reader must trust that validation-based scoring prevents information leakage. The provided content does not include an explicit leakage analysis, which would be a valuable follow-up.

Target Audience

This paper is most useful to AutoML and automated feature engineering researchers who are already familiar with Bayesian optimization and want to see how LLM-based feature generation can be integrated rather than stacked on top of it. It also serves practitioners building production tabular ML pipelines who want to judge whether LLM-assisted feature engineering plus joint HPO is worth the added inference cost. Secondary audiences include researchers studying LLM agents for data science tasks, MCTS/bandit-application researchers interested in a new scheduling domain, and PhD students looking for an example of a hybrid system that combines symbolic proposal (LLM), tree search (MCTS), and probabilistic modeling (BO) in one framework. Readers unfamiliar with Bayesian optimization, UCT/UCT-style tree search, or bandit policies will find the method section demanding.

Authors’ abstract

Feature Engineering (FE) is pivotal in automated machine learning (AutoML) but remains a bottleneck for traditional methods, which operate within rigid search spaces and lack domain awareness. While Large Language Models (LLMs) offer a promising alternative to generate unbounded operators with semantic reasoning, existing methods focus on isolated subtasks such as feature generation, falling short of free-form FE pipelines. Moreover, they are rarely coupled with hyperparameter optimization (HPO) of the downstream ML model, leading to greedy "FE-then-HPO" workflows that cannot capture strong FE-HPO interactions. In this paper, we present CoFEH, a collaborative framework that interleaves LLM-based FE and Bayesian HPO for robust end-to-end AutoML. CoFEH uses an LLM-driven FE optimizer powered by Tree of Thought (TOT) to explore flexible FE pipelines, a Bayesian optimization (BO) module to solve HPO, and a dynamic optimizer selector that adaptively interleaves FE and HPO steps. Crucially, we introduce a mutual conditioning mechanism that shares context between LLM and BO, enabling mutually informed decisions. Experiments show that CoFEH outperforms both traditional and LLM-based baselines in both standalone FE and joint FE+HPO settings.