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 provides limited support for data engineering and domain-knowledge integration, motivating a method that extends automation beyond model optimization. CAAFE uses LLMs to generate semantically meaningful feature code and explanations from dataset context, improving mean ROC AUC from 0.798 to 0.822 on 11 of 14 datasets. The approach demonstrates interpretable, context-aware feature engineering while retaining classical evaluation and exposes limitations involving dataset descriptions, large prompts, statistical testing, hallucinations, and code-execution safety.

  • Problem

    Existing AutoML systems support model optimization but leave data engineering and domain-knowledge integration largely to human practitioners.

  • Method

    CAAFE prompts an LLM with dataset context to iteratively generate Python code for semantically meaningful features and textual explanations, evaluating changes with classical algorithms.

  • Results

    Mean ROC AUC increased from 0.798 to 0.822, with improvements on 11/14 datasets when CAAFE used GPT-4.

  • Takeaways & Limitations

    CAAFE demonstrates that LLMs can automate context-aware feature engineering while preserving interpretable feature explanations and complementing existing AutoML methods.

  • Takeaways & Limitations

    Effectiveness depends on dataset-description quality, while large prompts, non-statistical feature testing, hallucinations, and incomplete execution safeguards remain limitations.

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

Existing AutoML systems optimize much of the modeling workflow but provide limited support for data engineering and domain-knowledge integration. CAAFE addresses this gap by combining LLM-based semantic feature generation with classical ML methods and interpretable explanations.

  • About 23% of data scientists’ time is spent on model selection, training, and scoring, leaving data engineering and cleaning insufficiently supported by AutoML tools.
  • CAAFE uses LLM-generated Python code to modify datasets, which classical algorithms then process for downstream prediction.
  • CAAFE iteratively generates semantically meaningful features and uses algorithmic feedback to improve prediction performance.
  • Each generated feature receives a textual utility explanation, supporting more interpretable and modifiable AutoML solutions.
  • Automated domain-knowledge integration could reduce latency and cost, evaluate a more informed solution space, and improve reproducibility.

2 Background

The background motivates semantic feature engineering as a way to incorporate context into tabular prediction while combining LLM flexibility with classical algorithmic reliability. Prior work shows both the promise of LLMs and their computational and interpretability limitations.

  • Large Language Models (LLMs): LLM-based tabular prediction can become computationally expensive because cost increases quadratically with the product of samples and features.
  • Large Language Models (LLMs): LLM tabular prediction performed best on tiny datasets with up to 8 samples but was outperformed on larger datasets.
  • Related Work: CAAFE differs from row-wise LLM data-wrangling methods by generating reusable code that can be applied to all samples through one LLM query.
  • Feature Engineering: Feature engineering constructs features from raw inputs to improve predictive performance, using transformations such as encoding, aggregation, and dimensionality reduction.
  • Feature Engineering: Explicit feature engineering remains useful for limited data, complex patterns, or model interpretability, despite deep models learning transformations from raw inputs.
  • Incorporating Semantic Information: Because the feature space is combinatorial, semantic and contextual information can prioritize meaningful features while reducing multiple-testing and computational burdens.

3 Method

CAAFE prompts an LLM with dataset context to generate feature-engineering code and explanations, then iteratively evaluates and retains transformations that improve validation performance. The method combines semantic guidance with classical evaluation and constrained code execution.

  • Method: CAAFE takes training and validation data plus a dataset-context description as input and prompts an LLM to generate feature modifications.
  • Method: Generated code is executed on both datasets, and a feature is retained only when the transformed data improves the evaluation criterion.
  • Prompting LLMs for Feature Engineering Code: The prompt asks for valuable features, utility justifications, and removal of unnecessary features whose information is already captured elsewhere.
  • Prompting LLMs for Feature Engineering Code: The prompt combines user-generated context with descriptive information such as missing-value percentages and sample dataset rows.
  • Illustrative Run: In an illustrative Tic-Tac-Toe run, validation ROC AUC rose from 0.888 to 1.0 in two feature-engineering iterations.
  • Prompting LLMs for Feature Engineering Code: A response template and chain-of-thought example guide the LLM from feature meaning and usefulness through input features, sample values, and code.
  • Technical Setup: CAAFE uses ten feature-engineering iterations, GPT-4 or GPT-3.5, and TabPFN for iterative evaluation of generated code.
  • Safety and Limitations: Syntax parsing and an execution whitelist restrict generated Python operations, but do not prevent infinite loops or excessive resource use.

