PaperlayerView exact revision on arXiv

Source-linked AI summary

Large Language Models for Automated Data Science: Introducing CAAFE for Context-Aware Automated Feature Engineering

Noah Hollmann, Samuel Müller, Frank Hutter

arXiv:2305.03403v5cs.AIcs.LG

TL;DR

AutoML leaves data engineering and domain-knowledge integration largely to human practitioners. CAAFE uses an LLM to generate semantically meaningful features and explanations for tabular data; with GPT-4, it raised average ROC AUC from 0.798 to 0.822 and improved performance on 11/14 datasets.

  • Problem

    Existing AutoML systems largely leave data engineering and domain-knowledge integration to human practitioners, despite data engineering and cleaning being among data scientists’ most time-consuming tasks.

    Source: §1 · p. 2

  • Method

    CAAFE uses an LLM to iteratively generate Python code for semantically meaningful tabular features from dataset context, evaluates feature alterations on validation data, and explains each feature’s utility.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    With GPT-4, CAAFE raised average ROC AUC from 0.798 to 0.822 and improved performance on 11/14 datasets.

    Source: §5 · p. 7

  • Takeaways & Limitations

    The evaluation supports CAAFE as a complement to existing automated feature engineering and AutoML methods.

    Source: §6 · p. 9

    CAAFE can generate features and explanations that appear significant and logical despite being inaccurate or invented, a risk for essential decision-making or research.

    Source: §6 · p. 9

Abstract

from arXiv · show

