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

AI summary

Overview

Research area: Relational databases and machine learning — specifically automated feature engineering on relational data using large language model (LLM) agents, positioned at the intersection of data management, tabular machine learning, and LLM agent design.

Technical level: Intermediate. The paper assumes familiarity with relational schemas, primary–foreign key joins, tabular models such as LightGBM, and the general idea of LLM agents, but its core ideas are explained conceptually rather than through heavy formalism.

Scope (one sentence): The paper introduces ReFuGe, an agentic framework in which three specialized LLM agents iteratively select schema elements, generate relational features, and filter them, in order to improve prediction accuracy on relational database (RDB) tasks.

What This Paper Is About

Prediction tasks on relational databases (for example, predicting whether a user will churn or click) are usually handled with only the columns of one target table, which throws away predictive signals scattered across related tables. The paper's goal is to automatically generate new features for that target table by reasoning over the relational schema, so that an ordinary tabular model can then be trained on an enriched, wider table. The difficulty is that the space of possible relational features is combinatorially large, the schema is complex, and there is no ground-truth set of "correct" features to learn from.

Key Contributions

1. A new problem formulation. The authors define the problem of generating relational features for prediction tasks on RDBs, formalizing the database as a set of tables connected by primary–foreign key relationships, with the goal of generating additional features for a designated target entity table and concatenating them with its original attributes to train a tabular model such as LightGBM.

2. An agentic framework (ReFuGe). They propose ReFuGe (RElational FeatUre GEneration), composed of a schema selection agent that identifies relevant tables and columns, a feature generation agent that proposes candidate features (using multiple LLM instances to encourage diversity), and a feature filtering agent that applies reasoning-based filtering followed by validation-based filtering.

3. An iterative self-improvement mechanism. Because no ground-truth feature supervision exists, ReFuGe loops through the three steps, generating natural-language feedback about how each evaluated feature changed validation performance and feeding that accumulated feedback back to the agents.

4. Empirical validation. Across seven real-world RDB benchmark datasets and eleven tasks, ReFuGe reports the best average AUROC (75.30) and best average rank (1.3), outperforming all baselines in 9 out of 11 tasks; code and datasets are released at https://github.com/K-Kyungho/REFUGE.

Main Findings

• Best overall predictive performance. ReFuGe achieves an average AUROC of 75.30 and an average rank of 1.3 across the eleven tasks, winning 9 of 11 tasks. For comparison, the strongest baselines in the same table are LLM-CoT at 70.81 (rank 4.2), GPT at 69.94 (rank 4.3), ICL at 69.42 (rank 4.7), Claude at 69.49 (rank 4.7), RT-Gemma at 67.12 (rank 5.9), Rel-LLM at 63.36 (rank 6.0), LightGBM at 63.77 (rank 6.8), XGBoost at 61.57 (rank 7.5), and FeatLLM at 52.51 (rank 8.7).

• Advantage on both simple and complex schemas. The paper notes ReFuGe performs best on RDBs with simple schemas, such as rel-amazon and rel-hm with 3 tables and 2 primary–foreign key relations, and also on the more complex rel-stack with 7 tables and 12 primary–foreign key relations.

• Ablation: reasoning-based filtering matters most. Removing the reasoning-based feature filtering (ReFuGe-FF, replaced by random selection) causes the largest performance drop, yielding an average AUROC of 70.56 (rank 3.0) versus ReFuGe's 75.30 (rank 1.5) in that table. Removing the feedback loop (ReFuGe-FB) causes the second-largest drop, at 70.88 (rank 3.1), and skipping schema selection (ReFuGe-SS) gives 73.77 (rank 2.3). ReFuGe outperforms all variants in 7 out of 11 cases and has the best average rank.

• Iteration with feedback helps. Evaluation performance generally improves over iterations across three tasks, and the loop stops when no additional feature is selected. ReFuGe performs an average of 2.4 iterations across all tasks.

• More generation agents tend to help. Using more LLM instances as the feature generation agent tends to improve performance across three tasks, with few exceptions; the authors fix the number of instances to three for simplicity and efficiency.