4 Experimental Setup

The evaluation compares CAAFE with multiple downstream classifiers and context-agnostic feature-engineering methods across established and lesser-known small datasets. Repeated common train-test splits and ROC AUC OVO results support comparison while tracking variation across splits.

  • Downstream Classifiers: CAAFE is evaluated with Logistic Regression, Random Forests, and TabPFN, while TabPFN evaluates the performance of added features.
  • Automated Feature Engineering Methods: Deep Feature Synthesis and AutoFeat are evaluated alone and after CAAFE to compare context-agnostic and context-aware feature engineering.
  • Datasets: The dataset collection includes recognized pre-September-2021 OpenML datasets and lesser-known post-September-2021 Kaggle datasets.
  • Datasets: The study focuses on small datasets with descriptive feature names and up to 2,000 samples, reducing training size when TabPFN otherwise achieves ROC AUC 1.0.
  • Results Reporting: Table 1 reports ROC AUC OVO results using TabPFN, with standard deviations across five splits and [R] marking reduced-data datasets.
  • Evaluation Protocol: Each dataset uses five repetitions with different random seeds and 50% train-test splits shared by all methods.

5 Results

CAAFE improves downstream classification by generating context-informed feature extensions, and it can be combined with classical automated feature engineering methods. Results also show that performance gains vary by dataset and that the method can recover from faulty generated features.

  • CAAFE improves TabPFN substantially, raising average ROC AUC from 0.798 to 0.822 with GPT-4 and improving performance on 11/14 datasets.This evaluation uses the downstream classifier TabPFN.
  • CAAFE generates feature extensions by combining features, binning numerical variables, transforming strings, removing superfluous features, and recovering from invalid code.The generated code and comments are based on user-provided dataset descriptions.
  • CAAFE’s improvement is similar (71%) to the average gain from replacing logistic regression with a random forest, comparing AUC 0.749 with AUC 0.783.
  • GPT-3.5 performs clearly worse than GPT-4, improving performance on only 6/14 datasets.
  • Improvement size varies with whether a dataset admits a useful code-expressible feature mapping and with the quality of its dataset description.
  • Across 14 datasets, five splits, and ten iterations, CAAFE generated 52 faulty features (7.4%) and recovered from them.
  • Classical AutoFE methods can follow CAAFE and further improve Logistic Regression and Random Forests, whereas the evaluated AutoFE methods do not improve TabPFN.Across all classifier and additional-AutoFE combinations, CAAFE improves performance on average.

6 Conclusion

CAAFE integrates domain knowledge into AutoML by generating semantically meaningful features and explanations with LLMs. Its conclusions also identify important limitations and support human-in-the-loop extensions.

  • CAAFE automates feature engineering for tabular datasets while generating semantically meaningful features and explanations of their utility.
  • Context-aware solutions are presented as important for robust outcomes, with LLMs potentially supporting additional data-science tasks.
  • Dataset descriptions constrain CAAFE because the study used only web-crawled text, while more accurate user descriptions could improve effectiveness.
  • Large feature sets can create prompts that challenge LLM processing, while non-statistical feature testing and hallucinations remain limitations.
  • Human-in-the-loop AutoML is a promising direction because CAAFE’s LLM inputs and outputs are interpretable and easily modified by experts.
  • Automating data-science tasks may displace workers, although routine-task automation could free data scientists for higher-level problem-solving and decision-making.

B.2 Replication of Biases