As the field of automated machine learning (AutoML) advances, it becomes increasingly important to incorporate domain knowledge into these systems. We present an approach for doing so by harnessing the power of large language models (LLMs). Specifically, we introduce Context-Aware Automated Feature Engineering (CAAFE), a feature engineering method for tabular datasets that utilizes an LLM to iteratively generate additional semantically meaningful features for tabular datasets based on the description of the dataset. The method produces both Python code for creating new features and explanations for the utility of the generated features. Despite being methodologically simple, CAAFE improves performance on 11 out of 14 datasets -- boosting mean ROC AUC performance from 0.798 to 0.822 across all dataset - similar to the improvement achieved by using a random forest instead of logistic regression on our datasets. Furthermore, CAAFE is interpretable by providing a textual explanation for each generated feature. CAAFE paves the way for more extensive semi-automation in data science tasks and emphasizes the significance of context-aware solutions that can extend the scope of AutoML systems to semantic AutoML. We release our $\href{https://github.com/automl/CAAFE}{code}$, a simple $\href{https://colab.research.google.com/drive/1mCA8xOAJZ4MaB_alZvyARTMjhl6RZf0a}{demo}$ and a $\href{https://pypi.org/project/caafe/}{python\ package}$.

1 Introduction

The introduction argues that AutoML leaves much data engineering and domain-knowledge integration to human practitioners, and presents CAAFE as a way to use LLMs for context-aware feature engineering with classical ML.

1 Introduction

  • Model selection, training, and scoring account for roughly 23% of data scientists’ time, while data engineering and cleaning receive limited AutoML support.Source: §1 · p. 2
  • LLMs offer contextual domain knowledge for data science tasks, but lack interpretability and verifiability, behave inconsistently, and can fail at simple calculations (Hendrycks et al., 2021; OpenAI Community, 2021).Source: §1 · p. 2
  • CAAFE uses LLM-generated Python code to iteratively add semantically meaningful features with algorithmic feedback, and supplies comments explaining each feature’s utility.Classical algorithms process the modified datasets, combining their robustness and predictability with LLM domain knowledge and creativity.Source: §1 · p. 2
  • Automating domain-knowledge integration is intended to reduce model-development latency and cost, evaluate a more informed solution space, and improve robustness and reproducibility.The introduction presents CAAFE as demonstrating the potential of LLMs to broaden AutoML toward data science tasks.Source: §1 · p. 2

2 Background

The background presents prompt-based LLMs as a source of context-sensitive task adaptation, while noting limitations of direct LLM prediction on tabular data. It motivates feature engineering as valuable for predictive performance and interpretability, especially when semantic context can guide the search for useful features.

2.1 Large Language Models (LLMs)

  • LLMs are pretrained next-word predictors; prompting enables task-specific in-context inference with frozen parameters, and GPT-4 is described as achieving state-of-the-art performance on tasks including coding.Source: §2.1 · p. 2
  • Hegselmann et al. (2023)'s textual tabular predictor scales quadratically with N·M, is difficult to interpret and consistency-assure, and performed best only on datasets with up to 8 samples.The method encodes the training dataset as text and predicts by continuing an unseen sample's description with the target column.Source: §2.1 · pp. 2–3
  • Narayan et al. (2022) and Vos et al. (2022) prompt per row for data wrangling, whereas CAAFE uses code so one query applies across samples; Figure 2 assigns model building to classical AutoML.Source: §2.1 · p. 3, Figure 2 · p. 3

2.2 Feature Engineering

  • Feature engineering constructs features from raw inputs to improve predictive performance; it remains useful for limited data, complex patterns, or interpretability, using transformations, encoding, clustering, aggregation, or dimensionality reduction.Source: §2.2 · p. 3
  • Automated feature-engineering methods include enumerative, heuristic, beam-search, learning-based, and reinforcement-learning approaches, but none automatically harness semantic information.The cited examples include DFS (Kanter & Veeramachaneni (2015)), Cognito, AutoFeat, and learning-based or reinforcement-learning methods.Source: §2.2 · p. 3

2.2.1 Incorporating Semantic Information

  • Semantic context can narrow the vast feature search to meaningful candidates, improving performance and interpretability while mitigating multiple-testing and computational-complexity issues.Source: §2.2.1 · p. 4

3 Method

CAAFE incorporates dataset context into iterative LLM-generated feature engineering, using validation performance to decide which generated changes to keep.

3 Method

  • CAAFE executes LLM-generated feature code on training and validation data, retaining a change only when its validation performance exceeds the baseline.Otherwise, both datasets remain unchanged.Source: §3 · p. 4

Prompting LLMs for Feature Engineering Code

  • The prompt requests useful features and explanations, removal of unnecessary features, and uses a code-and-explanation template with chain-of-thought steps from utility to implementation.Source: §3 · pp. 4–5
  • The prompt combines a user dataset description, feature names and types, train-split summary statistics including missingness, and 10 random dataset rows.Source: §3 · pp. 4–5

Technical Setup

  • CAAFE evaluates code changes using mean accuracy and ROC AUC changes across 10 random validation splits, keeping them when their average is above 0.The setup uses GPT-4 or GPT-3.5, 10 feature-engineering iterations, and TabPFN for iterative code evaluation.Source: §3 · p. 6

3 Method

  • Generated Python is syntax-parsed and restricted to whitelisted operations, but loops and list comprehensions can still enable infinite loops or excessive resource use.Source: §3 · p. 6

4 Experimental Setup

The experiments compare downstream classifiers and feature-engineering methods on small tabular datasets, using dataset descriptions as context for CAAFE. The setup also considers possible LLM training-data exposure and controls evaluation through shared train-test splits.

Setup of Downstream-Classifiers

  • Logistic Regression, Random Forests, and TabPFN serve as downstream classifiers, while TabPFN evaluates added features; preprocessing follows Hollmann et al. (2022).Missing values are mean-imputed, categorical inputs are one-hot or ordinal encoded, and features are normalized with categorical indicators passed where necessary.Source: §4 · p. 6

Setup of Automated Feature Engineering Methods

  • DFS and AutoFeat are tested alone and after CAAFE, using numeric addition and multiplication for DFS and one feature-engineering step for AutoFeat.DFS features beyond TabPFN’s 100-feature limit are evaluated at the no-feature-engineering performance.Source: §4 · p. 6

Evaluating LLMs on Tabular Data

  • To assess possible LLM benchmark exposure, the evaluation includes pre-September 2021 OpenML datasets and lesser-known post-September 2021 Kaggle datasets harder for web crawlers to access.GPT-4 and GPT-3.5 training data has a September 2021 cutoff, so established benchmark descriptions may have appeared in their training data.Source: §4 · p. 6

Dataset Selection

  • Datasets are small (up to 2,000 samples) and have descriptive feature names; OpenML task descriptions provide CAAFE context, and TabPFN-perfect datasets have reduced training sets.The authors focus on small datasets because feature engineering is most important and significant for them; datasets with numbered feature names are excluded. OpenML sources are cited as Vanschoren et al. (2013) and Feurer et al.Source: §4 · p. 6

Evaluation Protocol

  • Each dataset uses five repetitions with different seeds and 50/50 train-test splits, shared across methods to reduce split-related variance (Bouthillier et al., 2021).Source: §4 · p. 7

5 Results

With GPT-4, CAAFE improves mean ROC AUC and performance across most evaluated datasets; it can also precede classical AutoFE, while gains vary by dataset and description.

Performance of CAAFE

  • With GPT-4, CAAFE raises mean ROC AUC from 0.798 to 0.822 and improves 11/14 datasets, as shown in Table 2.The features generated by CAAFE are selected using TabPFN as the classifier.Source: Table 2 · p. 7, §5 · p. 7

Feature extension examples

  • Examples show CAAFE combines features, bins numeric variables into ordinal versions, transforms strings, removes superfluous features, and recovers from invalid-code errors.The code and comments are generated from user-provided dataset descriptions.Source: Table 3 · p. 8

Performance of CAAFE

  • GPT-4's gain is similar (71%) to the average gain from random forest (AUC 0.783) over logistic regression (AUC 0.749), while GPT-3.5 improves only 6/14 datasets.CAAFE with GPT-4 also improves performance on all of the new Kaggle datasets.Source: §5 · p. 9
  • Gains vary with feature-engineering amenability and description quality; runtime is 4:43 minutes per dataset, and 52 generated features (7.4%) were faulty but recovered.90% of runtime is spent on LLM code generation and 10% on evaluating generated features.Source: §5 · p. 9

Incorporating Classical AutoFE Methods

  • After CAAFE, classical AutoFE further improves logistic regression and random forests but not TabPFN; CAAFE improves averages for every tested classifier and additional AutoFE combination.The paper suggests TabPFN's larger hypothesis space may explain why the evaluated AutoFE methods do not improve it.Source: §5 · p. 9

6 Conclusion · B Broader Impact Statement · B.1 Social Impact of Automation

CAAFE uses large language models to generate semantically meaningful tabular features and explanations, integrating domain knowledge into AutoML. The conclusion notes limitations and future directions, while the social-impact discussion considers workforce displacement and adaptation.

6 Conclusion

  • CAAFE automates tabular feature engineering with semantically meaningful features and explanations, complementing existing feature-engineering and AutoML methods.The approach integrates domain knowledge by leveraging large language models.Source: §6 · p. 9
  • The authors anticipate LLMs could automate other data-science tasks, including data collection, processing, model building, and deployment, and that improved LLMs could increase CAAFE’s effectiveness.The conclusion emphasizes context-aware solutions for robust outcomes.Source: §6 · p. 9
  • Dataset descriptions are critical to CAAFE, and more accurate, detailed user-provided descriptions could improve its effectiveness.The study’s descriptions came solely from web-crawled text associated with public datasets.Source: §6 · p. 9
  • Large feature sets can create prompts that challenge LLMs, feature testing lacks statistical tests, and hallucinations may produce convincing but ungrounded features and explanations.The authors warn that misplaced trust in such outputs could be problematic in essential decision-making or research.Source: §6 · p. 9
  • Future work may explore prompt tuning, model fine-tuning, and other ways to incorporate domain knowledge, as well as human-in-the-loop AutoML.The authors note that CAAFE’s interpretable, modifiable LLM inputs and outputs make continuous expert interaction particularly feasible.Source: §6 · p. 9

B.1 Social Impact of Automation

  • CAAFE may contribute to data-science automation and worker displacement, while automating routine tasks could free data scientists for higher-level problem-solving and decision-making.The authors call for stakeholder awareness and workforce education and adaptation as AI technologies evolve.Source: §B.1 · p. 12

B.2 Replication of Biases … E.1 Semantic Blinding

CAAFE’s performance depends on semantic context, while the paper also discusses bias risks and safeguards, interpretability, capability risk, and its full prompt. Its authors report that concealing dataset descriptions and feature names reduces GPT-4’s average AUROC.

B.2 Replication of Biases

  • Because GPT-4 learns from web-crawled data containing social biases, CAAFE may generate biased features; the authors advise avoiding demographic or potentially discriminatory data, or carefully checking outputs.The paper also identifies potential bias at the LLM, generated-feature, and downstream-classifier levels, while noting that the classifier is not directly affected by CAAFE.Source: §B.2 · p. 12
  • CAAFE asks the LLM to explain feature usefulness to help users identify fairness-related harms and requires explicit user confirmation before generated code runs.The authors describe user confirmation as a way to keep the LLM in a decision-support role rather than allowing it to make final decisions.Source: §B.2 · p. 12
  • Although GPT-4 was selected with bias mitigation in mind, generated features are retained only when they improve cross-validation, and aligned LLM and dataset biases may still pass this filter.The authors present performance-based feature checking as a partial defense against biases present in the LLM but absent from the data, not as a guarantee against aligned biases.Source: §B.2 · pp. 12–13
  • In a deliberately biased doctor-or-nurse sample, GPT-4 proposed an `end_with_a` feature and said names ending in “a” might be more common in females, making its gendered reasoning visible.The sample associated male-associated names with doctors and used only female names ending in “a”; the authors say the usefulness comment made the bias easy to detect.Source: §B.2 · p. 13

B.3 AI Model Interpretability

  • CAAFE aims to improve interpretability by explaining model outputs and generating simple code, making automated feature engineering more transparent.Source: §B.3 · p. 13

B.4 Risk of increasing AI capabilities

  • The authors estimate that this work does not increase AI-capability risk because it demonstrates an application of LLMs rather than affecting their general capabilities.Source: §B.4 · p. 13

D Full LLM Prompt

E.1 Semantic Blinding

  • Concealing dataset descriptions and feature names drops GPT-4’s average AUROC from 0.822 to 0.8 across datasets.Table 4 compares CAAFE with and without semantic blinding; the passage reports the average values.Source: Table 4 · p. 15, §E.1 · p. 15

E.2 Per Dataset Results · E.3 Generated Prompts and Code

The appendix reports per-dataset ROC AUC OVO results across downstream classification methods, with CAAFE optimized for strong performance on TabPFN. It also points readers to the repository for generated prompts and LLM-generated code.

E.2 Per Dataset Results

  • Table 5 reports per-dataset ROC AUC OVO results across downstream classification methods, with CAAFE optimized for strong performance on TabPFN.The supplied passage gives no individual table values or winning method.Source: Table 5 · p. 15

E.3 Generated Prompts and Code

  • The repository contains the generated prompts and corresponding LLM-generated code.Source: §E.3 · p. 15

F Compute

Figure F shows that more feature-generation runs increase performance as well as inference cost and time. The evaluation averaged 4:43 minutes per dataset, while GPT-3.5 reduced time to about one-quarter and cost to one-tenth.

F Compute

  • 4:43 minutes was the average CAAFE evaluation time per dataset; GPT-3.5 reduced this time to about 1/4 and cost to 1/10.Code generation, which is performed per token, dominated CAAFE's evaluation time.Source: Figure 6 · p. 16, F Compute · p. 16
  • Increasing feature-generation runs increased performance, inference cost, and time spent.Source: Figure 6 · p. 16
  • TabPFN evaluation used one Nvidia RTX 2080 Ti and 8 Intel Xeon Gold 6242 CPU cores.Source: F Compute · p. 16

G Datasets

The section presents the test datasets used for the evaluation.

G Datasets

  • Table 6 identifies the test datasets used in the evaluation.The dataset descriptions are provided in Section 4.Source: Table 6 · p. 16

G.1 Dataset Collection and Preprocessing · G.2 Dataset Descriptions

The study evaluates CAAFE on constrained small OpenML datasets and four additional Kaggle datasets, using dataset descriptions as context. It documents how those descriptions were collected and cleaned for the prompts.

G.1 Dataset Collection and Preprocessing

  • OpenML datasets had descriptive, non-numbered feature names and at most 20 features, limiting prompt growth to fit GPT-4’s 8,192-token context.The authors cleaned task descriptions of redundant information before providing them as context; some descriptions were much more informative than others.Source: §G.1 · pp. 16–17
  • The study focused on datasets with up to 2,000 samples because feature engineering is considered most important for smaller datasets.Source: §G.1 · p. 17
  • For three datasets perfectly solvable by TabPFN at ROC AUC 1.0, training data were reduced: balance-scale to 20%, and breast-w and tic-tac-toe to 10%.Source: §G.1 · p. 17
  • The evaluation added four post-LLM-cutoff Kaggle datasets with string features, and dropped rows containing missing values.String features allow transformations such as separating names into first and last names for family grouping.Source: §G.1 · p. 17

G.2 Dataset Descriptions

  • Dataset descriptions were crawled from their respective sources; OpenML descriptions were stripped of uninformative source and reference-paper details.Source: §G.2 · p. 17
  • Figure 20 presents the parsed dataset descriptions used for each dataset, including examples for balance-scale and breast-w.Source: §G.2 · p. 17, Figure 7 · p. 17, Figure 8 · p. 17

Related papers

Loading 2305.03403v5…