From Features to Structure: Task-Aware Graph Construction for Relational and Tabular Learning with GNNs

TL;DR

This paper introduces auGraph, a framework that enhances graph structures for relational and tabular data by task-aware augmentation, leading to improved GNN-based learning performance.

Contribution

The paper proposes auGraph, a novel method for task-aware graph augmentation that integrates relevant attributes into nodes, improving GNN effectiveness on structured data.

Findings

auGraph outperforms schema-based graph construction methods.

Enhanced graphs better support relational and tabular prediction tasks.

Task-aware augmentation improves GNN learning efficiency.

Abstract

Tabular and relational data remain the most ubiquitous formats in real-world machine learning applications, spanning domains from finance to healthcare. Although both formats offer structured representations, they pose distinct challenges for modern deep learning methods, which typically assume flat, feature-aligned inputs. Graph Neural Networks (GNNs) have emerged as a promising solution by capturing structural dependencies within and between tables. However, existing GNN-based approaches often rely on rigid, schema-derived graphs -- such as those based on primary-foreign key links -- thereby underutilizing rich, predictive signals in non key attributes. In this work, we introduce auGraph, a unified framework for task-aware graph augmentation that applies to both tabular and relational data. auGraph enhances base graph structures by selectively promoting attributes into nodes, guided by scoring functions that quantify their relevance to the downstream prediction task. This augmentation preserves the original data schema while injecting task-relevant structural signal. Empirically, auGraph outperforms schema-based and heuristic graph construction methods by producing graphs that better support learning for relational and tabular prediction tasks.