Heterogeneous Sheaf Neural Networks

TL;DR

HetSheaf introduces a topologically principled framework using cellular sheaves for modeling heterogeneous graphs, enabling efficient, parameter-reduced predictions across various tasks.

Contribution

It presents HetSheaf, a novel sheaf-based approach for heterogeneous graph modeling, including a new pooling method for graph-level tasks, with improved efficiency and performance.

Findings

Achieves strong predictive performance on standard benchmarks.

Reduces parameters by up to 10x compared to state-of-the-art.

Supports multiple tasks like node classification, link prediction, and recommendation.

Abstract

Heterogeneous graphs, whose nodes and edges may belong to different types and feature spaces, arise in a wide variety of real-world domains such as biology, chemistry and computer networks. Existing methods typically address this heterogeneity by modifying the model architecture itself, which often results in specialized and parameter-intensive designs. To address this issue, we propose HetSheaf, a framework that models heterogeneous relational data through cellular sheaves, which provide a principled topological framework for encoding type-specific local feature spaces and their interactions directly in the data representation. We also introduce a family of heterogeneous sheaf predictors that learn restriction maps conditioned on node and edge types. To enable graph-level predictions, we further propose SheafPool, a graph pooling mechanism that aggregates node representations in stalk space while remaining invariant to local changes of basis, ultimately enabling stalk-space graph-level representations for the first time. HetSheaf achieves strong predictive performance on standard heterogeneous graph benchmarks, over numerous tasks such as node/graph classification, link prediction and recommendation, while reducing by up to 10x the number of parameters with respect to state-of-the-art baselines.