[Paper Review] ReFuGe: Feature Generation for Prediction Tasks on Relational Databases with LLM Agents

Kyungho Kim, Do Geon Lee+4 more|arXiv (Cornell University)|Jan 25, 2026

Data Quality and Management0 citations

TL;DR

An agentic framework (ReFuGe) uses specialized LLM agents to identify relevant schema, generate diverse relational features, and filter them via reasoning and validation, iterating until predictive performance converges, outperforming baselines on multiple RDB prediction tasks.

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.

Motivation & Objective

• Motivate the need for automated, informed feature generation over relational schemas for prediction tasks in RDBs.
• Propose an agentic framework (ReFuGe) with schema selection, feature generation, and feature filtering agents.
• Demonstrate iterative feedback-driven improvement and self-learning without ground-truth supervision.
• Evaluate ReFuGe across multiple real-world RDB benchmarks to establish efficacy.
• Provide guidelines and insights into when and how relational features improve predictive performance.

Proposed method

• Schema selection agent identifies relevant tables and columns from the RDB schema and task description.
• Feature generation agent creates diverse candidate relational features using multiple LLM instances to encourage variety.
• Reasoning-based feature filtering selects promising features based on semantic reasoning, augmented by prior iterations.
• Validation-based feature filtering temporarily augments the target table with features and trains a tabular model to empirically assess utility (e.g., AUROC).
• Iterative feedback loop where agents receive natural language feedback from previous iterations to guide future feature generation and selection, stopping when no new features are selected.

Experimental results

Research questions

• RQ1 How does ReFuGe compare to state-of-the-art baselines on RDB prediction tasks across diverse datasets?
• RQ2 What is the contribution of each key component (schema selection, feature filtering, iterative feedback) to overall performance?
• RQ3 Does the iterative, feedback-driven loop consistently improve predictive performance?
• RQ4 What representative features does ReFuGe generate in practice (case study)?

Key findings

• ReFuGe outperforms all baselines on most tasks and achieves best average performance and average rank across seven datasets.
• Ablation studies show each component contributes to performance; removing schema selection, filtering, or feedback degrades results, with feature filtering having a notably large impact.
• Performance generally improves over iterations, with an average of 2.4 iterations per task before convergence.
• A case study demonstrates how features derived from related tables (e.g., distinct ads clicked, geographic hierarchy level) can yield substantial gains.
• Increasing the number of LLM instances in the feature generation agent tends to improve performance across tasks.