U2B: Scale-unbiased Representation Converter for Graph Classification with Imbalanced and Balanced Scale Distributions

Read full text
Authors: Association for Artificial Intelligence 2026, Sheng Huang, Jiaming Ma, Guanjun Wang, Pengkun Wang, Yang Wang, Binwu Wang, Jianhao Zhang, Zhengyang Zhou
Cited by: 0

Abstract

Graph classification is a critical task in analyzing graph data, with applications across various domains. While graph neural networks (GNNs) have achieved remarkable results, their ability to generalize across graphs of varying scales remains a challenge. Conventional models often perform well on largescale graphs but struggle with distributions that are skewed towards small scales. Conversely, models tailored to address scale imbalances frequently prioritize small-scale graphs, leading to diminished performance in more balanced scenarios. To overcome these limitations, we introduce a Unbalanced-Balanced Representation Converter (U2B), which exhibits no explicit bias toward graph scales. U2B employs a two-step workflow: a distillation phase to extract base features from both node-level and graph-level representations, followed by a refinement phase to generate biased representations for improved balance. In the distillation phase, a static constraint guides node-level adjustments, improving the representation of nodes in small graphs. Simultaneously, a dynamic constraint in the graph-level process mitigates biases toward features from large graphs. To ensure harmony between the representations, a consistency alignment loss is introduced, aligning node-level and graph-level features to create more cohesive and balanced graph representations. Extensive experiments with 16 baselines across 8 datasets demonstrate that U2B achieves SOTA performance, boasting improvements of up to 22.19%. Additionally, we establish its strong compatibility with a range of other models. All associated code is provided in Supplement.

Fields of Study & Subjects

Advanced Graph Neural Networks Graph Theory and Algorithms Computational Drug Discovery Methods

Publication details

Journal / source
Underline Science Inc.
Publication date
7 January 2026
Language
EN