[Paper Review] GRATIS: Deep Learning Graph Representation with Task-specific Topology and Multi-dimensional Edge Features

Siyang Song, Yuxin Song+8 more|arXiv (Cornell University)|Nov 19, 2022

Advanced Graph Neural Networks9 citations

TL;DR

GRATIS is a novel, plug-and-play framework that learns task-specific graph topology and multi-dimensional edge features from any input data—whether pre-defined graphs or non-graph data—using joint vertex and global context modeling. It significantly improves performance across graph classification, vertex classification, and link prediction tasks by generating adaptive, complementary representations that outperform fixed-topology and single-feature baselines on 11 datasets.

ABSTRACT

Graph is powerful for representing various types of real-world data. The topology (edges' presence) and edges' features of a graph decides the message passing mechanism among vertices within the graph. While most existing approaches only manually define a single-value edge to describe the connectivity or strength of association between a pair of vertices, task-specific and crucial relationship cues may be disregarded by such manually defined topology and single-value edge features. In this paper, we propose the first general graph representation learning framework (called GRATIS) which can generate a strong graph representation with a task-specific topology and task-specific multi-dimensional edge features from any arbitrary input. To learn each edge's presence and multi-dimensional feature, our framework takes both of the corresponding vertices pair and their global contextual information into consideration, enabling the generated graph representation to have a globally optimal message passing mechanism for different down-stream tasks. The principled investigation results achieved for various graph analysis tasks on 11 graph and non-graph datasets show that our GRATIS can not only largely enhance pre-defined graphs but also learns a strong graph representation for non-graph data, with clear performance improvements on all tasks. In particular, the learned topology and multi-dimensional edge features provide complementary task-related cues for graph analysis tasks. Our framework is effective, robust and flexible, and is a plug-and-play module that can be combined with different backbones and Graph Neural Networks (GNNs) to generate a task-specific graph representation from various graph and non-graph data. Our code is made publicly available at https://github.com/SSYSteve/Learning-Graph-Representation-with-Task-specific-Topology-and-Multi-dimensional-Edge-Features.

Motivation & Objective

• Address the limitation of fixed, manually defined graph topologies and single-value edge features in existing graph representation learning methods.
• Enable task-specific, adaptive graph construction from arbitrary inputs, including non-graph data such as images or sequences.
• Improve downstream graph analysis tasks by jointly learning topology and rich, multi-dimensional edge features that encode task-relevant relationships.
• Provide a generalizable, robust, and flexible framework that integrates with various GNN backbones and works across diverse data types.
• Demonstrate that task-specific topology and multi-dimensional edge features are complementary and essential for optimal performance in graph learning tasks.

Proposed method

• The Graph Definition module constructs an initial raw graph from input data using predefined or learned vertex representations.
• The Task-specific Topology Prediction (TTP) module learns a probability-based adjacency matrix using vertex features and global context, enabling task-adaptive edge presence prediction.
• The TTP module also generates task-specific vertex features for non-graph data, which inform the final topology and improve representation quality.
• The Multi-dimensional Edge Feature Generation (MEFG) module produces task-specific, multi-dimensional edge features for each connected vertex pair using a learnable network.
• The framework integrates with any GNN backbone and uses end-to-end training with cross-entropy and contrastive loss to optimize topology and edge features jointly.
• The final graph representation is formed by combining learned topology and multi-dimensional edge features, which are used as input to downstream GNNs for classification or prediction tasks.

Experimental results

Research questions

• RQ1 Can a unified framework learn task-specific graph topology and multi-dimensional edge features from arbitrary input data, including non-graph data?
• RQ2 How do task-specific topology and multi-dimensional edge features compare to fixed topologies and single-value edge features in terms of performance on graph analysis tasks?
• RQ3 To what extent do the learned topology and multi-dimensional edge features complement each other in improving downstream task accuracy?
• RQ4 How robust is the framework to variations in hyperparameters such as the number of vertices or the scaling factor C?
• RQ5 Can the framework be effectively combined with various GNN backbones and still improve performance across diverse datasets and tasks?

Key findings

• GRATIS consistently improves performance across 11 datasets in graph classification, vertex classification, and link prediction, with clear gains over fixed-topology and single-feature baselines.
• The combination of task-specific topology and multi-dimensional edge features yields the best performance, demonstrating their complementary nature in encoding task-relevant cues.
• The framework achieves stable performance across different values of the hyperparameter C (with less than 3% variation in F1 scores on BP4D and DISFA datasets), indicating robustness to hyperparameter settings.
• Multi-dimensional edge features outperform single-dimensional edge features, confirming their importance in capturing complex, task-specific relationships between vertices.
• GRATIS improves performance across diverse GNN backbones and predictors, showing its generalization and compatibility with existing GNN architectures.
• The framework is effective even on non-graph data, such as facial Action Units, where it learns a superior graph representation compared to existing methods using fixed topologies and single-value features.