• Concrete generated features (case study). On the rel-avito dataset for the user-clicks task — predicting whether each user will click on more than one advertisement in the next 4 days — ReFuGe first derives the number of distinct ads clicked per user from the VisitStream table, and in the next iteration adds the average geographic hierarchy level of users' search locations by joining the SearchStream and Location tables.

• What feedback looks like. Feedback is textual, for example "[Ad View Diversity] increased validation AUROC from 65.06 to 66.30" or "[Direct Search Ratio] caused validation AUROC to drop from 66.68 to 66.40".

• Not reported. The paper content provided does not specify which LLM backbone ReFuGe itself uses (only that the same model is used across instances with a positive temperature), nor dataset sizes, nor the number of candidate features generated.

Methodology in Plain English

The input is a relational database: several tables linked by primary–foreign key relationships, plus a designated target table and a task description. The output is that same target table with new columns added. ReFuGe runs three agents in a loop.

First, a schema selection agent reads the schema — table names, column names, column types, and key relationships — along with the task description, and decides which tables and columns are relevant, discarding the rest. This shrinks the search space before any features are proposed.

Second, a feature generation agent proposes candidate features over that reduced schema, drawing on reasoning about the relational context and domain knowledge. Multiple LLM instances generate features separately so that the resulting pool is diverse; the pool is then merged.

Third, a feature filtering agent prunes the pool in two stages. In reasoning-based filtering, the agent judges which candidates look semantically promising for the task; in the first iteration it relies on semantics alone, and from the second iteration it also uses performance of previously generated features. In validation-based filtering, each surviving feature is temporarily added to the target table, a tabular model (LightGBM) is trained — in practice on a sampled subset of the training set for efficiency — and the feature is kept only if validation performance improves; otherwise it is discarded.

At the end of each round, ReFuGe writes natural-language feedback describing how each evaluated feature affected validation performance. This feedback accumulates across rounds and is given back to all agents, so the system progressively learns which directions are promising. The loop repeats until no additional features are selected, and a standard tabular model is then trained on the enriched table.

Why This Matters

Impact on research. The paper reframes RDB prediction as a feature generation problem and shows that LLM agents can reason over schemas to produce useful relational features without any ground-truth feature supervision. It also offers an ablation-grounded argument that reasoning-based selection and cross-iteration feedback each carry meaningful performance weight, rather than being decorative.

Real-world applications (as described or implied by the paper's domains):

• Online retail platforms — enriching customer tables with aggregated purchase statistics or recent transaction counts, as the paper's motivating example describes.
• Social networking services — one of the web application types the paper lists as relying on RDBs.
• Forum platforms — likewise listed as a core RDB use case.
• Ad and event prediction — the rel-avito and rel-event tasks, such as predicting repeated user engagement or whether a user clicks more than one advertisement in the next 4 days.

Industry relevance. Because the output is a wider tabular table, the generated features plug directly into established tabular learners such as LightGBM and XGBoost, meaning teams can adopt the framework without replacing their modeling stack. The reported advantage on both simple (3-table) and more complex (7-table, 12-relation) schemas suggests applicability across organizations with differing data-model complexity.

Future Directions

• Extending beyond classification. The authors explicitly state they aim to extend the framework to additional downstream tasks such as regression and link prediction.
• Scaling and tuning the agent population. The parameter analysis shows gains from more LLM instances in the generation agent, but the study fixes the count at three; how far this scaling continues, and how to choose it per task, remains open.
• Reducing the cost of validation-based filtering. The filtering step trains a model per candidate feature, on a sampled training subset for efficiency; making this process cheaper or more sample-efficient is an unresolved practical question.
• Understanding when schema selection helps versus hurts. ReFuGe-SS was the closest-performing ablation variant, which raises the question of under what schema or task conditions the schema selection step is unnecessary or could be replaced by a lighter mechanism.

Target Audience

Researchers and practitioners working on relational data management, automated feature engineering, tabular machine learning, and LLM agent systems will benefit most. It is also relevant to data scientists at web-scale companies who maintain multi-table databases and want to improve predictive models without hand-engineering relational aggregates, and to readers interested in agentic workflows that operate without explicit supervision.

Authors’ 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.