CAAFE can reproduce bias through LLM-generated features, but its prompt explanations make this risk visible and its workflow includes user checks before execution.

  • Bias risks: CAAFE can inherit biases from GPT-4 and generate biased features, especially when demographic information could support discrimination.The authors advise against using CAAFE for such data or recommend extensive checking of generated features.
  • Mitigations: The system asks users to confirm generated code is free of dangerous actions or biases before execution, keeping the workflow human-in-the-loop.The authors frame this as a design choice aligned with using LLMs as decision-support tools rather than final decision makers.
  • Mitigations: CAAFE retains generated feature operations only when they improve cross-validation, discarding biases present in the LLM but absent from the data.This provides somewhat more defense against bias than directly using LLM outputs.
  • Bias replication: In a doctor-versus-nurse example, GPT-4 generated a name-ending feature that reflected gender-associated patterns in the intentionally biased sample.The generated feature tests whether names end with “a.”
  • Bias replication: The prompt’s required usefulness explanation exposed the model’s bias in this non-cherry-picked example.The explanation linked names ending in “a” with female gender and potential occupational bias.
  • Interpretability: CAAFE’s interpretability goal is supported by simple generated code and explanations for model outputs and engineered features.The authors present this as making automated feature engineering more transparent.

E.2 Per Dataset Results

The per-dataset evaluation reports ROC AUC OVO results for each dataset and downstream classification method, with CAAFE optimized for TabPFN.

  • Per-dataset results: Table 5 reports ROC AUC OVO results by dataset and downstream classification method.The table evaluates CAAFE in combination with different downstream classifiers.
  • Per-dataset results: CAAFE was optimized for strong performance on TabPFN in the reported per-dataset results.The supplied table description identifies TabPFN as the classifier used for CAAFE optimization.

F Compute

More feature-engineering iterations increase performance but also raise computational cost and runtime; code generation dominates CAAFE’s evaluation time.

  • Compute trade-offs: More feature-engineering iterations increase performance while also increasing cost and time spent.Figure 6 presents mean ROC AUC OVO alongside inference cost and runtime as iterations increase.
  • Compute trade-offs: 4:43 minutes is the average CAAFE evaluation time per dataset, with LLM code generation dominating runtime.GPT-3.5 takes about one quarter of this time, and its cost is about one tenth as of the paper’s writing.

G Datasets

Table 6 lists the test datasets used for evaluation and points to Section 4 for their descriptions.

  • Test datasets: Table 6 lists the test datasets used for the evaluation.The table directs readers to Section 4 for descriptions of the datasets.

G.1 Dataset Collection and Preprocessing

The study uses small, descriptively named OpenML datasets and supplements them with four Kaggle datasets containing string features. Preprocessing limits dataset size and feature count to accommodate the method’s context-based prompting and evaluation design.

  • OpenML datasets: OpenML datasets have descriptive feature names, task descriptions supplied as user context, and no more than 20 features.The feature limit reflects prompt-length constraints for standard GPT-4.
  • OpenML datasets: The evaluation focuses on datasets with up to 2,000 samples because feature engineering is considered especially important for smaller datasets.
  • OpenML datasets: For balance-scale, breast-w, and tic-tac-toe, training data were reduced to 20%, 10%, and 10%, respectively, because TabPFN perfectly solved the original tasks.
  • Kaggle datasets: Four post-cutoff Kaggle datasets add string features that support transformations such as separating names into first and last names for family grouping.Rows containing missing values were dropped for these evaluations.

G.2 Dataset Descriptions

Dataset descriptions provide the contextual information used for feature engineering, covering OpenML and Kaggle tasks across demographic, medical, software, game, environmental, and transportation domains.

  • Description construction: The dataset descriptions were crawled from their respective data sources and cleaned for OpenML by removing uninformative details such as source information and reference papers.The parsed descriptions used for each dataset are shown in the paper’s figures.
  • OpenML descriptions: The collection spans prediction tasks involving contraceptive-method choice, diabetes-related patient data, balance scale, breast cancer, and software defects.
  • OpenML and Kaggle descriptions: Other descriptions cover eucalyptus soil conservation, wine, airlines, Jungle Chess endgames, pediatric pharyngitis, and the Spaceship Titanic competition.
  • Dataset-specific caveat: The eucalyptus dataset description notes that changes in trial methods over time may contribute to less significant results.
Loading 2305.03403v5…