Distributionally-robust heterogeneous graph neural networks via relation-aware grouping

Abstract

Out-of-distribution (OOD) generalization remains a key challenge for heterogeneous graph neural networks (HGNNs), which often exhibit substantial performance degradation under real-world distribution shifts. While Group Distributionally Robust Optimization (GroupDRO) has proven effective in promoting OOD generalization in various domains, its application to heterogeneous graphs remains largely unexplored. This gap stems from a fundamental obstacle: the complex, multi-relational nature of heterogeneous graphs renders conventional group definitions, which rely on explicit node attributes, inadequate. To bridge this gap, we propose DroHGNN, a GroupDRO-based framework for OOD generalization on heterogeneous graphs. DroHGNN introduces a relation-aware soft grouping mechanism coupled with a two-stage optimization procedure to enhance OOD robustness. Specifically, DroHGNN learns soft group assignments through learnable prototypes over relation-specific node representations. These assignments serve as optimization-oriented group proxies for GroupDRO, which minimizes worst-case group risk. To encourage structural coherence of the inferred assignments, we further introduce a heterogeneous graph structural consistency constraint that regularizes the group assignments of structurally connected target nodes. Extensive experiments on multiple real-world benchmarks show that DroHGNN consistently achieves strong OOD performance compared with competitive baselines. Notably, under degree-distribution shifts on the ArXiv dataset, DroHGNN outperforms NaiveDRO and the leading baseline CaNet by 23.3% and 6.5% in Macro-F1, respectively. Moreover, relative gains of up to 3.2% in average Macro-F1 under the inductive temporal setting further demonstrate its effectiveness under heterogeneous distribution shifts.