Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ReFuGe: Feature Generation for Prediction Tasks on Relational Databases with LLM Agents

Kyungho Kim, Geon Lee, Juyeon Kim, Dongwon Choi, Shinhwan Kang, Kijung Shin

TL;DR

The paper tackles predicting outcomes on relational databases by introducing ReFuGe, an agentic framework that leverages specialized LLM agents to generate and refine relational features for a target table. It decomposes the problem into schema selection, diverse feature generation, and two-stage feature filtering, all within an iterative feedback loop that self-improves based on validation outcomes. Empirical results across seven RDB benchmarks show ReFuGe consistently improves predictive performance, outperforming multiple baselines and ablations highlight the importance of schema-aware generation and reasoned filtering. The work demonstrates that LLM-driven relational feature engineering can effectively exploit complex database schemas, with practical impact for real-world predictive tasks on structured data.

Abstract

Relational databases (RDBs) play a crucial role in many real-world web applications, supporting data management across multiple interconnected tables. Beyond typical retrieval-oriented tasks, prediction tasks on RDBs have recently gained attention. In this work, we address this problem by generating informative relational features that enhance predictive performance. However, generating such features is challenging: it requires reasoning over complex schemas and exploring a combinatorially large feature space, all without explicit supervision. To address these challenges, we propose ReFuGe, an agentic framework that leverages specialized large language model agents: (1) a schema selection agent identifies the tables and columns relevant to the task, (2) a feature generation agent produces diverse candidate features from the selected schema, and (3) a feature filtering agent evaluates and retains promising features through reasoning-based and validation-based filtering. It operates within an iterative feedback loop until performance converges. Experiments on RDB benchmarks demonstrate that ReFuGe substantially improves performance on various RDB prediction tasks. Our code and datasets are available at https://github.com/K-Kyungho/REFUGE.

ReFuGe: Feature Generation for Prediction Tasks on Relational Databases with LLM Agents

TL;DR

The paper tackles predicting outcomes on relational databases by introducing ReFuGe, an agentic framework that leverages specialized LLM agents to generate and refine relational features for a target table. It decomposes the problem into schema selection, diverse feature generation, and two-stage feature filtering, all within an iterative feedback loop that self-improves based on validation outcomes. Empirical results across seven RDB benchmarks show ReFuGe consistently improves predictive performance, outperforming multiple baselines and ablations highlight the importance of schema-aware generation and reasoned filtering. The work demonstrates that LLM-driven relational feature engineering can effectively exploit complex database schemas, with practical impact for real-world predictive tasks on structured data.

Abstract

Relational databases (RDBs) play a crucial role in many real-world web applications, supporting data management across multiple interconnected tables. Beyond typical retrieval-oriented tasks, prediction tasks on RDBs have recently gained attention. In this work, we address this problem by generating informative relational features that enhance predictive performance. However, generating such features is challenging: it requires reasoning over complex schemas and exploring a combinatorially large feature space, all without explicit supervision. To address these challenges, we propose ReFuGe, an agentic framework that leverages specialized large language model agents: (1) a schema selection agent identifies the tables and columns relevant to the task, (2) a feature generation agent produces diverse candidate features from the selected schema, and (3) a feature filtering agent evaluates and retains promising features through reasoning-based and validation-based filtering. It operates within an iterative feedback loop until performance converges. Experiments on RDB benchmarks demonstrate that ReFuGe substantially improves performance on various RDB prediction tasks. Our code and datasets are available at https://github.com/K-Kyungho/REFUGE.
Paper Structure (7 sections, 4 figures, 2 tables)

This paper contains 7 sections, 4 figures, 2 tables.

Table of Contents

  1. Introduction & Related Works
  2. Task Formulation
  3. Proposed Framework
  4. Experiments
  5. Experimental Settings
  6. Experimental Results
  7. Conclusions

Figures (4)

  • Figure 1: An example of feature generation by ReFuGe across iterations. ReFuGe iteratively generates and appends relational features to the target table.
  • Figure 2: Overview of ReFuGe. Through (1) schema selection, (2) feature generation, and (3) feature filtering, each performed by a specialized LLM agent. ReFuGe iteratively generates relational features that enrich the target table and improve performance on RDB prediction tasks. At each iteration, feedback is generated and passed to the agents to guide their next rounds.
  • Figure 3: (RQ3) Evaluation performance tends to improve over iterations across all three tasks. The $\star$ indicates the best validation performance over the iterations.
  • Figure 4: (RQ5) Evaluation performance tends to increase as more LLM instances are used as feature-generation agents. The $\star$ indicates the best test performance.