One for Two: A Unified Framework for Imbalanced Graph Classification via Dynamic Balanced Prototype
Authors: Guanjun Wang
Binwu Wang
Jiaming Ma
Zhengyang Zhou
Pengkun Wang
Xu Wang
Yang Wang
Venue: ICLR 2026 Oral (2026)
Keywords: Graph classification; graph imbalance learning; graph neural networks; Graph data mining; long-tail learning
Primary area: learning on graphs and other geometries & topologies

Abstract: Graph Neural Networks (GNNs) have advanced graph classification, yet they remain vulnerable to graph-level imbalance, encompassing class imbalance and topological imbalance. To address both types of imbalance in a unified manner, we propose UniImb, a Unified framework for Imbalanced graph classification. Specifically, UniImb first captures multi-scale topological features and enhances data diversity via learnable personalized graph perturbations. It then employs a dynamic balanced prototype module to learn representative prototypes from graph instances, improving the quality of graph representations. Concurrently, a prototype load-balancing optimization term mitigates dominance by majority samples to equalize sample influence during training. We justify these design choices theoretically using the Information Bottleneck principle. Extensive experiments on 19 datasets-including a large-scale imbalanced air pollution graph dataset AirGraph released by us and 23 baselines demonstrate that UniImb has achieved dominant performance across various imbalanced scenarios. Our code is available at GitHub.

Published as a conference paper at ICLR 2026
ONE FOR Two: A UNIFIED FRAMEWORK FOR IMBAL- ANCED GRAPH CLASSIFICATION VIA DYNAMIC BAL- ANCED PROTOTYPE
Guanjun Wang', Binwu Wang Jiaming Ma',
Zhengyang Zhou! Pengkun Wang' Xu Wang Yang Wang
I University of Science and Technology of China (USTC), Hefei, Anhui, China
2 Suzhou Institute for Advanced Research, USTC, Suzhou, Jiangsu, China falways, JiamingMa]enan @maal. emst.uustem...........m..o wuw..uud..ccu........ooooo edu. (bbw2224, z2y0929, pengkun, wx309, suryaann.ust..ddu.coo
ABSTRACT
Graph Neural Networks (GNNs) have advanced graph classification, yet they remain vulnerable to graph-level imbalance, encompassing class imbalance and topological imbalance. To address both types of imbalance in a unified manner, we propose Unilmb, a Unified framework for Imbalanced graph classification. Specifically, Unilmb first captures multi-scale topological features and enhances data diversity via learnable personalized graph perturbations. It then employs a dynamic balanced prototype module learn representative prototypes from graph instances, improving the quality of graph representations. Concurrently, proto- type load-balancing optimization term mitigates dominance by majority samples to equalize sample influence during training. We justify these design choices theoretically using the Information Bottleneck principle. Extensive experiments on 19 batasets-including a large-scale imbalanced air pollution graph dataset AirGraph released by us and 23 baselines demonstrate that Unilmb has achieved dominant performance across various imbalanced scenarios. Our code is available at GitHub
INTRODUCTION
Graph Neural Networks (GNNs) have achieved remarkable performance graph classification tasks by iteratively aggregating local neighborhood information (Wei et al. 2023 Nguyen et al. 2022 Wu ct al. 2020 Fan et al. (023). However, the success of most architectures relies heavily on balanced graph datasets (Ma al. 2025 In real-world scenarios, graph data often exhibits imbalanced distributions, which can be mainly categorized into two aspects:
0 Class Imbalance. A few classes contain most of the samples while many classes severely underrepresented, leading GNNs to favor head classes during training (Tang & Liang 2023). (Tang
Liang 2023 A promising approach to mitigate class imbalance is the graph-of-graphs tramework Wang et al. 20226 Tang & Liang 2023). which treats each input graph as node in a higher-order meta-graph. Edges in the meta-graph are formed based on graph-level feature similarity, enabling information propagation among similar graphs and thereby alleviating data sparsity in tail classes.
Topological Imbalance. Graph datasets frequently show a highly skewed distribution of node counts: small number of large graphs contain most nodes, while the majority of graphs are substantially smaller. This imbalance often causes GNNs focus disproportionately on larger graphs during training, leading to diminished performance small-scale ones QQnn et al. 2025 To address this challenge, topology imbalance methods identify underrepresented small-scale graphs during training and enhance their contributions through reweighting or data augmentation, yielding more
consistent performance (Liu et al. 2022c Xu ct al. 2024
Despite remarkable progress, existing imbalanced graph learning methods focus unilaterally either class imbalance or topological imbalance isolation. Class-imbalance graph learning methods
BBinw Wang and Yang Wang the corresponding authors.

---

Published as a conference paper at ICLR 2026
enhance the expressiveness of underrepresented classes but often overlook the structural heterogeneity within graphs. Conversely, methods focused on topology imbalance adapt to small graphs but typically ignore imbalanced class distributions. This narrow focus limits their ability to address the complex, intertwined data imbalances commonly found in real-world graph datasets. To remedy this, we propose Unilmb, a Unified Imbalance framework for imbalanced graph learning. The core idea of Unilmb is to extract semantically shared prototype features via a balanced extraction that equalizes the influence of tail graphs-boll bots minority-class instances and small-scale graphs. These prototypes enrich graph representations and improve robustness to complex, intertwined imbalances.
Specifically, Unilmb first encodes multi granularity topologiaal features and applies personalized graph perturbation module, which injects controlled randomness into graph features and connections with graph dependent intensity. Graphs from smaller or minority classes are augmented with different strengths than larger, frequent ones. We then design a local prior based GNN to encode topological information and obtain graph level representations. Next, we introduce the Dynamic Balanced Prototype (DBP) module, consisting of a set of learnable prototype embeddings that compactly represent shared semanties. Graphs are softly assigned to prototypes based on similarity shared embedding space, and prototypes are jointly perceived in a collaborative manner. To regulate this process, we derive a prototype regularizer based on the information bottleneck principle, which encourages prototype actívation probabilities to approach a uniform distribution. This allows tail graphs, such as minority class or small scale graphs, to exert comparable influence during prototype learning. Finally, the learned prototypes are used enhance and regularize underrepresented graph representations, thereby improving discriminative performance on imbalanced datasets.
Contributions, 1 Unified Insight. We propose Unilmb, a unified framework for imbalanced graph learning that effectively addresses both class imbalance and topological imbalance in graph classifica- tion tasks. @ Practical Solution. Unilmb integrates series of advanced strategies with a Dynamic Balanced Prototype mechanism, effectively enhancing the representational capacity for tail graphs. ? Theoretical Explanation. Our design incorporates an information Nottleneck-driven regularization, which ensures balanced influence from tail classes during training. Empirical Study. Beyond the standard graph classification datasets commonly used in mainstream benchmarks (Qin et al. 2025 we further introduce 3D conformer graph datasets and AirGraph, a large-scale real-world air- pollution graph dataset with naturally occurring long-tailed imbalance, to spur further research and practical progress in this area. We conduct systematic evaluation on 19 datasets, and the results show that Unilmb consistently achieves stable and highly competitive performance in challenging imbalanced graph classification scenarios. Moreover, Unilmb is well compatible with a wide range of graph representation backbones and remains efficient while maintaining strong performance.
2 RELATED WORK
Graph classification is a foundational task in graph data analysis (Zhang et al. 2018 with GNNs emerging as the dominant architecture for this task /Velickovic et al. 2001 Xu et al. 2018). More re- cently, inspired by the success of Transformer architectures in natural language processing, researchers have extended these approaches to graph representation learning (uun et al. 2019 Rampášek et al. Wang et al. 2024b). However, these methods encounter challenges with graph imbalance
2022
distribution (Liu et al. 2025). Imbalanced graph learning has emerged as a dedicated field to address this issue, comprising two main branches: class imbalance and topology imbalance (Qin et al. 2024 Class-inbalanced graph learning methods aim to correct skewed class distributions (Mao et al. 2025
pioncering work is G2GNN (Wang et al. 2022b which models each graph instance as node in a "ggaph of graphs" and connects instances for data augmentation, thereby improving learning under imbalanced conditions. ImGKB (Tang& & Liang 2023) develops the restricted random walk kernel with the global graph information bottleneck to eftectively capture inter- -raeh supervision for minority-class graphs. Conversely, topological imbalance methods aim to improve performance
on small-scale graphs (Qin et al. 2025). For example, SOLT-GNN (Liu et al. 2022c identifies recurring subgraph patterns to balance learning across graphs of varying sizes. Topolmb (Zhao
et al. 2022) models topological groups and employs modulator to assign greater importance to tail- topology graphs. However, existing methods often struggle in scenarios with complex imbalances.
The development of a unified imbalanced graph learning framework remains an unexplored challenge.
2

---

Published as a conference paper at ICLR 2026
Imbalance Distribution
Graph Perturbation Strategy
Leariable
GNN
Embedding H Embedding
2
Arelhas
Protetype
Balance
Bolanced octivation
Morop edges Meak meder
(V.
Perceive
DOG
Balance
Weaken
Boost
Input Graphs
Local Tapelogy Encoding
Global Tepology Encoding
2N
Dynamic Balanced Prototype(DBP)
K-hhetefysse
Sallent Attention
2
Prototypes
Topk;
x Prototypes
The activation times prototype
(lis
Similarity
2N x TTokks
Graphs
Average frequency Prototypes
Select Topk
2N Graphs
Grophs
T (5H))/4.)
Load Balance Optimization
Prototypes
Mixing Graphs
15+ 6, to, to 3
Figure I: The overall architecture of Unilmb which enhances graph representations by extracting prototypical features in uniform manner.
The detailed discussions of related work graph representation learning, node-level imbalance learning on graph, and edge-level imbalance learning on graph are given in Appendix
PRELIMINARIES
Let 9 = (V,E,, be graph. where is the node set with | nodes, & is the edge set, A E IRIV is the adjacency matrix, X E RR [V]xd is the node feature matrix with d-dimension.
Given a graph set G
,.. GN with N graph instances where the graph instance G; (V8, 89. As. X..). and their corresponding label set Y {y1, 322 .... yN the goal of the graph-level classification task is to learn mapping function JF G - R to map any graph to low-dimensional vector h E R This representation is subsequently passed through classifier to produce the predicted label distribution, yielding the final output prediction for the graph instance.
4 METHODOLOGY
As illustrated in Figure ! Unilmb frist adopts a personalized perturbation strategy with multi-scale graph topology encoding to enhance individual graph instances. It then employs GNNs to learn graph-level representations. Finally, we propose a dynamic balanced prototype mechanism to extract discriminative prototype features, thereby enhancing the model's capability to represent tail graphs and promoting unbiased learning.
4.1
GRAPH REPRESENTATION LEARNING
4.1.1 GRAPH TOPOLOGY ENCODING
Local Topology Encoding We adopt learnable node position learning local information (Dwivedi et al. 2021, which leverages random walks to perceive local structural information around each node, Specifically, y, for graph G,, we first calculate its random walk operator M 1-11 99.
where A Si is the adjacency matrix of G. D is the corresponding degree matrix and D., = Σ, A f.j
With steps random walk, for the node vy of graph G.. we can get its local topology encodings:
LE (M )ái (M [ (M )3,1l e R
(1)
here we adopt a low- complexiy usage of the random walk matrix by considering only the landing probability of the node ny to itself. i.e., (M9:), )..j) And can get the initial node position encodings
3

---

Published as a conference paper at ICLR 2026
LE E RI Subsequently, The position encodings are integrated with GIN for end-to-end updates, enabling the capture of accurate graph topology information.
Global Topology Encoding. The Laplacian matrix of graph contains rich information about the graph's 's global structural properties, such as subgraph frequency and connectivity (Dong et al. 2016). To this end, we introduce a global topology encoding strategy to learn this information. Specifically, for graph G, we first extract the eigenvalues and eigenvectors its laplacian matrix L De: - A which is denoted as [Λ1, A2, >=] and [h1, h2, , E [RIV[| can get the global topology encoding as follows,
Then we
GE = + (le(h,, A,) + ((-,,, A.) RRIV
(2)
where (p.)) and ((-) are permutation- invariant neural networks (e.g., MLP layer) used map the input to a sigher-dimensional hidden space.
ionuutationnnnvariant
4.1.2 PERSONALIZED GRAPH PERTURBATION STRATEGY
Graph perturbation strategies, such as edge drop and feature mask (Wang et al. 2022a can effectively improve the generalization performance of GNNs. However, traditional methods apply uniform perturbations to all graph samples, which can increase the difficulty of learning tail graphs in imbalanced scenarios. To address this issue, we propose personalized graph perturbation strategy that customizes learnable perturbations based on the unique characteristics of each graph.
Edge Drop. For the graph 93 with | nodes and |& edges, we first caleulate the average degree 2|g61 of the graph: ds, = |
R. Then, we create personalized edge dropout distribution B B(ag.) where a. = a (MLP (d)))) E R is learned by one ML.P layer. And we sample the edge drop matrix M where any value m M follows Bernoulli distribution m. B(a ), and m. means that the corresponding edge is dropped. Finally, we apply the learned perturbations to the adjacency matrix: A A M where @ means Hadamard product.
Feature Mask. We use a similar operation to personalized feature mask distribution B(08)) ); for every node vj in graph G; we also sample feature mask mi E {0, ~ ((8))), if m², is equal to O, the corresponding node feature value in X to 0. And mask matrix M and apply it to node feature matrix: X9: 9:
get personalized feature
x @ M
4.1.3 GRAPII REPRESENTATION LEARNING VIA GNN
We leverage the powerful graph representation capabilities of GNNs to generate graph-level represen- 2018 Khoshraftar & An 2024). Specifically, we use an MLP layer to map each graph G. into a high-dimensional space, and the output is then integrated with the global topology encoding as the initial output of the GNN, denoted as X (0).0. |xd,
tations (Xu et al.
For local node position encodings LE E R [x] we also take into one MLP layer, and the output is denoted as LE E glV., xdh This transformed encoding, together with X x. is fed into the GNN layer for graph representation learning. In this work, use GIN (Xu al. as the primary backbone to exemplify our approach and validate the method across several other backbones, as shown in Table di For node By, we can aggregate features from the neighborhood as follows,
X
(
(
x ( MLP(()
(1 (1++)
Σ
.Ve (1,2, ........11. L} LE
(0.9.
LE
(EN(")))
where [.] means the concatenation of two matrices in the column dimension. & is a harmonic coefficient, N(r;)) represents the neighbors of node vy in A of the graph G,. And we also employ an independent GNN layer update the local topology encoding, thereby capturing deeper graph topology information, which is denoted as follows,
LE LE ((+1). MLP(()
(( (1+))LE(.....
+ Σ
LE(0.0:
(4)
NEN
4

---

Published as a conference paper at ICLR 2026
After L graph convolution layers, each node aggregates information from its neighborhoods within radius of up to L hops. Subsequently, a readout function is applied to aggregate the resulting node representations into unified graph representation Hg, for the graph Gi as follows,
g READOUT((X) (1 Uy V).)) E R
Finally, we can obtain the graph representations graphs, denoted as H IR Nxds To enhance the
feature diversity of tail graphs, employ Feature Mixup strategy (Ju et al. 2025 This approach involves randomly rearranging the order of the representation vectors of graph instances, thereby generating new matrix, denoted as H E 2
The details are provided in Appendix C.O.1
4.2 DYNAMIC BALANCED PROTOTYPE BASED ON INFORMATION BOTTLENECK
We further propose a Dynamic Balanced Prototype (DBP) strategy, which unbiasedly extracts representative prototype features from graph data. These prototypes are universally applicable, enabling them effectively enhance the representational capacity of represented tail graph. Specifically, prototypes are defined as set of learnable embeddings, denoted by 181, 82, sk] E KKdd, where the hyperparameter K represents the number of prototypes and satisfies K < N.
for under-
Prototype Perception. Prototypes are perceived from the graph data based on the attention coefficients computed between the prototypes and graph representations, as follows:
Hs- = Softmax (TopK, (SH /V HW, kxd,
where W. Xdd, xd, is the learnable parameters. Next, we utilize prototypes to enhance the representation learning of the graph.
Prototype Balance. We employ an attention mechanism to compute the affinity between each graph and the prototypes, thereby selecting the most relevant prototype enhance its representational capacity. Specifically, this process is formulated as follows:
H === Sigmoid (Top
HS
/v /vda dp
+ +7) Hs.
where H E R2Nxdd represents the enhanced representation. Here, we employ sigmoid function together with a Iearnable vector E R x to generate discriminative similarity scores, rather than producing values concentrated around the mid-range.
During the iterative training process involving two alternating stages, the model learns to extract prototype features. However, these learned prototypes may become biased due imbalanced graph data distributions, which can negatively affect their generalization capability. Ideally, we should reduce the dependency the prototypes on the input data in order to enhance both robustness and generalization performance. To this end, we introduce the information bottleneck theory to guide the learning of more generalized prototype representation.
Theory 1. Information Bottleneck. For the training graphs G, prototype features S, and labels Y. the objective of the information bottleneck theory is expressed as minimizing the following objective function Tishby et al. 2000
min lin $I(S; Y)
(8)
where B is the Lagrange multiplier. We can optimize the mutual information between the hidden feature and the labels I(S: Y) through supervised learning. Thus, the above objective is simplified to minimizing the mutual information between the training graph set G and the prototype feature S, i.e., reducing the redundancy of imbalanced input information. For more details, please refer to Appendix D.I
Proposition 1. Prototype Load Balancing. Let (p1, P2.... ... PKk denote the activation probability distribution of the K prototypes over the training graph set G and U = {ur U2, Uk
'A prototype is considered activated once if its attention coefficient any graph is non-zero
5

---

Published as a conference paper at ICLR 2026
Table 1: Macro-F1 and Micro-Fl scores on class imbalance datasets with extreme imbalance degree under different U distributions. The best results are marked and the runner-ups are underlined
PROTEINS
D&D
NC11
COLLAB
Distribation
Macro-Fi Fl
Micro-FI F1
Macro-FI
Micro-FI
Macro-FI
Micro-Fi FI
Maumo-Fi
Micro-
Assell 2001
73.29
1.58
69.25
Exponeene oatctal [967
71.99
73.51
Poisson Cotow
6
2
41.02 #3.2
71.98
27.67
5.0
2:08
Uniffrmn
Kipaow Niccereeer 2012
70,44
74.30
68.30
4.27
76.34
represent an empirical distribution. Minimizing the mutual information between S and G is equivalent the following objective:
min G(: = Σ p(9)p(s(9) ( p(s) = min (L[(P]|U) R min
1
K
Pk
ux)
2 k=1
G.
where KL means Kullback-Leibler Divergence. The detailed proof is provided in Appendix D.2
Proposition 2. Prototype-balancing Optimization. For 2N graphs, each graph hits TopK2 proto- types, and we have a total hit count K=2NX TopK2. The actual number hits for each prototype where n.= k (1. {1..... KJ. Subsequently, the regularization term minimizes the discrepancy between the actual attention activation ny and the prior distribution ug To achieve smoother and more effective optimization of the attention generation process, introduce an intermediate modulation term ry:
is N {m1, 12,, ... nk).
H Sigmoid (TopK, [HST
IIST/Val /Vdn - -n) + +7) Hs.
(10)
where 7) E RKK directly affects the attention generation process. Nevertheless, the Topk operation non-differentiable practice. To address this, we introduce specific constraint loss LM that iteratively models the prototype activation process during training.
2 N= TopK2
K
2 N TopK2
LM 2
StopGrad(wu= 1/)
(1I)
where (topGrad denotes the stop-gradient operator (Chen & He 2021 which behaves as identity in the forward pass but yields zero in the backward pass, and this a common approach used for non non-differentiable operations. Building on Eq. 11 we further derive the optimization objective for with the learning rate y as follows:
1 1 VVLM 1 yp sgn (n4- 2 N TopK2 *ug),
(12)
To identify appropriate prior distribution, conduct experiments evaluate model performance under four different distributions of U: Zipf uk
Exponential
Uux
Σ e
Poisson
and Uniform (ux k) The complete analysis and more details
R:
are provided in Appendix D.3 The experimental setup described Section As shown in observe that when the empirical distribution U follows uniform distribution, the model achieves the highest accuracy. In other words, whenever the activation count of any prototype exceeds the average level, 17 imposes penalty reduce that prototype's 's activation priority in the subsequent optimization step, thereby enforcing load balancing. This mechanism effectively ensures balanced influence from tail graphs during learning.
Table I
Finally, we input the representations of the N input graph instances (i.e., the first N rows of H into decoder composed of layers MLP for generating the predicted class labels Y.
EXPERIMENTS
5.1 EXPERIMENTAL. SETUP
Datasets. We evaluate the model's performance on 19 datasets under class imbalance and topological imbalance scenarios. Following existing protocols Qin al. 2024 the datasets processed

---

Published as a conference paper at ICLR 2026
into three levels of imbalance (low, medium, extreme). Due space limitations, we primarily report results on six highly imbalanced datasets in the main text: D&D Shervashidze et al. 2011 NCIL Waae et al. 2008 COLLAB (Leskovec et al. 2005), PROTEINS, REDDIT-B, and IMDB-MULTI Yanardag & Vishwanathan 2013 We also introduce a large-scale dataset, Airgraph (described below), and three 3ID conformational graph datasets MoleculeNet (Wu al. 2018): BBBP for blood-brain barrier permeability prediction, BACE for ß-secretase inhibition prediction, and HIV for antiviral activity prediction.
Large-Scale AirGraph Dataset. comprises hourly PM2.5 measurements from the China National Environmental Monitoring Center collected at 1,341 stations across 28 provinces in mannland China between 2021 and 2023. For each province the data are represented as graph in which monitoring stations are nodes and stations are connected based on distance. Each node encodes the preceding 24 hours of PM2.5 readings, and each graph is labeled by the spatially averaged PM2.5 over the subsequent 24 hours, categorized into three levels: low, medium, and high pollution. With 30,660 graph instances, AirGraph is substantially larger than most existing graph classification datasets and exhibits a natural class imbalance: the high pollution category accounts for only 6.86% of instances. More details are provided in Appendix G
Baselines. We compare our model against four categories of baselines, encompassing a total of 23 representative baselines: O Classic Graph Classification Models: GIN (Xu et al. GCN (Kipf &
Welling 2016). GraphSAGE (Hamilton et al. 2017 InfoCrraph (Sun et al. 2019 and GraphC You
al. 2020g g @ Graph Transformer Models (G7's:: GraphGPS (Rampášek al. 2022), Exphormer
(Shizzad et al. 2023 and Graph-Mamba Wang et al. 2024b); @ 3D Graph Models: SchNet
(Schutt et al. 2017), DimeNet (Gasteiger et al. 20200, and SphereNet (Liu et al. [002bb; 1
Class-imbalance Methods; upsampling (Kubat et al. 1997 and reweighting strategies (Yuan & Ma
applied to various backbones, well as G'GNI with edge and node masking. ImGKB TTag
& Liang 2023), and DataDec (Zhang et al. 2023);
Topology-imbalance Methods: SOLT-GNN
ILu et al. 2022c), ImbGNN at.
2024 and Topolmb (Zhao al. 2022).
Experiment Setting, We implement our proposed model on an 40GB NVIDIA A100 GPU with Pytorch. For the iabalaneblen graph classification task, split the dataset into 25% for training and 25% for validation, and the remaining data is used for testing. For the topological-imbalnnee graph classification task, split the dataset into 10% for training and 10% for validation, and the remaining data is used for testing. We use the Adam optimizer (Kingma & Ba 2014 with learning rate 0.001. All experiments are repeated 20 times, and we report the average and standard deviation. For the REDDIT-B. IMDB-MULTI, and COLLAB datasets, we employ one-hot encodings of node degrees as node features (unn et al. 2019 You et al. 2020 We use widely metrics in the graph imbalance tasks including Macro-Fi and Micro-FI (Xia et al. 2014 Lipton et al. 2014). And is set to 0.001. The number of prototypes set to { [16, 16, 24, 24, 24, 32) on Six datasets. We stack
L= 5 layers of GNN. Other key hyperparameters of each dataset are shown in Table K
Table 2: Performance on class imbalance datasets with extreme imbalance degree. The best results are marked and the runner-ups are underlined We report the average and standard deviation over 20 runs Numbers marked with indicate that the improvement is statistically significant compared with the best baseline (Wilcoxon Signed-Rank Test with p-value 0.05).
PHOTEENS
D&D
NCI
RESDFLE
COLLAR
INDDE-BULTT
Matte
Turthuse
1633
51.00
57 10-20
53 27
Brghht
33
3510
31A2
00.22
11:30
15
18 441
13.5 3331
12.0 1387
11654775
12:11 513
12.33
2100 12.86
1430
25.44, 1578
Uls
Gay8-MMrri
31.12
780
61241122 11.36
1494
12.77 1246
11/0 14.1
42.50 11:3
15.55 2437
25.97 1247
41 13.4
11.00 1835
SURPE 518
10:00:331
23.57
15.11 RUM
10.08 03.0
413
12:20+555
21.74
18
16.00
43.30
11 33
+3.07
13 1110
119
0079 SJ8
30
62.22 2/6
54 2-14
31 44.30
16/00330
G'ONN
mave sky
1915-191
N6 35-
15 15
11
13 47
14
15 +1.5
15 15.50
55 15
11 71
57.33 41
1114
16:04011
13
1334
74.38
43.
11.52
3454
Treviess
4M91
12491
9499
4HN'
7.891
11
34191
17,8%
14.%1
4182%
1144%1

---

Published as a conference paper at ICLR 2026
Table 3: Performance comparison on topological imbalance datasets with extreme imbalance degree.
PROTTRINS
NC31
REHOrEA
COLLAB
EMPR-MERTI
Method
Ml
Mi
Macr:FF1
Mixo-FF
Micro-FF
MinuEl
NacooF]
Micr=FF
Maats
Mi
53.48 1248--153
50 00 10
<7 5.51
10:00 -1:0
(0--113 30
61 58
66 600 3:37
67.41 1:3
61.00 2110
70 -011
14:73 -1
4.11
2.11
53 455
1:33
35.35
36-42 LiB
56
Grspaciis
55.50
6335-110
03.85 -10
13.06-35: 3.51
2::00-3::
3111-83:
26
16.97
24-64
58 68:06=141) 141
12.00+316
13.65+977)
11.18-953
6310-3333
03.00 m316
12:47
63:0-144
13.0-540
20.37 1130
1743+
13.31+603
Tepoleh
64.19
10.59
07.20
2.75
06.97
6:33
18:21 m100
67 172
01 16.6-.111 113
68:19.6:1
52.31 m/
22.85 500
13.91 1:5
SOLIFINN
236
39.50
54397 3.24
0:309
02 1:44
80.34
2110
Unilib
194 194
11:99
110
76
11
76.33
06.90 4:0
65.76 671
17:10F 1535
76332
15:58 110
18:05 -1
5.2 PERFORMANCE ANALYSIS ON IMBALANCE GRAPH CLASSIFICATION
Owing to space limitations, we present the experiment results on all datasets under medium and low imbalance degrees in Appendix 1.1.1 and 13.11
O Class Imbalance Graph Classification. As shown in Table 2 classic graph representation models perform the worst due to their inability to address class imbalance. While Graph Transformers (GTs) augment GNN with global attention mechanisms capture long-range dependencies, they remain ineffective in mitigating the issue of class label imbalance. Upsampling methods require resampling tail elasses, which may lead to overfitting on tail features and result in loss of generalization. Topolmb and ImGKB exhibit highly unstable performance, performing even worse than naive models on most datasets (Qin et al. 2024. DataDec inadvertently introduces data leakage by using test samples to fine-tune the SVM classifier. To ensure a fair evaluation, we restrict all fine-tuning to the training set. Under this setting, DataDec shows significantly degraded performance. G2GNN and ImbGNN utilize the G2G framework for data augmentation. However, they underperform compared to simple upsampling baselines on multi-labe datasets such as COLLAB. Our model achieves the best performance across all metrics.
Topological Imbalance Graph Classification. As shown in Table $ GIN consistently achieves the worst performance across all topologically imbalanced datasets. GraphCI. further improves the performance through contrastive lcarning with graph augmentations, yielding significant gains over GIN. Graph Transformer models, including GraphGPS, Exphormer, and Graph-Mamba, build upon GNN by integrating global attention mechanisms, which help address topological imbalance. Topolmb alleviates substructure imbalance by incorporating topology-awe aware modeling and instance weight adjustment. ImbGNN constructs an enhanced G² based on topology coaelgy-friendly random walks, which helps mitigate the impact of aopology-imbllanee in the data. Our Unilmb integrates personalized graph perturbation strategy, a comprehensive graph structure encoding module, and the DBP module. The DBP module effectively extracts prototype features, significantly enhancing the representation capability for graphs of various scales. Our model achieves competitive performance.
5.3 COMPREHENSIVE EVALUATION OF UNIIMB: APPLICABILITY AND SCALABILITY
Applicability with various backbones. We evaluated various GNN variants and Graph Trans- formers as backbone networks under class-imbalanced settings. The results in Table 4 show that all backbone networks achieved significant performance improvements when integrated with Unilmb. This validates the plug-and-play applicability of our method, demonstrating its ability to significantly enhance the imbalanced graph classification performance of GNNs. For experiments on the applicable scenarios of topological imbalance, see the Appendix Li
Scalability on the large graph dataset. Figure illustrates the prediction accuracy (Macro-F1) and complexity (training time and memory usage) of several advanced models on the large-scale dataset AirGraph, which is naturally imbalanced. We observe that Unilmb consistently outperforms strong baseline models, achieving superior performance in predicting air pollution levels. This provides valuable insights for environmental protection and human health. Furthermore, Unilmb demonstrates good computational efficiency. Compared with methods specifically designed for imbalance handling, such as Topolmb, ImbGNN, and G2GNN, Unilmb achieves more significant improvements in both performance and computational complexity.
8

---

Published as a conference paper at ICLR 2026
Table 4: Graph classification performance of various back- bones combined with Unilmb under the Class Imbalance setting with extreme imbalance degree.
BIDI SEE
(Ours)
ITBI SIB
IMDI-MILITE
Macmeii
Mime-11
Marm-n
Meno-11
Mane-D
Mims-11
Explar
Mettiory-Digge
11.38.
31.19
3612
GUN Oare
GUN
4188 441s
2:30
7623
1100
10.25+1.0
38.52
1.
Go Graph Manih
ImbGNN
GrighSage
608-3m
1400
DataDee
Tapolmib
Grapheing Our
42.09
19.36 600
10.39 8258
85880 6323
32.29 23
25.00
GraphGiSS Ours
48.17
68.37
1687
4300
Psphoener
3081 41
79 1:
Training Spend (slepoch)
Explurmer
GimpoMammta
Criph-Maats
42:38
7.0
12 mus
71:83
$2.38
16.00 +557
15.0
41.02
Figure 2: Model performance and effi- ciency comparison on AirGraph dataset.
5.4 PERFORMANCE ON INTERTWINED CLASS AND TOPOLOGY IMBALANCES SCENARIO
We further evaluated the effectiveness of various models in complex scenarios involving class and topological imbalance. Details of the data processing are provided in Appendix G.1.5 Additionally. AirGraph also exhibits combined case of categorical and topological imbalance. Figures and compare our method with several advanced models. We found that methods designed to address topological imbalance generally perform worse than those targeting class imbalance. This because the distribution of categorical features has crucial impact the decision boundary, directly influencing model performance. Our method achieved the best overall performance.
TopeDa
D0C
Bulle
Figure 3: Macro-FI (left) and Micro-F1 (right) Figure 4: on intertwined imbalance datasets.
Ablation experiments on class imbalance (upper) and topological imbalance (lower).
5.5 ABLATION STUDY
To evaluate the effectiveness of the key components of the model, we create five variants: (1) wwo ImMix' means that we remove imbalanced feature mixup strategy; (2) 'w/o TopoEnc' means that we remove graph topology encoding; (3) "w/o Pertu" means that the graph perturbation strategy strategy is completely removed; (4) 'w/o DBP means that the dynamic balanced prototype strategy is removed; (5) 'w/o Ba1Opt' means that the we remove the load balancing optimization.
As shown in Figure the w/o TopoEnc variant achieves notably lower accuracy, indicating that incorporating structural information enables the model to better capture graph-level features. Among the proposed bad-balanced prototype learn-
all ablated versions, the w/o DBP variant-withe ing exhibits the worst performance, demonstrating that explicitly modeling prototype features is critical for handling imbalanced graph data. The w/o BalOpt variant also suffers performance drop, particularly on class-imbalanced datasets, which further confirms that the proposed had-bblanning strategy facilitates the extraction of generalizable prototypes. Overall, the inferior performance of all ablated variants compared to Unilmb provides strong evidence for the effectiveness and necessity of each component our framework. For more ablation studies, see the Appendix and 1.5
5.6 HYpERPARAMETER SENSITIVITY ANALYSIS
We further evaluate the model's sensitivity to the number of prototypes. As shown in Figure 5 observe that model performance first increases and then decreases with respect to K. When K is smaller than the optimal value, is insufficient to capture enough discriminative prototype information. On the other hand, when K exceeds the optimal value, having too many prototypes

---

Published as a conference paper at ICLR 2026
prevents the model from focusing representative ones, leading performance degradation. For sensitivity studies on TopK₁ and TopK2, see the Appendix L.6
NOH
MO
COCLAR
(a) GIN
(b) ImbGNN
Notber aiTrooope
(e) Before DBP
After DBP
Figure 5: Sensitivity study class imbalance (upper) and topological imbalance (lower).
Figure 6: Representation visualiza- tion on the PROTEINS dataset with extreme imbalance degree.
5.7 EVALUATION ON 3D CONFORMATION DATASETS
Beyond conventional graph datasets, we introduce 3D conformational graph datasets from Molecu- leNet: BBBP for blood-brain barrier permeability prediction, BACE for 3-secretase inhibition activity prediction, and HIV for antiviral activity prediction. To accommodate this graph structure, we incorpo- rate specific models, including SchNet, DimeNet and SphereNet, which serve backbone networks for graph imbalance methods. We follow established protocols for handling extremly imbalanced class distributions. And we select SphereNet as backbone of G2GNN, Topolmb, and ImbGNN. For our model, we retain only the core DBP module. As shown in Table S the proposed method achieves excellent performance on these 3D Conformation datasets, validating its applicability to domains with complex structural dependencies.
Table Performance comparison on class imbalance datasets with extreme imbalance degree.
HIV
BRBP
BACE
Model
Macro-F1
Micro-FI
Macro-Fi
Micro-FI
Macro-Fi
Micro-FI
SchNer
29 1.16
96.57 +3:30
43.48
74 7m
57.09 4118
60.00 1.5
SpbercNet
50.68 1.30
95.32
=9.27
18.62
76.27 2612
45.25
=2.12=528
DimeNe:
23.13
a362
2138
2436
=2.8
Topolmh
42.40
+3.04
-5.05-385
-0.1
ImbGNN
56 2232
97.63
2200
63.92
82.74
56.622243
63.50-336
Ours
60.29 1:18
08.95 2m
76:34
86,00 1735
62.30 1W
63.34-1W
5.8 GRAPH REPRESENTATION VISUALIZATION
We extract the node-level representations from the PROTEINS dataset using GIN and ImbGNN, as well as the representations before and after applying the DBP module. These are then visualized in Figure The representations generated by the naive GIN (sub-figure (a)) for imbalanced graphs are highly entangled and poorly separated, making it difficult for the classifier to distinguish between classes. While ImbGNN a specialized method for imbalanced learning ssows relatively better separation, many samples still exhibit overlapping feature distributions. We observe that, after incor- porating the Dynamic Balanced Prototype mechanism, the learned label-wise representations become significantly more discrininative. This enhancement greatly benefits the classifier in distinguishing between different classes and improves overall performance.
CONCLUSION
In this paper, we propose unified framework for addressing imbalance graph classification task, effectively tackling both class imbalance and topological imbalance. Our core contribution lies in the introduction of a dynamic balanced prototype strategy, which extracts prototype features from graph data to enhance robust representation learning under imbalanced scenarios. We further impose a prototype- seootype-balancing optimization strategy to encourage balanced activation across all prototypes. Additionally, encoding, graph perturbation and feature mixup to comprehensively improve model performance. Extensive experiments over 19 datasets demonstrate the impressive effectiveness of framework.
incorporate several advanced techniques, including graph topology
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
This work addresses graph learning and does not involve individual-level sensitive data or direct interaction with users. We do not identify any inherent ethical risks in the method itself. Nevertheless, downstream applications sensitive domains may introduce harms, and practitioners should assess ethical implications before deployment.
REPRODUCIBILIIY STATEMENT
The main text formalizes the model architecture with equations. The Appendix contains full imple- mentation details, including dataset descriptions, evaluation metrics, model architectures, hyperpa- rameters, training procedures, and experiment configurations. The code is provided at anonymous link.
ACKNOWLEDGMENT
This paper partially supported by the National Natural Science Foundation of China (No. 12227901) and the Natural Science Foundation of Jiangsu Province (BK20250482). The Al-driven experiments, simulations and model training were performed on the robotic Al-Scientist platform Chinese Academy of Science.
REFERENCES
Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.004411 2016.
Emily Alsentzer, Samuel Finlayson, Michelle Li, and Marinka Zitnik. Subgraph neural networks. Advances in Neural Information Processing Systems, 33:8017-8029, 2020.
Robert L Axtell. Zipf distribution of us firm sizes. science, 293(5536):1818-1888 2001.
Deng Cai and Wai Lam. Graph transformer for graph-to-sequence learning. In Proceedings of the
AAAI conference on artificial intelligence, volume 34, pp. 7464-7471, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 15750-15758, 2021.
Xueqi Cheng, Yu Wang, Yunchao Liu, Yuying Zhao, Charu C Aggarwal, and Tyler Derr. Edge classification graphs: New directions in topological imbalance. In Proceedings the Eighteenth ACM International Conference on Web Search and Data Mining, pp. 392-400, 2025.
Prem C Consul and Gaurav C Jain. A generalization the poisson distribution. Technometrics, 15 (4):791-799, 1973.
Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999.
Asim Kumar Debnath, Rosa L Lopez de Compadre, Gargi Debnath, Alan J Shusterman, and Corwin Hansch. Structure-activio relationship mutagenic aromatic and heteroaromatic nitro compounds. correlation with molecular orbital energies and hydrophobicity. Journal of medicinal chemistry, 34 (2):786-797, 1991.
Tyler Derr, Yao Ma, and Jiliang Tang. Signed graph convolutional networks. In 2018 IEEE international conference data mining (ICDM), pp. 929 934.934.. IEEE, 2018.
Xiaowen Dong, Dorina Thanou, Pascal Frossard, and Pierre Vandergheynst. Learning laplacian matrix in smooth graph signal representations. IEEE Transactions on Signal Processing, 64(23)) 6160-6173, 2016.
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. arXiv preprint arXiv:21 10.0077. 2021.
11

---

Published as a conference paper at ICLR 2026
Shaohua Fan, Xiao Wang, Chuan Shi, Peng Cui, and Bai Wang. Generalizing graph neural networks on out-of-distrbuttio graphs. IEEE transactions on pattern analysis and machine intelligence, 46 (L):322-337, 2023.
Xingcheng Fu, Yuecen Wei, Qingyun Sun, Haonan Yuan, Jia Wu, Hao Peng, and Jianxin Li. Hy- perbolic geometric graph representation learning for hierarchy-imbalance imbalance node classification. In Proceedings of the ACM Web Conference 2023, pp. 460-468, 2023.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.00 123, 2020.
Zhaojie Gong, Yijun Duan, and Qiang Ma. Edge classification on imbalanced multi-relational graphs. In International Conference Advanced Data Mining and Applications, pp. 257-264. Springer, 2025.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs Advances in neural information processing systems, 30, 2017.
William L Hamilton. Graph representation learning. Morgan & Claypool Publishers, 2020.
Mingxuan Ju. Tong Zhao, Wenhao Yu, Neil Shah, and Yanfang Ye. Graphpatcher: mitigating degree bias for graph neural networks via test-time augmentation. Advances in Neural Information Processing Systems, 36:55785-55800 2023.
Wei Ju. Zhengyang Mao, Siyu Yi, Yifang Qin, Yiyang Gu, Zhiping Xiao, Jianhao Shen, Ziyue Qiao, and Ming Zhang. Cluster- guided contrastive -analannbledceed graph classification. In Proceedings of the AAAI Conference Artificial Intelligence, volume 39, pp. 1192--19932 2025.
Z
Shima Khoshraftar and Aijun An. A survey on graph representation learning methods. ACM Transactions on Intelligent Systems and Technology, 15(1):1-55, 2024.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 2014.
Thomas N Kipf and Max Welling, Semi-supervised classification with graph convolutional networks. arXiv preprint 1009.:660.0000000 2016.
Miroslav Kubat, Stan Matwin, et al. Addressing the curse of imbalanced training sets: one-sided selection. In leml, volume 97, pp. 179. Citeseer, 1997.
Lauwerens Kuipers and Harald Niederreiter. Uniform distribution of sequences. Courier Corporation, 2012.
Jure Leskovec, Jon Kleinberg, and Christos Faloutsos. Graphs over time: densification laws, shrinking diameters and possible explanations. In Proceedings of the eleventh ACM SIGKDD international conference on Knowledge discovery in data mining, pp. 177-187, 2005.
Xiao Li, Li Sun, Mengjie Ling, and Yan Peng. A survey of graph neural network based recommenda- tion in social networks. Neurocomputing, 549:12644, 126441, 2023.
Zachary C Lipton, Charles Elkan, and Balakrishnan Naryanaswamy. Optimal thresholding of classifiers to maximize fl measure, In Machine Learning and Knowledge Discovery Databases: European Conference, ECML PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings. Part Il 14, pp. 225-239. Springer, 2014.
Jiaying Liu, Feng Xia, Xu Feng, Jing Ren, and Huan Liu. Deep graph learning for anomalous citation detection. IEEE Transactions Neural Networks and Learning Systems, 33(6):2543 (566)2243-25577,, 2022a.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference on Learning Representations (ICLR), 2022b.
Zemin Liu, Wentao Zhang, Yuan Fang, Xinming Zhang, and Steven CH Hoi. Towards locality-aware meta-learning of tail node embeddings on networks. In Proceedings of the 29th ACM international conference on information & knowledge management, pp. 975-984 2020.
12

---

Published as a conference paper at ICLR 2026
Zemin Liu, Trung- Kien Nguyen, and Yuan Fang. Tail-gnn: Tail-node graph neural networks. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining, pp- 1109-1119, 2021.
Zemin Liu, Qiheng Mao, Chenghao Liu, Yuan Fang, and Jianling Sun. On size-oriented long-tailed graph classification of graph neural networks. In Proceedings of the ACM web conference 2022,
pp. 1506-1516, 2022c.
Zemin Liu, Yuan Li, Nan Chen, Qian Wang, Bryan Hooi, and Bingsheng He. A survey of imbalanced learning on graphs: Problems, techniques, and future directions. IEEE Transactions on Knowledge and Data Engineering. 2025.
Yihong Ma, Yijun Tian, Nuno Moniz, and Nitesh V Chawla. Class-imbalanced learning on graphs: A survey. ACM Computing Surveys, 57(8);1-16, 2025.
Zhengyang Mao, Wei Ju, Siyu Yi, Yifan Wang, Zhiping Xiao, Qingqing Long, Nan Yin, Xinwang Liu, and Ming Zhang. Learning soowledge-diverss experts for long-tailed graph classification. ACM Transactions on Knowledge Discovery from Data, 19(2):1-24, 2025.
Albert W Marshall and Ingram Olkin. A multivariate exponential distribution. Journal of the American Statistical Association, 62(317):30 1967.
Christopher Morris, Nils M Kriege, Kristian Kersting, and Petra Mutzel. Faster kemels for graphs with continuous attributes via hashing. In 2016 IEEE 16th International Conference on Data Mining (ICDM), pp. 1095-1100. IEEE, 2016.
Dai Quoe Nguyen, Tu Dinh Nguyen, and Dinh Phung. Universal graph transformer self celf-attention networks. In Companion Proceedings of the Web Conference 2022, pp. 193-196, 2022.
Dung Nguyen and Anil Vullikanti. Differentially private densest subgraph detection. In International Conference on Machine Learning, pp. 1140-8151. PMLR, 2021.
Emmanuel Noutahi, Dominique Beaini, Julien Horwood, Sébastien Giguère, and Prudencio Tossou. Towards interpretable sparse graph representation learning with laplacian pooling. arXiv preprint arXiv: 1577, 2019.
Francesco Orsini, Paolo Frasconi, and Luc De Raedt. Graph invariant kernels. In Proceedings of the twenty-oourth international joint conference on artificial intelligence, volume 2015, pp. 3756-3762. IJCAI-INT JOINT CONF ARTIF INTELL,,
Jiawen Qin, Haonan Yuan, Qingyun Sun, Lyujin Xu, Jiaqi Yuan, Pengfeng Huang, Zhaonan Wang, Xingcheng Fu, Hao Peng, Jianxin Li, et al. Igl-bench: Establishing the comprehensive benchmark for imbalanced graph learning. arXiv preprint arXiv:2406.09877 2024.
Jiawen Qin. Pengfeng Huang, Qingyun Sun, Cheng Ji, Xingcheng Fu, and Jianxin Li. Graph size- imbalanced learning with energy -grrded structural smoothing. In Proceedings of the Eighteenth ACM International Conference on Web Search and Data Mining, pp. 457-465, 2025.
Liang Qu, Huaisheng Zhu, Ruiqi Zheng, Yuhui Shi, and Hongzhi Yin. Imgagn: Imbalanced network embedding via generative adversarial graph networks. In Proceedings of the 27th ACM SIGKDD conference on knowledge discovery & data mining, pp. 139001339, 2021.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Newral Information Processing Systems, 35:14501-14515, 2022.
Kaspar Riesen and Horst Bunke. Iam graph database repository for graph based pattern recognition and machine learning. In Structural, Syntactic, and Statistical Pattern Recognition: Joint IAPR International Workshop, SSPR & SPR 2008, Orlando. USA, December 4-6, 2008, Proceedings, pp. 287-297. Springer, 2008
Robert Sanders. The pareto principle: its use and abuse. Journal of Services Marketing, 1(2):37-40, 1987.
13

---

Published as a conference paper at ICLR 2026
Ida Schomburg, Antje Chang. Christian Ebeling, Marion Gremse, Christian Heldt, Gregor Huhn, and Dietmar Schomburg. Brenda, the enzyme database: updates and major new developments. Nucleic acids research, 32(supp 1):D431-D433 2004.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Blaus-Robert Müller. Schnet: continuous-filee convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mehlhorn, and Karsten M Borgwardt. Weisfeiler-lehman graph kernels. Journal of Machine Learning Research, 12(9), 2011.
Hamed Shirzad. Ameya Velingker, Balaji Venkatachalam, Danica Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. In International Conference on Machine Learning, pp. 31633-33366 PMIR, 2023.
Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv: metp::773000000 2017.
Chengcheng Sun, Cheng Zhai, Qihan Feng, Xiaobin Rui, and Zhixiao Wang. Heterogeneous graph neural network with relation-aware label propagation for unbalanced node classification. Physica A: Statistical Mechanics and its Applications, 660:13036) 2025.
Fan-Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semi- respphhstttrpretsettaaattooo
supervised graph-level representation learning via mutual information maximization. arXiv preprint 12Xiv:1908.01000 2019.
Jeffrey J Sutherland, Lee A O'brien, and Donald F Weaver. Spline-fitting with genetic algorithm: A method for developing classification structure- activity relationships. Journal of chemical information and computer sciences, 1306):1906-1915 2003.
Hui Tang and Xun Liang. Where to find fascinating inter-graph supervision: Imbalanced graph classification with kernel information bottleneck. In Proceedings of the 31st ACM International Conference on Multimedia, pp. 3240-3249, 2023.
Bishal Thapaliya, Anh Nguyen, Yao Lu, Tian Xic, Igor Grudetskyi, Fudong Lin, Antonios Valkanas, Jingyu Liu, Deepayan Chakraborty, and Bilel Fehri. Ecgn: A cluster-aware approach to graph neural networks for imbalanced classification. arXiv preprint arXiv:2410.11765, 2024.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pp. 1-5. Icce, 2015.
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physicx/0004057, 2000.
Hannu Toivonen, Ashwin Srinivasan, Ross D King, Stefan Kramer, and Christoph Helma. Statistical evaluation of the predictive toxicology challenge 2000-2001. Bioinformatics, 19(10): 1ath-118111999 2003.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv: 10903, 2017.
Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio, et al. Graph attention networks, stat, 1050(20):10 18-5500000-88550 2017.
Hippolyte Verdier, Maxime Duval, François Laurent, Alhassan Cassé, Christian I. Vestergaard, and Jean-Baptiste Masson. Learning physical properties of anomalous random walks using graph neural networks. Journal of Physics A: Mathematical and Theoretical, 54(23):234001, 2021.
Nikil Wale, Ian A Watson, and George Karypis. Comparison of descriptor spaces for chemical compound retrieval and classification. Knowledge and Information Systems, 14:347-375, 2008
Binwu Wang, Yudong Zhang, Xu Wang, Pengkun Wang, Zhengyang Zhou, Lei Bai, and Yang Wang. Pattern expansion and consolidation evolving graphs for continual traffic prediction. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 22322, 2023.
14

---

Published as a conference paper at ICLR 2026
Binwu Wang. Jiaming Ma, Pengkun Wang, Xu Wang, Yudong Zhang, Zhengyang Zhou, and Yang Wang. Stone: A spatio-temporal ood learning framework kills both spatial and temporal shifts. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 2948-2959) 2024a.
Chloe Wang, Oleksii Tsepa, Jun Ma, and Bo Wang. Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv preprint (rXi::2402.00789, 2024b.
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. loxiliary-loss-fii free load balancing strategy for mixture-of-expeeee arXiv preprint arXiv:2408. 5664, 2024c.
Yu Wang and Tyler Derr. Degree-related bias in link prediction. In 2022 IEEE International Conference on Data Mining Workshops (ICDMW), pp. 757-758. IEEE, 2022.
Yu Wang, Charu Aggarwal, and Tyler Derr. Distance-wise prototypical graph neural network in node imbalance classification. arXiv preprint arXiv:2110.12033, 2021.
Yu Wang, Wei Jin, and Tyler Dem. Graph neural networks: Self-supervised learning. In Graph Neural Networks: Foundations, Frontiers, and Applications, pp. 001399140 Springer, 2022a.
Yu Wang, Yuying Zhao, Neil Shah, and Tyler Derr. Imbalanced graph classification via graph-of- graph neural networks. In Proceedings of the 3 ACM international conference on information & knowledge management, pp. 2067-2076, 2022b.
Lanning Wei, Huan Zhao, Zhiqiang He, and Quamming Yao. Neural architecture search for gnn-based graph classification. ACM Transactions on Information Systems, 42(1):1-29, 2023.
Shiwen Wu, Fei Sun, Wentao Zhang, Xu Xie, and Bin Cui. Graph neural networks in recommender systems: survey. ACM Computing Surveys, 55(5):1-37, 2022.
Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science, 9()):133-530 2018.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengyi Zhang, and Philip S Yu. A comprehensive survey graph neural networks. IEEE transactions neural networks and learning systems, (2(1):4-20 2020.
Rongkai Xia, Yan Pan, Lei Du, and Jian Yin. Robust multi-view spectral clustering via low-rank and sparse decomposition. In Proceedings of the AAAI conference on artificial intelligence, volume 28, 2014.
Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arKiv preprint arXiv: 1810000800 2018.
Wei Xu, Pengkun Wang, Zhe Zhao, Binwu Wang. Xu Wang, and Yang Wang. When imbalance meets imbalance: Structure-driver learning for imbalanced graph classification. In Proceedings of the ACM Web Conference 2024, pp. 905-913, 2024.
Pinar Yanardag and S Vishwanathan. Deep graph kernels Proceedings of the 21th aem sigkdd international conference on knowledge discovery and data mining, 1365-1374. ACM. New York, 2015.
Cheng Yang, Chunchen Wang, Yuanfu Lu, Xumeng Gong, Chuan Shi, Wei Wang, and Xu Zhang. Few-shot link prediction in dynamic networks. In Proceedings of the fifteenth ACM international conference on web search and data mining, pp. 1245-1255, 2022.
Kuo Yang, Zhengyang Zhou, Wei Sun, Pengkun Wang, Xu Wang. and Yang Wang. Extract and refine: Finding a support subgraph set for graph representation. In Proceedings of the 29th ACM SIGKDD conference on knowledge discovery and data mining, pp. 2953-99666 2023.
15

---

Published as a conference paper at ICLR 2026
Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877-28888 2021.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. Advances in neural information processing systems, 33: 5812-5823, 2020.
Bo Yuan and Xiaoli Ma. Sampling+ reweighting: Boosting the performance of adaboost on imbal- anced datasets. In The 2012 international joint conference neural networks (IJCNN) pp. 1-6. IEEE, 2012.
Haonan Yuan, Qingyun Sun, Junhua Shi, Xingcheng Fu, Bryan Hooi, Jianxin Li, and Philip S Yu. Graver: Generative graph vocabularies for robust graph foundation models fine-tuning. arXiv preprint arXiv:2511.05592 2025a.
Haonan Yuan, Qingyun Sun, Zhaonan Wang, Xingcheng Fu, Cheng Ji, Yongjian Wang, Bo Jin, and Jianxin Li. Dg-mamba: Robust and efficient dynamic graph structure learning with selective state space models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 22272-22280, 2025b.
Seongjun Yun, Minbyul Jeong, Rachyun Kim, Jaewoo Kang, and Hyunwoo Kim. Graph transformer networks. Advances in neural information processing systems, 32, 2019.
Sukwon Yun, Kibum Kim, Kanghoon Yoon, and Chanyoung Park. Lte4g: Long-tail experts for graph neural networks. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pp. 2434-2443, 2022.
Chunhui Zhang, Chao Huang, Yijun Tian, Qianlong Wen, Zhongyu Ouyang, Youhuan Li, Yanfang Ye, and Chuxu Zhang. When sparsity meets contrastive models: Less graph data can bring better class- balanced representations. In International Conference on Machine Learning, pp. 41113-41155 PMLR, 2023.
Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learning architecture for graph classification. In Proceedings the AAAI conference on artificial intelligence, volume 32, 2018.
Lingxiao Zhao, Wei Jin, Leman Akoglu, and Neil Shah. From stars to subgraphs: Uplifting any gnn with local structure awareness. arXiv preprint (xXiv22/000037:: 2021a.
Tianxiang Zhao, Xiang Zhang, and Suhang Wang. Graphsmote: Imbalanced node classification graphs with graph neural networks. In Proceedings of the 14th ACM international conference on web search and data mining, pp. 833-841, 2021b.
Tianxiang Zhao, Dongsheng Luo, Xiang Zhang, and Suhang Wang. Topoimb: Toward topology-level imbalance in learning from graphs. In Learning on Graphs Conference, pp. 37-1. PMLR, 2022.
16

---

Published as a conference paper at ICLR 2026
Appendix
Table of Contents
Use of LLM
18
Related Work
18
B.
Graph Representation Learning
18
B.2
Node-leve Imbalance Learning on Graph
18
Edge-level Imbalance Learning on Graph
8.4
Subgraph Extraction Graph
Method
19
Theory Prool
19
D.I
Explanation for Theory
19
D.2
Proof for Proposition
20
Prototype-balanchi Optimization
22
D.4
Comparison Dilferent Distributions
24
E
Algorithm Pseudocode
24
Delinition of Graph Class Imbalance and Topological Imbalance
25
Class Imbalane
25
Topological Imbalance
25
Dataset
26
Introduction
26
Experiment Setting
30
Hyperparameter Setting
30
Eva Evaluation Metries
30
Additional Experimental Results
31
1.1
Class imbalance Performance Analysis
................................
31
Multi-class Datasel
32
I.3
Topological imbalance Performance Analysis
33
4
Ablation experiment with Micro-Fi score
35
1.5
Additional Experiment of Graph Perturbation Strategy
35
1.6
Hyperparameter Experiment
36
1.7
Applicability analysis Unilmb
36
1.8
Applicabliity analysis of Dynamic Balanced Prototype
37
199
Craph Structure Visualization
38
imitation Discussion and Future Work
38
17

---

Published as a conference paper at ICLR 2026
A USE OF LLM
We employed large language models solely to refine wording and improve readability and fluency. LLM assistance was limited to language editing and did not alter model design, data processing, or experimental results.
B RELATED WORK
B.1
GRAPH REPRESENTATION LEARNING
Graphs serve as fundamental data structures for modeling complex relationships among entities Yuan
al.... 2025a Wu et al. 2022 Derr al. 2018), with broad applications in social networks (Li et al.
energy systems Yuan et al. 2025b). transportation (Wang et al. 2023 2024a), and beyond. A central task in graph machine learning graph representation learning, which aims to extract latent information from structured graph data and encode it into vectorized representations that support downstream reasoning (Hamilton 2020 Early work in this area primarily focused on spectral methods based on the Laplacian operator Noutahi et al. 2019) and random walks (Verdier et al. 2021). In recent years, the success of deep learning on Euclidean data has spurred the development of GNNs, which leverage a message-passing paradigm to effectively integrate topological structure with node features, yielding richer and more informative representations. Several GNN variants have since been proposed to address distinct challenges: graph attention network (Velickovic et al. 2017 introduces an attention mechanism to dynamically weight neighbor contributions GraphSAGE employs neighborhood sampling for scalable learning on large graphs (Hamilton et al. 2017); and GIN el 2018) enhances expressiveness by refining the aggregation function during message passing. However, most existing approaches focus node-level embeddings, which may be insufficient for capturing higher-order structural patterns. Moreover, their use of topological information is often implicit and coarse. Recently, inspired by the transformative impact of Transformer architectures in NLP and vision, researchers have begun adapting them to graph representation learning (Yun al. 2019 Cai & Lam| 2020). For example, Graphormer (Ying et 2021) uses single dense self-attention mechanism in Transformer but adds structural features in the form of centrality and spatial encodings. GraphGPS [Rampášek al. 2022 proposed general framework for combining paasageepssingg networks with attention mechanisms, while allowing for the mixing and matching of positional and structural embeddings. These emerging approaches better capture global dependencies and structural properties of graphs, marking a significant step forward in the evolution of graph learning methodologies.
2023
B.2 NODE-LEVEL IMBALANCE LEARNING ON GRAPH
Unlike graph classification, the node classification task aims to predict fine-grained labels for individ-
ual nodes (Liu et al. 2020 Ju et al. 2023 Yun et al. 2022 Fu et al. 2023). Class imbalance in node classification arises when the distribution of labeled nodes uneven across different categories. To address this issue, researchers have explored various cutting- methods, which can be broadly categorized into algorithm-level and data-level approaches. Algorithn-level methods focus de- signing representations optimized for minority classes. For example, ECGN (Thaaalyya et al. 2024 enhances graph node classification by integrating cluster-aware node updates with synthetic minority class node generation. RLP-HGNN (Sun et al. (2025) designs a relation-aware label propagation mechanism to obtain pseudo-labels for nodes in heterogeneous graphs DPGNN (Wang al. 2021
adopts class cottotype-driven training strategy to balance losses across classes and utilizes metric learning to determine the positional relationship between nodes and their corresponding prototypes. At the data augmentation level, GraphSMOTE (Zhao et al. 2021b) generates synthetic minority-class node representations through interpolation between minority- node pairs. GraphENS synthesizes complete ego-networks for minority classes (tail nodes and their one-hop neighbors) based on simi- larity. Tail-GNN (Liu et al. 202T) enhances tail node representations by transferring neighborhood patterns from high-degree nodes and locally adapting them to each target node. ImGAGN (Qu et al. 202T mitigates class imbalance in node classification by using a graph generator to produce synthetic minority-clee class nodes, thereby balancing the class distribution.
In contrast, graph classification assigns a single label to each graph, which means there may be no node-level labels. Techniques that rely on node annotations node-level augmentation
18

---

Published as a conference paper at ICLR 2026
cannot be directly transferred to this scenario, demanding specific strategies to handle class imbalance at the graph level.
B.3 EDGE-LEVEL IMBALANCE LEARNING ON GRAPH
Edge-level class imbalance is an important and understudied problem with practical relevance to tasks such as link prediction, edge-level anomaly detection, and inductive recommendation (Liu al. 2025). The work (Wang & Dem 2022) were among the first to expose degree-related biases in link prediction. The influential Topo (Cheng et al. 2025 formally introduced the notion of topological imbalance for edge classification and proposed topological entropy as mechanism for weighting edge importance, which catalyzed subsequent work in this area. Follow-up approaches include EdgeSMOTE (Gong et al. 2025), which synthesizes minority-class edges by generating high-quality node features to rebalance the class distribution. In application domains where edge imbalance is inherent- e.g., few-shot link prediction (Yang et al. 2022) and edge-level anomaly
detection (Liu et al. 2022a researchers have pursued improvements via architecture adaptations and tailored loss functions.
B.4 SUBGRAPII EXTRACTION ON GRAPII
Subgraph extraction a fundamental technique for identifying the most influential local structures within a graph. It plays pivotal role in fields such as protein analysis, biomedicine, and community detection (Yang et al. 2023 Recent work has proposed a variety of subgraph-orientee methods. Density-based discovery approaches have been adapted to network science, biological analysis, and graph databases (Nguyen & Vullikanti 2021 GNN-AK (Zhao et al. 2021a leverages star pattern subgraphs with base graph neural network kernel to perform convolution across all subgraphs. thereby producing diverse and informative node embeddings. SubGNN (Alsentzer et al. 2020 decomposes graph's 's topology into three channels that reflect position, neighborhood, and structure, and it designs corresponding subgraph patterns to capture each aspect.. Despite these advances, the issue of subgraph imbalance remains largely unaddressed. In many graphs, critical structures may occur infrequently but play a decisive role, as exemplified by key molecular motif's in chemical compounds. Accurately extracting and characterizing such minority subgraphs remains an important and underexplored challenge.
C METHOD
C.0.1 IMBALANCED FEATURE MIXUP
Given the representation of N graphs denoted as H. the feature mixup strategy randomly shuffle the order of the representation vectors of the graph instances, and the output denoted as e xdx, Then, sample mixing ratio À [0, 1] from Beta(a, a) distribution, and then perform linear interpolation between H and H using the mixing ratio as follows,
A, = AH; + (1 A)H, E Rd [N.
(13)
Finally, we concatenate the mixed instances with the previous graph instance representations, and the output is denoted as H A H E R2dxd, In this way, the representation of the tail graphs can gain supplementary information from the head graph representation.
D THEORY PROOF
D.1
EXPLANATION FOR THEORY
For the training graphs G. prototype features S, and labels Y, the objective of the information bottleneck expressed as minimizing the following objective function:
min I(S; G) -I(S; Y)
(14)
The Lagrange multiplier 3 determines the level of relevant information captured by the representation S, I(S; Y). For the prototype representation variable S, which is obtained via nonlinear mapping
19

---

Published as a conference paper at ICLR 2026
from the graph representation G, its statistical properties be characterized either by its joint distribution with G and the label Y. or equivalently by the encoder distribution P(S and the decoder distribution P( P(Y|S). Given the joint distribution P(G Y). the pr representation S can be uniquely mapped to a point on the information plane (Shwart--Ziv & Tishby 2017), with coordinates (I(G; I(; Y)), which represent the mutual information between the prototype and the input graph, and between the prototype and the label, respectively. Based on this perspective, we -hterrttiintttttticcccoo analysis to our proposed model, Unilmb, whose architecture satisfies
apply the the Markov chain condition (Tisbby & Zaslavsky 2015). Accordingly, the variable S can be mapped to point on the information plane and forms a unique information path together with the input G, the label Y, and the predicted label Y:
Y
(15)
This path satisfies the following Data Processing Inequality (DPI) chains
Cooer 1999::
IG Y) 2 ((S 2 I(; Y) and H(G) 2 ((; 2 I(; Y)
(16)
These inequalities indicate that the prototype representation compresses the input information while retaining task-relevant signals, thus revealing fundamental trade-off between representation capacity and generalization from an information-thooreee
perspective. Therefore, we can maximize the mutual information between the predicted labels and the ground-truth labels, I( (Y; Y), through supervised learning, which in turn maximizes the mutual information between the prototype representation and the labels, I(S; Y). Based on this objective, the optimization process can be further simplified to minimizing the mutual information between the training graph set G and the prototype feature S, I(S; G), i... reducing the redundancy in the input information caused by maas--iballace..
D.2 PROOF FOR PROPOSITION
According to Theorem 1, our goal is minimize the mutual information between the prototype distribution and the input graph data G, in order to reduce the dependency on the imbalanced distribution of the training data. This ensures that the prototypes not overly tailored to specific graphs, but instead are shared in a more general and consistent manner across the entire dataset.
Formally, the mutual information between prototype variable S and the graph identity G is defined as follows:
I(S G) = Σ (())p(s()) p(slg) p(s)
(17)
Directly computing the marginal distribution of the prototype variable, p(s) = Σi pss 9.) =
EN p(s()))p(G..) is generally intractable. To address this, we introduce auxiliary distribution r(s) to approximate the true marginal. By the non-negativity of the Kullback-Leibe divergence, i.e.,
kL(p(s)][((s)] = Σ p(s) log p(s) r(s)
20
(18)
we obtain the inequality
Σ p(s) logp(s) 2 Σ p(s) rss),
(19)
with equality holding and only if p(s) r(s) for every s E S. This result leads to an upper bound on the mutual information, which allows us replace the intractable marginal term with a more manageable surrogate rss). Therefore, we arrive at an upper bound on the mutual information:
I(S G) s Σ ((C)p(s()) log p()
(20)
r(s)
G.8
20

---

Published as a conference paper at ICLR 2026
Then approximate the true joint distribution p(G. v) p(G) using the empirical data distribution (Alemi et al. 2016 as follows:
2N Σ
p(G v)
dz. (9) 6y. (v)
(21)
2N
where the Dirac delta function 8(-) used to express the empirical distribution of the joint random variables G and y, approximating the true joint distribution p(G, y). Therefore, can obtain this formula:
2N
P(s r(s)
I(; G) $ 2N
Σ E p(s log
(22)
Assuming that our choice of p(s|G.) and the auxiliary distribution r(s) allows for an analytical evaluation Kullback-Leibler divergence, we can directly incorporate this term into the objective. Thus, obtain the following optimizable objective, which aim to minimize:
2N Σ
min ((; G) = min 2N
Σ Eg. [[))(()))))))
(23)
In the paper, the prototypes are derived from the interaction between graph representations and similarity coefficients. Therefore, the formula can be written as:
2N
min (S; G)
min
Eg. KT. (p (x19a) ||r (x))
2N
1=1
min
Esg (p (sx())
2N
(24)
2N
K
Eg, [KL (p (9)) 1r (ue (9)) [8)
min
2N
2N
min
Eg. [KL (p&f (G.) (G)) || (uxf (G.)) [g,)
2N
1=11
where (-) means the mapping function of the graph. In the last line, to align the two terms, we introduce auxiliary variable G. which denotes any graph. And P = {Pi, (P1.P2..... PK} represents the similarity probability distribution of K prototypes. From step three to step four, we set & to Gi. if so, we still satisfy the distribution r(s) rather than the general case, by modulating the auxiliary similarity probability distribution U {u1, u2...., j.
Thus, the original objective can be optimized by minimizing the Kullback Leibler divergence between the similarity probability distribution P and the probability distribution U:
Pk: log Pk uk
Iin G) min KL(P || U) min
(25)
Then we let:
K
Ps = ui + 5k
with Σ
6k = 0
(26)
km
Therefore, have the following formula,
K
1 8k
KL(P
U) =
Σ k=
(uk + 8k) log
(27)
21

---

Published as a conference paper at ICLR 2026
We aim for Pk uks and 5%- + O,
&
0, vk = 1,2,.... ......... K, so the Taylor expansion for
us
log(1 + æ)) around x = O. Apply this to log(1 + 6L)
us
5k 1
8%
0
1
log
+
+
O
(28)
tkkk
2
u.
3
Wk
11kk
Substitute into the KL divergence:
K
1
8%
8k
KL(P
U)
+ 5k)
+
O
2
Uk
Wk
k=1
1
02
1
83
81
+
O
(29)
2
Uk
6
k
Wk
Since the first-order term Σ 1== 6ş = 0. the KL divergence simplifies to a second-order term 6
Ek=1 16 m a third-order correction Ek=1 and forth-order term
Thus, the approximation of the KL divergence is:
K
K
KL(P U) ~
(30)
Ukk
k=1
k=1
Since &g is a relatively small quantity (i.e., Pa. is close to the uniform distribution }) we only retain first two terms as an approximation. Since the first-order term zero, we have:
82
KL(P |
2
W.
k=1
(31)
K
(pk-
Uk
2
Since our goal is to minimize KL(P U). we can observe that the KL divergence between the probability P and the auxiliary distribution U its minimum when these two distributions are equal. For intuitive optimization, we convert it into hit counts. For 2N graphs, each graph hits TopK2 prototypes, and we have total hit count of K=2Nx TopK2. The actual number of hits for each prototype = {n1, 722, ... T.k]. where =C= k E [I. , K]. If the attention coefficient of each prototype and a graph is greater than zero, then we consider this prototype to be hit once. In order word, Each value in N E k equal to the number of values greater than zero in each column
of TopK,
HS /Vd+ +
22X
However, we aim to optimize N to approach KUU through loss function. Nevertheless, the TopK operation is oon-differentiable in practice. To address this, we introduce regulation coefficient E R into the similarity computation process as follows.
ñ Sigmoid (TopK, (iiST) / /Vdn - + +7) Hs.
(32)
D.3 PROTOTYPE-BALANCHE OPTIMIZATION
To meet the above objectives, introduce specifie constraint loss LM that iteratively models the prototype activation process during training.
1
K
2 N TopK,
LM = 2
k=
35
(33)
K
2 * N TopK
Σ
17 + StopGrad (ny - -n)
2 A=
22

---

Published as a conference paper at ICLR 2026
where ng Rk represents the sequence of activation counts for K prototypes. StopGrad()) denotes the stop- hoo-gradient operator (Chen & He 2021). which behaves as identity in the forward pass but yields zero in the backward pass, and this common approach used for non-diffeeemmiaale operations. The gradient of the loss LM with respect to y is:
2 * N TopK2
VPLM= n
(34)
Therefore, be optimized by the following objectives:
2 * N TopK2
" 1 sgn (V,fM) = - sgn 71g
(35)
where yp the learning rate for updating 7). The sign function sgn (-) used to control the attenuation direction. Next, we select several commonly used discrete distributions as the auxiliary distributions U show.
Uniform Distribution. when the auxiliary distribution U follows a uniform distribution Kuipers
Niederreiter 2012. i...
our load balancing optimization goal can be written
K
2 * TopK2
- リ- sgn (V,LM) = sgn
(36)
K
When the load Mai for prototype is higher than the average, it will be penalized to reduce its activation frequency. This mechanism encourages a balanced distribution of attention between head and tail graph information during training.
Zipf Distribution. when the auxiliary distribution follows Zipf distribution Axtell can be written as,
2001
Uk
1/k= Z
where Z
Σ ;=1
1 ja
8>1,
EE K}
(37)
where s is the parameter of zipf distribution and our optimization goal is as follows,
2 * N TopK
- 17 40 sgn
11g
(38)
where yp is the learning rate, and Z = -=>ド FF
3 Exponential Distribution. when the auxiliary distribution U follows exponential distribution Marshall & Olkin 1967,, it can be written as,
Ak
K =
where
Z
Aj
A>0 0
(39)
Z
where the decay parameter governs the steepness of decline. A larger implies a stronger preference toward early-index prototypes, effectively enforcing sparsity and focus in the prototype activation pattern. And our optimization goal is as follows.
2 N TopK,
7 + sgn
(40)
Z
Poisson Distribution. when the auxiliary distribution follows a poisson distribution (Conull & Jain 1973 it can be written as,
23

---

Published as a conference paper at ICLR 2026
e kZ
Ne-A j!
Uk
where ==
(41)
5=1
where acts as both the mean and the mode of the distribution. When À is small, attention is concentrated on early prototypes: when increases, the most attended prototypes shift to middle indices, promoting more balanced activation pattern centered around a specific region in the prototype space, And our optimization goal as follows,
TopK2 A k.
1 +
sgn
Tik
(42)
Z
In Appendix D.4 we provide extensive experimental evidence demonstrating that the model achieves optimal performance when the auxiliary distribution U follows a uniform distribu-
tion, i.e., U
This indicates that when each prototype has balanced load, the
K
extracted prototype features are more generalizable.
In fact, prior study, DeepSeek (Wang et al. 2024c). also follows similar intuition and applies the concept of balancing to the design of large language models, claiming that balanced expert allocation can enhance model performance. While their work differs fundamentally in task setting, our core contribution lies rigorously justifying this intuitive principle from the perspective of information bottleneck theory- particulllll in the context of imbalanced graph learning. Furthermore, we provide more detailed theoretical analysis the balancing optimization strategy, along with formal derivations that establish its effectiveness. This complements and extends previous work by offering principled formulation the balancing objective and demonstrating its benefits through both theoretical perspectives.
Finally, our load balancing optimization goal can be written as:
TopK2
- - sgn (V, = 17 sgn 7i
(43)
K
D.4 COMPARISON ACROSS DIFFERENT DISTRIBUTIONS OF U
In this section, we present evidence for setting U to an uniform distribution. Using four class- imbalanced datasets with extreme imbalance ratio example, compared five settings of U, including Zipf distribution Axtell 2001). Exponential distribution (Marshall & Olkin 1967). Poisson distribution (Consul & Jain 1973) and Uniform distribution (Kuipers Niederreiter 2012 These settings correspond to different optimization losses described in Appendix D.3 As shown in Table and Table g observe that the model achieves the best performance when U follows uniform distribution - that is, when prototypes are expected to be activated uniformly. This setting encourages better feature extraction for tail classes, leading improved generalization.
Table Macro-F1 and Micro-F1 scores on class imbalance datasets with extreme imbalance degree under different distributions. The best results are marked and the runner-ups are underlined We report the average and standard deviation over 20 runs.
PROTEINS
D&D
NCII Micro-Fi
COLLAR
Distribution
Macne-FI
Micro-FI
Macro-Fi
Micro-FI
Macro-FI
Macm-FFi
Micro-FI
Zipl Axell 2011
67.50 3.11
73:29 400
44.05 #229
79.00 7.38
65.69 2.11
79.78 =
6 1026
69.79 1008
Expexemete onrriaal Olkin 1967
66.45 3.34
71.99 1494
44.99
2:05
42.51 9.36
68.03 2.41
79,85 2.09
75.51 1.75
74.47 1.00
Paisson (oro w.. inniin 1974
68.00
73.04 1:00
41.62 5.23
71.96 + 10.6
67.67
80.00 800
25 25.86427000
Unifurat
Kuipets Nuevoreber 2012
70.44
74.50 4.99
16.63
83,00
80.68 4.22
75.73
76.34
E ALGORITHM PSEUDOCODE
We present the pseudocode for Unilmb in Algorithm and comprehensive implementation details available anonymized code repository link.
24

---

Published as a conference paper at ICLR 2026
Table 7: Macro-Fl and Micro-Fl scores on topological imbalance datasets with extreme imbalance degree under different U distributions.
PROTEINS
D&D
NCH
COLLAB
Distribution
Macro-F1 F1
Micro
Macro F1
Micco F1
Macro F1
Micco-F1
Macro-FI
Micro-Fi
7ipl
Astell 2001
39 1.54
74.01 a0
74.32 m0
76.13
63.08 5.20
63.94
-9.51=608
22.47
Expengmeen vie mill Olkin 1967
6
24.57 2090
14.17 = 1.30
63.09
6 15.75
69.30 11.20
Poisson CMiore
69.30 1.31
74.35
74.41 s0m
26
58.65 1914
59.88 8.27
68,09 9.21
70.33 18.54
Unifurm
2012
71.32
74,89
74,69 2
64,99
65.76
73.51
x EAS
Algorithm 1: Unilmb for Imbalanced Graph Classification
Input: Input Graph G (91 G2... .g....,
................................ GNN with graph instances where the graph instance
G3 = (Vw, g A si X and the corresponding label set Y {81, U2....... ................................ . UN }N
Output: Well-trained parameters 1i of Unilmb.
while maximum epochs nor reached or not converged do for patience = 1.2..... , do
P
Computing graph topology encoding LE and GE in Eq and Eq 2 Implementing personalized graph perturbation strategy;
H < ReadOut GNN
X, S,G.ELE).
// Cenerate graph-level
zepresentations
Hs = Softmax
TopK,
SH /Vdh)) HW::
// Prototype feature
perception
H Sigmoid
Topk2
HS
/vd, + +1) + Hs;
// Dynamic balanced
prototype
O θ - vo
// Update the parametners of Uni using cross-ent classification loss
2>N=Topk,>
+ 1 9 sgn
71
K
// Update the balance requlation
coefficient using prototype-balancing Optimization loss end
10
11 end
F DEFINITION OF GRAPH CLASS IMBALANCE AND TOPOLOGICAL IMBALANCE
F.1 CLASS IMBALANCE
Definition. There is a notable imbalance in the number of labeled samples (nodes or graphs) across different classes, resulting in a long-tailed distribution of quantities (Qin et al. 2024
Definition 1 (Graph-Level Class Imbalance Ratio). Graph-level Class Imbalance occurs when there is a disproportionate number of labeled graphs across different classes, resulting in long-tailed distribution. Given set of labeled graphs G {9I, G2..... GNN with N graphs, where each G, represents the graphs in class i, the imbalance ratio is defined as the ratio between the size of the largest class and the smallest class, i.e.,
9max
p
(44)
where 9max denotes the number of graphs in the largest class and Gmin represents the number of graphs in the smallest class.
F.2 TOPOLOGICAL IMBALANCE
Definition. The imbalance facilitated by the uneven graph size (the number of nodes) distribution QQin et al. 2024
25

---

Published as a conference paper at ICLR 2026
Definition 2 (Graph-Level Topological Imbalance Ratio). Following by the work QQnn et al. 2024). Graph-level Topology Imbalance refers to the uneven distribution of graph sizes (1.c., number of nodes) across dataset, which often leads to biased model performance. Let G {91. ....... 9N with graphs denote the set of labeled training graphs. To capture the long-tailed nature of graph size distributions, labeled graph set G into two parts: the top 20% of graphs with the largest number of nodes constitute the head subset Hgg while the remaining 80% form the tail subset 7g. Then we sample the training set generated by the graph based on different settings. The opplogy-imbalance ratio is then defined as the ratio between the average graph size the head set and that in the tail set:
.....
follow the Pareto principle (i.e., the 20/80 rule) (Sanders) 1987) and divide the
p [Hg] Eg.eno [
(45)
17gl Dg,eTc ETO [100]
where VS. denotes the number of nodes in graph Gs Such imbalance is especially prevalent in graph classification tasks such as molecular property prediction or protein function analysis. Larger graphs typically encode richer structural information, making them more expressive and leading models to perform better on them. As result, models may become biased toward large graphs, while underperforming small ones, which contributes performance disparity.
G DATASET
G.1 INTRODUCTION
G.1.1 COMMON GRAPH CLASSIFICATION DATASETS
We conducted extensive experiments on 12 widely-adopted graph datasets, including nine binary-clasifffcttion datasets and three multi-class classification datasets. As shown in Ta- ble 8 these datasets cover a range of domains: Chemical compounds: PTC-MR, NCII, MUTN AUTAG FRANKENSTEIN, and AIDS. Protein compounds: PROTEINS, D&D, and
DHFR. Social and collaboration networks: REDDIT-B, IMDB-MULTI, and COLLAB. 4
Synthetic graphs: Synthie.
Table 8: Statistics of benchmark datasets for imbalanced graph classification.
Datasel
Domain
# Graphs
Arg-Node
Asg-EEdee
Fentures
Classes
MUTAG
Dehnath et 1991
Chemical
188
17.93
19.79
2
Lorvonee
PROTEINS
Yanardag Vishwanathun 2015
Protein
1113
39.06
72.82
218
D&D
Shervastnize alal. 2011
Protein
1178
284.32
715.66
89
2
231
AMM
Chemical
4110
29.87
33.30
DHF
Sutheriand
2003
Protein
756
42.43
44.54
AIDS
Kuesen & Bunse 2008
Cbemical
2000
15.69
16.20
38
2
FRANKENSTEIN Chmm
Chemical
4337
17.88
2
757
REDDIT-B
Yonardag Vnlowasation
2015
Social
2000
429.63
197.75
2
400
CO1 LAR
L.ckovee et 2005
Collabomtion
5000
74.49
2457.22
3
991
Synthe Moms ct 2010
Synthetic
400
95.00
172.93
IMDB-MULIT Yamardag
Social
65.94
G.1.2 AIRGRAPH DATASET
To evaluate the practicality of the model, release air quality graph dataset, AirGraph, for research on air pollution level prediction. This dataset contains 30,660 graphs, significantly surpassing the size of most commonly used benchmark datasets in the graph learning domain. Below, we provide a detailed introduction to the dataset.
Data Source. The dataset collected by the China National Environmental Monitoring Center, covering 1,341 monitoring sensors across provinces in mainland China from 2021 to 2023, with hourly PM2.5 readings recorded.
26

---

Published as a conference paper at ICLR 2026
Graph Generating. We generated the graph structure and node features of the dataset based the following rules to align with the graph classification task.
Node construction: All monitoring stations within each province are treated as nodes the graph;
Edge construction: Edges established hased the geographic distances between monitoring stations, reflecting spatial dependencies;
Node features: Each node's features consist of the past 24 hours of PM2.5 observations, capturing temporal information;
Graph labels: Each province-level graph label represents the air quality level corresponding to the average PM2.5 index of all sensors in the province over the next 24 hours, categorized into three levels: low pollution (< 25 μg/m²), medium pollution (25-75 ug/m²), and high pollution 75 pg/m²), following international air quality standards.
Based on this processing procedure, we constructed large-scale dataset comprising 30,660 graphs (28 provinces 365 days years), which significantly exceeds the scale of most commonly used benchmark datasets in the graph learning field.
Graph Imbalance Characteristics. As shown Figure 7 AirGraph exhibits a naturally im- balanced distribution of pollution levels: high pollution accounts for 6.86%, medium pollution for 42.84%, and low pollution for 50.30%, with the ratio between the largest and smallest classes reaching 7.33. Due the heterogeneity in the distribution of monitoring stations across provinces, the graph topologies in AirGraph also exhibit naturally imbalanced distributions, providing a realistic and ideal resource for studying imbalanced graph learning problems.
Relative percentage per day (%)
80
Low Pollution Level
High Pollution Level
Medium Pollution Level
High Pollution Level
.866
40
Pollution Level
50.30%
42.84%
Medium Pollution Level
20
2021
2022
2023
2024
Figure 7: Distribution of Pollution Levels in the AirGraph Dataset: High (6.86%), Medium (42.84%), and Low (50.30%).
G.1.3 PROCESSING OF CLASS IMBALANCE GRAPH DATASETS
Dataset Setting. We perform graph classification tasks on twelve modified Class Imbalanced graph datasets. Among them, nine datasets are used for binary classification tasks, and three datasets are used for multi-class classification tasks. Therefore, we adopt two different processing methods for the imbalance settings follows. Final detailed iaassiibalaance dataset summary can be found in Table 2 For the binary classification task, we follow the setup the GNNN paper Wang et al. 2022b), where the dataset is split into 25% training and 25% validation, with the remaining samples assigned to the test set, resulting in specified numbers of graphs in each subset. We treat the class with more graphs as the head class and the other the tail class. The imbalance ratio p is defined the ratio between the number of head and tail graphs. According to three different ratios highly imbalanced setting =10 (extreme), moderately imbalanced setting p x2 (medium). and balanced setting p =2 (low) - we sample the desired number of graphs for the training and validation sets. All remaining graph samples allocated to the test set.
27

---

Published as a conference paper at ICLR 2026
@ For the multi-class classification task, allocate 25% of the dataset for training, 10% for validation, and the remaining 65% for testing, allowing us to obtain the desired sizes for each subset. Then, sort the classes in specific order and sample training graphs for each class following an exponential decay: N; = No e -xx where No means the ratio of the largest class, which we set to 1. N; represents the number of sorted i-th class graphs. We control coefficient À meet three different imbalance degree: highly imbalanced setting with ~ 100 (extreme), moderately imbalanced setting with p =22 (medium), and balanced setting with p Re (low). In the validation set, we maintain an equal number of samples across all classes. All remaining graph samples are assigned to the test set.
Table 9: Statistics of the Class imbalance datasets for graph classification. In the graph-level class imbalance setting, the imbalance ratio p is defined the ratio between the number of graphs in the Head Class and that in the Tail Class. For binary classification tasks, the validation set shares the same class distribution as the training set. For multi-class classification tasks, the number of graphs per class in the validation set is equal to ensure fair evaluation.
Dataset
Task
Level
# Graphs (Hrad Class)
Graphs (Tail Cluss)
Imhalance Ratio
extreme
45
5
9:00
MUATG Debnath et al.
[99]
Binary
medium
35
2.33
extreise
81
9
9.00
PTC-MR Toivonen et al.
2003
Binary
mellium
63
27
2.33
low
45
45
extremee
PROTEINS Yonardag Vishwanthhan
2015
Hinary
medium
210
90
2.33
low
150
150
L.00
270
30
9:00
D&D
Shervashseee al.
Binsry
mediium
210
2.33
low
150
130
1.00
extreme
900
100
9:00
NCIL
Wale et al. 2008
Binary
mediium
700
300
2.33
500
500
extreine
108
12
9000
DHFR Sotherland et all
2003
Binary
medium
84
36
2.33
Low
60
60
1,00
extremee
Bunke
Binory
mociium
350
150
2.33
low
250
250
1,00
900
100
9:00
FRANKENSTEIN
K Jnsini er
2015
Rinary
low
500
500
L.00
extreene
450
50
9:00
REDIIT-B Yammlag & Vishwanttion
2015
Binary
modiiom
350
150
2.33
250
1100
11
100.00
COLLAB Leskooe 2005
Multi-Class
moddium
980
low
417
417
1.00
extrense
100
100.00
Synhhe Marris al. 2016
Multi-Class
msedium
60
20.00
25
extrecnc
300
3
100.00
IMDE-MULLTI Yanardag Vishwarathan
Multi-Class
mediium
300
15
20.00
125
125
1.00
G.1.4 PROCESSING OF TOPOLOGICAL IMBALANCE GRAPH DATASETS
Dataset Settings. We perform graph classification tasks on eight modified Topological Imbalanced graph datasets, including six for binary classification and two for multi-class classification. Following the experimental settings of IGL-Bench (Qin et al. 2024). we adopt a unified processing approach for both binary and multi-class graph classification tasks. Since the test sets of some datasets may not include all classes when generating topological imbalance data under the predefined settings, four datasets are excluded from the evaluation on topological imbalance performance. The detailed topological imbalance dataset settings are summarized in Table LO
Specifically, following the setting in the IGL-Bench Qin et al. 2024), we partition the dataset into training and validation sets with a 10%/10% split, while the remaining 80% of samples are assigned to the test set. To ensure fair evaluation under imbalanced graph distributions, we adopt class-balancee splitting strategy, where each class is represented by equal number of graphs in both the training and validation sets. As formally defined in Definition 2, we define the head graph subset as the top 20 % of graphs with the largest number of nodes, while the remaining 80%
28

---

Published as a conference paper at ICLR 2026
constitute the tail graph subset (Sanders sample from the two subsets to construct the corresponding training set. To comprehensively evaluate the robustness of Unilmb under varying degrees of topology-imbalancee we create three representative splits with different levels of balance: extreme, medium, and low. Due to the complexity of the graph size distribution, it is challenging to strictly satisfy the target sparsity ratio in practice. After multiple attempts, we obtain the final splits whose imbalance levels closely approximate the intended scenarios. The resulting appoloyyiiballnge statistics are summarized in Table TC reflecting the practical difficulty in achieving exact sparsity control under real-world graph distributions.
1987.
According to the desired imbalance ratio p.
Table 10: Statistics of the modified Topological Imbalance datasets for graph classification. In the graph-level topological imbalance setting, the imbalance ratio p is defined as the ratio between the average size of the head graphs and the average size of the tail graphs. To ensure class balance, the number of graphs per class is equal in both the training and validation sets.
Dataset
Tusk
Level
+ Head Graphs
(per Class)
Tall Graphs
(per Class)
Avg, + Size (Bead Graphs)
Avg. # Size (Tail Graphs)
Imbalance Ratio
56,00
11.37
4:93
PTC.MR Talsusen
Binuy
meslium
34.25
11,80
2.90
Jow
29.50
18.4
1.60
extreme
276,00
22.76
12:1
PROTEINS Yanindag&
Rinrry
medium
90.25
22.36
4.05
kow
93.50
45.81
2044
essreme medium kow
1765.00 00
198.11
891
D&D Shervshisee
Binayy
53
524.83
202.7
2.3
$48.00
350.58
1.56
Calicunc melium
83.62
NCN Walc al Ma:
Rirery
21
185
49.90
low
50.45
36,06
77:55
FRANKENSTELN Onini et al. 2015
Binzry
22
195
33.09
2442
222.73
10097
REDDIT-R Y Yaunday& & Velosnathan 2015
Rinary
medium
10
90
1097.75
225.04
4.92
1180.75
425.25
2.79
509.41
50.77
9:09
COLLAB LeAkINe ml XX1
Mulii-Cles
masliun
150
147.32
50.04
kow
50.14
1.76
56.33
10,10
IMDI-MULTI
Yansrding de Vishwanation 2015
Multi- Calii-Ches
medium
48
25.33
2:55
26.13
G.1.5 PROCESSING OF DATASETS WITH INTERTWINED CLASS AND TOPOLOGICAL IMBALANCES
Dataset Settings. To simulate the complex scenario of Class Imbalance and Topological Imbalance interwoven in real-world graph classification tasks, we conducted experiments on six modified graph datasets that exhibit both class-imbalance and topological-imbalance These datasets include four binary classification datasets and two multi-class classification dataset. We divided both the training and validation sets with 10%/10% ratio, and the remaining samples were allocated to the test set. For these two types of datasets, we adopted two different processing methods. The detailed dataset configurations are provided in Table
For the hinary classification task, we adopt a Class Imbalance ratio similar to that used in GPGNN. Specifically, the imbalance ratio in the training and validation sets is set to 1:9, with the remaining samples allocated to the test set. Within each class, the topolog--imbalnn ratio of the average graph size (in terms of node count) between the head and the tail set, as outlined in Definition 2. To achieve this, we employ the Pareto principle (20%/80%) to partition the dataset into head and tail sets. Consequently, for each class in the training set, we allocate 10% of the graphs to the head set and 90% to the tail set based on their graph size, thus creating a scenario where both class-imballnce and topology-imbalnne are interwoven.
ratio is defined as the
@ For multi-class classification task, similar to the aforementioned class-imbalance setting, we assume that the number of graph samples per class in the training set follows an exponential decay distribution. Given the total number of training samples and class-imbalance ratio p. we determine the number of samples per class such that the ratio between the Head Class and the combined Tail Classes is 9:1, thereby creating extremely Class Imbalanced training environment. Building this, we further introduce Topological Imbalance within each class. Specifically, we rank the graphs in each class by their size (i.e., number of nodes), and partition them into head and tail subsets
29

---

Published as a conference paper at ICLR 2026
with 1:9 ratio. That is, only 10% of the graphs in each class belong to the structurally complex head subset, while the remaining 90% belong to the simpler tail subset. This strategy enables us to construct challenging experimental scenario where class- muss-imbalaane and eapllagical-imbllaaee are interwoven. Investigating such intertwined imbalance settings is significant importance, as closely mirrors the complex distributions commonly encountered in real-world graph classification tasks. In many practical applications, such as social network analysis, protein-proteie interaction prediction, and recommender systems, class-imbalance often coexists with disparities in structural complexity- some classes consist of graphs with more nodes and complex structures, while others are relatively simple. By conducting experiments under this intertwined imbalance setting, we can comprehensively evaluate the robustness and generalization ability of Unilmb in realistic and challenging scenarios.
Table 11: Statistics the modified Intertwined Class and Topological Imbalanced datasets for graph classification. In the graph-level flss-imbalance setting, the imbalance ratio p is defined the ratio between the number of graphs in the Head Class and that in the Tail Class. And in the graph-level appoloyy-imbalane setting, the imbalance ratio p is defined the ratio between the average size of the head graphs and the average size of the tail graphs.
Graphs
Size
Inbullane topelags)
Hoad Grapls
Tail Grapts
a Head Guzaks
Tal Graas
(llad
PIOTEINS Sazandig& Vietomatoon 2013
Bisary
21.53
2.7.7
SCTI
Wide cl 200
Bissay
57
333
$2.27
26:45
9022 8: 1.95
Valworaha 2013
Binay
IMDS
10
H EXPERIMENT SETTING
H.1
HYPERPARAMENEE SETTING
We summarize in Table the detailed hyperparameter configurations of Unilmb across the six datasets discussed in the main text. To ensure fair comparison, the remaining models are based on publicly available optimal hyperparameters.
Table 12: Detailed hyperparameters of Unilmb with Class Imbalance datasets.
Dataset
Learning rate
Hidden dim
GNN layers
Batch size
K
Topk,
Top
PROTEINS
Yanardac & Vishwanathan 2015
0.901
128
32
16
8
D&DD
Shervashiize et 2011
0.001
128
64
16
32
8
NCTI
Walc 2008
0.001
128
64
24
32
12
REDDIL 1
Yanarday & Vishwanatham 2015
0.001
128
64
32
32
16
COLLAB
Leskovec al. 20055
0.001
128
64
32
32
16
IMDR-MI MULT
Yanardag & Vishwanathan 2015
0.001
128
32
32
16
8
Table 13: Detailed hyperparameters of Unilmb with Topological Imbalance datasets.
Dataset
Learningratt
Hiddon
GNN Jayera
Batch_size
TopK,
Top
PROTEINS
Yanardag Vishwanathan 2015
0.001
128
32
24
8
D&D
Shervashidze
2011
0.001
128
32
4
NCTI
Wale
0.001
128
32
4
REDDIT-B
Yanardag & Vishwanathan
2015
0.001
128
32
16
COLLAB
Leskovec
2005
0.001
128
32
8
IMDB MULIT Yanardag & Vishwanathan 2015
0.001
128
32
16
8
H.2 EVALUATION METRICS
Following existing work in Class Imbalanced graph classification, adopt two evaluation metrics, Macro-Fl and Micro-F1 (Wang et al. 2022b), to assess the performance of Unilmb and other baseline models. Macro-F1 calculates the accuracy for each class independently and then averages these values,
30

---

Published as a conference paper at ICLR 2026
treating all classes equally. In contrast, Micro-Fi computes accuracy over all samples collectively, which may result in the head class dominating the evaluation. A more detailed explanation of these evaluation metrics can be found follows.
Macro-FI FiXX ct al. 2014 is a commonly used metric for evaluating multi-class classification performance. It computes the F1 score for each class independently and then averages them, treating all classes equally regardless of their size. The formula is given by:
N 2 Precision; Recall,
Macro-F1
(46)
Precision; + Recall;
where Precision, = FP, and Recall; TP.4PN," Macro-F1 equally considers all classes, regardless of their frequency, and emphasizes the balance between precision and recall, is especially suitable for imbalanced datasets. Since each class contributes equally, Macro-F1 can be overly influenced by rare classes and may not reflect the overall performance class distribution is extremely skewed.
TP.
TP.
Micro-F1 (Lipton al. 2014) aggregates the contributions of all classes to compute global precision and recall, and then calculates the Fl score. Unlike Macro-F1, it does not treat classes separately but considers the total number of true positives, false positives, and false negatives across all classes. The formula is:
Precision Recall
Micro-F1 Precision + Recall
(47)
where Precision
E.P E. FFP:+ and Recall Σ, (T::+NN Micro-F1 provides a global measure
ETT.
of performance and is robust to class imbalance, it gives more weight larger classes. Therefore, it may overlook poor performance on minority classes, as these contribute less to the overall count of predictions.
1 ADDITIONAL EXPERIMENTAL RESULTS
1.1 CLASS IMBALANCE PERFORMANCE ANALYSIS
I.1.1 COMPLETE EXPERIMENTS ON ALL DATASETS WITH VARING IMBALANCE DEGREE
The performance on the 12 class-imbalanced datasets is presented in Table 14 Figure T6 and Table Classic graph representation models perform the worst due to their inability to address the challenge of class-imbalance. Oversampling methods involve resampling tail classes, which may lead to overfitting of tail features and consequently degrade generalization performance. DataDec identifies an informative subset for model training through dynamic sparse graph contrastive learning and leverages abundant unlabeled data to achieve limited improvements. It is worth noting that the original code of DataDec provided by its authors uses the test set for classifier fine-tuning. For a fair comparison, we only use the training set for fine-tuning, which may result in performance slightly lower than what was originally reported. Topolmb and ImGKB, as well as G2 GNN and ImbGNN, improve representation learning under class-imbalance by selecting similar graphs for data augmentation within tail classes, thereby achieving higher prediction accuracy. Our model achieves the best performance across all metrics on the dass-imbalanced datasets. This superior performance is attributed to our prototype-basee strategy, which effectively enhances the model's 's ability to learn robust and discriminative representations from dlas-imbblaaneed graph data.
16
We further compare our model with bascline models under scenarios of medium and low class- imbalance (i.e., close to balanced). As shown in Table 15 and Table 16 we find that as the class distribution approaches balance, classic graph classification models perform increasingly better, while models specifically designed for imbalanced settings tend to become less effective. For instance, GNNNN constructs the graphs of graphs that can easily lead to information redundancy, thereby degrading performance. In contrast, DataDec performs increasingly well and becomes the best- performing baseline in the low-imbalance setting. On the other hand, our model achieves consistently strong performance across all scenarios, demonstrating its robustness in handling more complex and varying class distributions.
31

---

Published as a conference paper at ICLR 2026
Table 14: Macro-FI and Micro-Fl score on class imbalance datasets with extreme imbalance degree. We report the average and standard
The best results are marked and the runner-ups are underlined deviation over 20 runs.
PROTEINS
D&6
SCH
PTC-HR
DiIFi
Madel
Bockbore
18:0
12:44.640
21.3
21A1
37
31.00
56.28
2100+
50.50 18:31
56.78
18:11 65
65
25:6:130
51.0 2330
Expheree
45:39:438
16.33
Capb Maana
11:33-500
1411+
54
42:48-111
68.00
71 11:3--8-
53
11.0810033 1938
3124
265641 546
2010 2415
36/66 1486
53.786600 35.161347
50
$1.00 +:.0
1070-191
58 16
04.75 -167
(4:57 11
malk sok
8301 270
133397714
01.31_199
15:8
4356 5.8
1076
14.75120
46.51 28.37
26.701 14.11
36/2...00
13.0
$1.25.1.
16.00
30.1 4.25
36.2
40:10
quanty
-1
6:
34.1
42:38
BeNONS
3748
2333
1339
01301100
54:01437
Ualted
16.70
10.0011331
REDOLT-3
ABBS
FRASKENTTINN
COLLAB
DIDB-HUIT
Spank
Babbeee
Naar-Fi
MacrFFi
Mavs81
Shan-Fi
Mamne
Stenei
MuaoFi
Mhan-F
Mact-FI
05 lo 104
07.80
(3:5 23
13 23
35.06 w3.0
14
DTC
1x272 13.46
17.02 3:47:
243
25522 1.06
413121111
16.9 24.5
26.99 11247
3159 25.55
GIN
42.09 +1:5
43.374981
03:0 43.34
6331
9001 145
575
10:28 4230
35.33 64:2
40.594229
62.17
6574 9
11 05
18:00-9
14:10 10
(6:30-101
GYEN
N167-10
11:0 +1.9
1140-111
72A
9316141
14.9
2455
2633
1.38
41:3
21:11140
Tackn
94
14:0+111
2000+ 201
05.71 2
1634
1854.190
3615
-7:5--8
43.33
Table 15: Macro-Fl and Micro-Fl score on class imbalance datasets with medium imbalance degree.
MUTFM
PRITFEINS
D&B
NCH
PICAKR
DHIPR
Madel
Bacibuass
Mi
Mire:
Marra
Mires
Mario 51
Viee
2016--11 5.15
80.0 23.4
06.15 2.30
$1.55 2446
MUAI 234
62.11:12 3.2
06,03 319
45.00 21.3
80:22
$2.0
GraaGis
12.84.6.1
$7.29 -:33
53
36 1324
4197 15.29
5600 353
55.4
CTC
Ough Mari
1800 11277
2439
65.21 130
22:9
55 2475
M4Grogh
3110 3m
4112-4211
54.8
75.03 4.0
07.37
55
51.29 1622
35.25 24
-3.51-421
43.37 112
Gupach
GYONN
51.36 2:6
01.02 2250
$5.29 1.76
07.51
55 1187
BUS LM5
40.12
37.40
35.25-112 912
4524
4534415
80-301125 13.84
96.58 1150
3016 13.41
31.9 3-42
3616 1635
Tocacn
68 -.4
67.51 11
67 111
5530
210
53.0
11
2080
14
Unileds
80.18
50.9
PRANKENSHLIS
DNIB-ULITY
Systte
Mudil
Tucthuss
WimoFF
Mkve-F1
WinoFF Pi
MaroFi
Mi
ManoFi
Mi
Marn-h
Mi
Marn
GNN
(3.55 1012
90.55 +00
13
1130-11
vodia
60:29:2.44
66.8 1131
1674400
2027
0900 A021
52.2
430
34.7
436
1535-327
Iapherac
2833 1436
1330
31301A5
1130
43 23
30.3
5885
0001
23
5355
4024
(C.1045.9
GBI
7183 4::0
41.30
07.32 4033
1910 A030
61.2 61:1
6148 4344
60 60
16.5 18.9
34.7. 3
Dig GighCt
12.05 6321
76.38 6337
100
63.9 1803
(111 RUR
44.204430
19
45
CIGNN
52.30 13
445547
101-18
15.00+011
-011
$4:00-1111
3051+180
500
31811191
quidiy
Ivckn
10
06:30
11
14
315-19
6737
29.75
1437 2639
39,81 15.0
3010
55.00 13:0
1294421
Uniieb
16.00 12.6
10.18
09300
2015
21311131
1.2 MULTI-CLASS DATASET
To further validate the performance of Unilmb in multi-class settings, incorporate additional datasets that cover diverse domains. Specifically, we include Enzymes (Schomburg et al. 2004 (a bioinformatics dataset with 6 classes) and Letter-high (Riesen & Bunke 2008) and Letter-low
32

---

Published as a conference paper at ICLR 2026
Table 16: Macro-Fl and Micro-FI score on class imbalance datasets with low imbalance degree.
MUTAG
PROTEISS
D&D
NCTI
PICMR
DRFE
Mipue1
MactoFi
Micsot
Naawei
Mizue1i
Mact-Pi
Micuet
Ma
Maswelt
HamoFii
Min-Fi
14
$1.26
69 76
70.99 2223
99 1010
72:33 486
13.99.3 2480
73,99 23344
55.30 55.29
56.2 285
45,20 2479
23.95
2.13
2225
7339
73.49
22.99
301661936
01.25 15:44
1338 1935
67.34 175
72.83 4239
05324560
00:0 7:46
70024 185
45.00
5035 855
CTC
Gooph Mario
Paphereer
23200.550 505
12.99.129
04.06
126.38
454
45.72
10A39
11:164-544
80.41
67.41 12.14
12:18
04.77 +1:00
66.45 12:4
3001 15:39
47.02.137
46:78-1222
1451
Destac
Sysari spandty
33.58x0 561
11.25 5:34
03. 63
2:33
+10
08 + 1:55
13:
73.33+6:1
m6:3
73.31
+635
13:30333
54.85 3:
06-38+333
1-3
7108 0180
10061317
1183
61.7
InAKNN
73.93+181
08
61.31
01007 199)
1155
REDDIESS
ABIS
HRANKESSTEIN
COLLAB
INDREMBLLTI
Systlic
Miro F1
Mim=F1
Ham-Fi
Min
Mim-
Ma
Mim
MinoFi
82.44 3,07
$3.0 4425
-330
00.1 324
05
-0.3
26.14.500
13.35-441
Ernghtt
43.13 1570
3482 174
SC72-17
4515
63N0 13.35
15.00 2317
30.55 2.5
04.5-113
0h
GrghiKaaka
07.34 11.24
03.27 11:44
4023 -13
19.30
38.55 1667
$2.5 11.21
35.35 1435
34022 3.99
56.57 13M
47.00 1004
10.33 2.5
55,00 2.54
58
35.15
31.93
37:90=133 1:36
(7.51
14-22 64607
90
33.4
25004-44 4/4
4000
13.14
26.55
S1.36-241
15:33 1457
47:25:43
Tegeleh
13.51 +0.11
09.99 -0:33
50 1.
14:00 16:97
13.01 16.35
15.25 S1
36.66
15 16:00-781
+3 02
1:94
10:00 6111
tyra saudyy
95.26 115
68.75
14.00
8147
07.07
6:30
45.22 432
buckn
34.55
33.56
11.0
1.36
10
TeIKONN
40:72:440
55.15 19-25
11 116
45.55 6246
26 11.00
05.9
112
64 396
4471-117
Uiilibh
01:34-655
16816-000
(6.11 -115
08.33-1111
6567
06:97-199
37 435
Riesen & Bunke 2008 (two visual-domain datasets, each containing 15 classes). These datasets
introduce greater label granularity and structural diversity, enabling a more comprehensive evaluation of Unilmb's generalization ability in complex multi-class graph imbalance learning scenarios experimental results indicate that ImbGNN exhibits inferior performance,
As shown in Table 7 achieving performance close to GIN, whereas model Unilmb maintains competitive performance on multi-label datasets.
Table 17: Macro-FI and Micro-Fl score multi-class datasets with extreme imbalance degree.
ENZYMES
Letter-high
Letter-low
Model
Macro-F1
Micro-F1
Macro-F1
Micro-F1
Macro-F1
Micro-F1
GIN
10.15 a 4.36
12.28 2.00
25.62 2.7
27.84 a 1.70
23.08 a: 3.83
40.85 a 2313
GPGNN
17.94 + 1.38
19.33
29.40 +5.39
33.67 +4.92
31.28 2.47
47.92
4..94
Topolmb
4.76
8.34
18.27 +9.26
19.41
21.94
6.28
39,22
7,36
ImbGNNN
14.76 2.06
17.19 3330
31.96 +3.47
35.84
+
27.29 6.54
45.27
5.43
Unilmb
41.38 2:37
56.86 3,14
39.62 1:38
46.94 2.16
43.86 2,48
64.97 3,59
1.3 TOPOLOGICAL IMBALANCE PERFORMANCE ANALYSIS
I.3.1 COMPLETE EXPERIMENTS ON ALL DATASETS VARIOUS IMBALANCE DEGREES
we report the performance gains of our model over other baselines on both large-sacle and small-scale graphs under extreme topological imbalance. And the performance across the 8 datasets is shown in Table 19 Table 20 and Table A SOLT-GNN identifies co-occurring topologies in larger graphs (or "head" graphs) and transfers this information to smaller graphs, thereby improving their performance. It achieves strong results, demonstrating the effectiveness of knowledge transfer mechanisms in alleviating imbalanced classification. ImbGNN designs a topology-aware G²G construction method that is friendly to imbalanced structures, showcasing its potential in handling Although Topolmb was primarily proposed to address uneven substructure distributions, our results also show that it can effectively mitigate datasets. Our model achieves competitive performance regardless of whether the graphs are large-scale or small-scale, which can be attributed to two main factors: first, our topological local encoding and graph perturbation strategies enhance the model's robustness to data variations; second, the proposed prototype-base based strategy improves the representation learning of small-scale graphs.
In Table B8
-mbalannee across multiple
33

---

Published as a conference paper at ICLR 2026
Table 18: Comparison of Micro-Fl performance topological imbalanced datasets with extreme imbalance degree across graphs of different scales.
PROTEINS
NCI1
IMDB-MULE
Model
Small-scale
Large-scale
Small-scale
Large-scale
Small-scale
Large-scale
GIN
56.12 3.188
64.28 3.52
58.39 = 3.12
70.88 = 6.01
32.19 4.41
33.98 = 5.06
InfoGraph
58.42 2.67
65.10 3.19
60.76 = 3.29
70.64 = 5.17
34.22 4.61
37.29 = 2.13
GraphCT.
59.24 + 2.28
65.88 + 3.41
61.33 # 4.67
74.75 + 5.02
35.98 * 5.19
GraphGPS
66.02 2.85
80.40 6.46
60.81 # 4.77
68.18 + 7.61
34.14 1.10
31.38 3.01
Exphormer
65.42 2.84
73.91 + 9.07
61.96 12.87
68.83 + 2.48
33.59 + 3.45
33.63 # 3.79
Graph-Mamba
66.12 # 3.61
85.25 + 3.52
62.49 = 3.60
68.60 6.11
30.45 + 3.19
34.07 = 1.08
Topolmb
53.36 6.68
62.87 36.75
61.72 = 1.24
72,92 = 1.57
38 18.40.0 2.56
47.65 2.74
ImbGNN
68.30 0.88
63.23 + 3.07
58.68 = 1.86
70.89 = 1.56
32.76 # 6.13
34.07 = 7.33
SOLT-GNN
67.34 # 3.55
83.88 + 3.64
58.52 = 1.91
68.77 = 7.96
36.64 # 1.39
42.24 # 2.76
Unilmb
69.68 2.15
87.23 2.04
64.68 1.44
76.19 2.36
42,48 1.81
48.23 2.72
Table 19: Macro-F1, Micro-Fl on lopological imbalance datasets with extreme imbalance degree.
PROTEINS
NCII
PTC-MR Accuracy
Model
Macro Fl
Micro-F1 Fl
Macro.Fi
Micmo.Fi
Macro-Fi F1
Micro. Fl
Macro-Fil
GIN
53.48 2.03
58.00.4.10
57.98 a 5.51
60.68 6.89
61.60 2.20
61.84 2.29
34,56 6.32
48.41 7.077
InfoGraph
55.31 3.30
60.68 3.52
59.46 4.03
61.22 +4.11
62.25 + 1.53
62.69 + 2.01
44,68 - 8.14
48.33
GraphCL
57.62
62.01 5.18
63.29
4.682
GraphGPS
65.54 4.222
69.26 2:48
63.33 a 12.97
63.85 2.08
62.96 3.51
64,02 2:40
9.77
50,46 1.45
Exphormer
64.01
67.33 2.78
8.06
60.26
16.16
63,23
49.14
51.15
Graph-Mamiba Mambba
3.911
72.09 + 3,16
43.86 = 9.73
$4.45 995.3
63.09 + 282
63.63 2.14
49.72 12.58
51.01 5.66
Topolibb
44.79 14.19
54,89 13.58
63.97 a 2.78
2.96
63.57 1.69
64,17
51.65 1
51,96 1.16
ImhGNN
66.87 22559
67.07 + 2.75
68.67 + 6.33
70.71 6,99
61.24 1.72
61.66 1.17
43.87
4.75
52.66 4:53
SOLT-GNN
20.70:22.20
72.14 2.18
64.97
1.44
62.13
40.70 3.27
51.74
UniTmb
71.32 1.88
74.89 1,12
74.49 1,13
76.73 e 1.04
64,99 9.58
65.76 3
$3.86 s 2.59
56.38 1.888
Table 20: Macro-F1, Micro-Fl on topological imbalance datasets with medium imbalance degree.
PROTEINS
D&D
NCI
PTC-MR
Model
Macro-F1
Micro-F1
Macro-F'l
Micro-F1 -11
Macro-F1
Micro-F1
Macro-F1
Micro-FI
GIN
2.90
62.14
60.89 1 1.84
61.46:2..4
$8.37
60.47 3.60
38.58 5:97
InfoGraph
58.87
63.07
61.78 1.99
62.37
63.21 1:33
64.29 3.94
45.28 19133
4
GraphCL
61.4822.88
64.43
63.41 a 4.03
6
64718
66.02 2.13
43.16
4.55
GraphGPS
59.74
58.72
1.05
49.12
47.59
Exphormer
60.32 4.666
63.19 3.61
60.74 + 2.13
61.04 2.70
58,93 1.28
61.49 + 3.02
39.25 3.42
49.16
Graph-Mamba
62.41 3.28
62.85
1,70
57.44
3,16
42.56
50.78
Topolmb
59.03 4..35
64.03 1-1.43
64.65 1.76
65.99 1,25
66.94 1.79
67.16 1.71
49.17 0.95
51.59 1.300
ImbGNN
69.91 =: 1.61
70.70 + 1.99
53.71 # 17.20
58.42 = 17.49
$8,26 4.48
59.83 2.47
49.64 6:20
52.92 4.40
SOLT-GNN
70.58 2.03
71.95 1:36
58.67 1.91
63.33
1.86
62,76 1 1.29
62.81 1.28
43.20 0.36
53.04 1 3.91
Unilmb
70.62 :: 1.60
74.20 0.94
72.70 3.03
74.56 3.18
69.53 1.33..33 3.50
70.07 a 2.84
52.36 x. 2.62
$3.99 3.07
Table 21: Macro-F1, Micro-F1 topological imbalance datasets with low imbalance degree.
PROTEINS
D&D
NCH
PTC-MR
Model
Macro Fl
Micro F1
Macro-F F1
Micro- F1
Macro- F1
Micro F1
Macro Fl
Micro F1
GIN
59.67 1.87
63.58 : 1.76
63.39 : 084
65.01 = 0.69
61.89 : os
63,30 3.19
48.26 3.11
52. 4.36
InfoGraph
2:22
64.27
1.58
62,56 1 1.09
63,83 a 2.29
3,44 1.12
64.69 1-100
47.19
51,07 181
GrahC11.
63.72 = 3.03
4.29 65.08
64.54 a 3.08
64.28 x 4.66
65.08 3.27
66.82 = 43.19
48.57 2.23
$2.90
GraphGPS
65.95 2.46
65,96 3.23
64.69 1.33
65.94
63.20 0.94
63.46 U.8
49.81 2.95
50.07 288
Exphormer
51.77 6.39
54.73 5.82
38.95 + 1.13
41.48 + 1.25
49.17 8.89
54,77 + 3.48
44.99 33.29
45,97 + 2.60
Graph-Mamha
60.16 2.54
60.88 2:20
3.35 59.48
61.42
64.49 1.92
64,70
47.85
4 2.1
Topolmb
59.94.24.33
61.19 4.61
64.97 4 1665
65 4
62.41 1.66
62.46 0.66
48.83
49.71 1.98
ImbGNN
62.41 11.37
64.33 4.665
50.74 x 7.30
58.04 2.77
62.01 0.94
62.09 0.99
48.72 1.222
50.33 1.226
SOLT-GNN
64.60 4.66
63.56 4 4.33
62.34 a 0.66
63.78 a 1.06
58.75 4 0.73
39,59 a 0.77
43.63 à: 1.61
47.54 4.33
Unilmb
65.80 3:51
63. 3.39
65.25 3.49
53.50
4.12
50.85
53.32 +
34

---

Published as a conference paper at ICLR 2026
1.4 ABLATION EXPERIMENT WITH MICRO-FI -1 SCORE.
In the main text, we report the Macro-Fl score from the ablation studies. Here, we supplement the re- sults with the Micro-Fl score for completeness. shown in Figure the findings remain consistent: each component contributes positively to the overall performance, confirming the effectiveness of our design choices.
ImMix
TopoEnn
DalOpt
NL. 245
PROTEINS
DD
NCNI
COLLAR
Figure 8: Ablation experiments of class-imbalance (upper half) and topologicallimbalanee (lower half), evaluated by Micro-F1 score.
Ours
Mask mode
Drop vdlee
DoOA
445
Marre1 (N)
PROTLINS
DD
NCII
COLLAB
Figure 9: Ablation experiments on class-imbalance (upper half) and topological-imbalanee (lower half), evaluated by Macro-F1 score.
Ours
Mosk mode
Drop edge
DoOA
51.
112
MenoF1 [N]]
Minoee psa
22.
PROTEINS
NCII
Figure 10: Ablation experiments on class-imbalance (upper half) and topological-imbalane (lower half), evaluated by Micro-F1 score.
I.5 ADDITIONAL EXPERIMENT OF GRAPH PERTURBATION STRATEGY
further evaluate the effectiveness the two graph augmentation strategies. The Mask Node variant denotes our approach with randomly dropped nodes, while the Drop Edge variant refers to the method that randomly masks edges. In both cases, the perturbation operations are applied in a non-learnable and static manner. Degree-oriented Optional Augmentation (DoOA) (Xu et al. 2024 improves upon these baselines by customizing the node masking and edge dropping probabilities for each graph based its structural Spperties--seciffcclly the node degree distribution. As shown
35

---

Published as a conference paper at ICLR 2026
in Figures 9 and T0 all variants underperform compared our full model to varying degrees. This demonstrates the effectiveness of our proposed learnable node and edge perturbation strategies, which adaptively learn data- augmentation policies from the imput graphs. These mechanisms facilitate the learning of richer topological representations and contribute to improved overall performance.
PROTHINS
P80
SC
CUELAS
Tapk,
Topk
Tapk;
Number of Topk,
Number Taps;
Number Tapk;
Number Taps
Figure 11: Sensitivity study of TopK1 on class-imbalance (upper half) and topological-imbalanee (lower half).
PROTEINS
D&D
SCH
COLLAB
Minze Dill M Su
Tops;
Number Topk:
Number Tooks WCll
Number Topka
PROTEIN
D40
OELAR
Bigl
Mitmaa Mi
Stuca
Numbbr of Topk,
Number TTpp,
Number of Topk,
Number Tapk,
Figure 12: Sensitivity study of TopK2 on class-imbalance (upper half) and topological-imballnce (lower half).
1.6 HYPERPARAMETER EXPERIMENT
We further evaluate the sensitivity of two key hyperparameters, TopK, and TopK2, which are involved in the purification operations during the attention computation. Specifically: TopK, denotes the number of most-matching graphs selected for each prototype during the prototype-awarenes stage. TopK2 represents the number of most-affine prototypes selected for each graph during the prototype-balance stage. In this analysis, we fix to its optimal value for each dataset and report the performance under different settings of TopK, and TopKg. As shown in Figure II and Figure 12 we observe that TopK, typically performs well within the range of 16 to 64. Values of TopK, that too large may lead to the inclusion of redundant prototype features, which can degrade model performance. On the other hand, TopK; generally achieves strong results when set between 4 and 12. Furthermore, we conducted hyperparameter experiments on the D&D dataset under an extreme class-imbalance scenario, exploring different combinations of TopK, and TopKg. The model achieves the best performance when TopK, is set to 32 and TopK2 to 8, as illustrated Figures I and 14 These results indicate that carefully selecting these hyperparameters is essential for capturing representative patterns and achieving balanced learning.
1.7 APPLICABILITY ANALYSIS OF UNIIMB
We further evaluate the impact of different GNN architectures and Graph Transformers (GTs) as backbones on model performance. Specifically, we experiment with GCN (Kipf & Welling 2016). GraphSAGE (Hamilton al. 2017 , GraphGPS (Rampásek et al. 2022 Exphormer (Shirzad et al. 2023), and Graph-Mamba ( (Wang et al. (024b) as alternative graph learning backbones in place of the GIN architecture used in our original design. The results are shown in Table 2 and Table
36

---

Published as a conference paper at ICLR 2026
0.65
4652
0..5
08.66
3786
62.70
Topky
Topks
Tapk
Figure 13: Macro-F1 on the D&D Dataset for TopK; and TopK2 Hyperparameter Combina- tions.
Figure 14: Micro-F1 on the D&D Dataset for TopK, and TopK2 Hyperparameter Combina- tions.
We find that both GraphSAGE and Graph- achieve superior performance: GraphSAGE benefits from its efficient neighborhood aggregation that preserves informative local structures, while Graph-Mamba excels at jointly modeling node features and long-range structural dependencies. Importantly, when integrated with our framework, all backbone variants demonstrate significantly improved performance on imbalanced graph classification tasks. This result validates the effectiveness and generalizability of our method across various backbone architectures.
23
Table 22: Graph classification performance of various backbones wiht our Unilmb under Class Imbalance datasets with extreme imbalance degree.
PROTIINS
REDOITEII
IMDB-MI 131
Methol
Ma
Mico-F1
MacraF11
Macs-FF
MiceFF
MamnFI
MicrFF
Mice-Fi
GIN
23:33
25.30.533 4:82
5099 7.46
41.35
41.2
-3.55
51.3144112
15:25 c::0
15922-54
GUS
21.37 14,7
25.61 33577
2.20
2.27
19:25
12.881150 1530
15.1911.36
18.161344
55.3514.30
1435
12.92
Cregiigg geg
3473-4.11
2.5
1144
23225-2211
36.77
12:6-172
1691 5332
10133 4344
mm
7 9m
+2.1
25.99
11.36
1131-191
Unikal
Exptermer
24.38
0.79
4.28
10.79
27.33
13.61 17:41
42.00
20081 +5:5
2434
Eeghamser Tallmb
66585
00.89
6433
68.19
Gragh-Stanbe Liiiish
16882
75.88
12.154444
28,88
Table 23: Graph classification performance of various backbones wiht our Unilmb under Topological Imbalance datasets with extreme imbalance degree.
Methed
Macn-Fi
MicsFFi
MlawFi
MiesFi
Mixo-Fi
MireFi
Ma
Ml
MavoFFl
MinvFi
MaxeFFI
MinoFi
GIN
55445 2.10
36,00
2418
51.95 455)
4485
6100--30
01
043222211
6748 4428
20000
455
24
Tiillab
76.89
9111
764991111
61.9-954
9699
7633
1558-110
0685-19
OCN
$1,20
58.02
1500
1254
1835
60.9
13.12-144
2100
63511931
66.10 10.6
66.00
68.99
1128
1204
11
GUN Unllmb
6617
76.99
115
6192-330
7130
76.55
09556133
133
72.28
Gogbhlig
$2.34 606
06.00 4411
54.57 180
62,07
03
65.09
4311
10.0
23,0 464
GriphSegr Uniaas
09.59 37A
79665
9191
71.01
03.16
7579
71.38
1334
1913 113
151991331
1425
60.25
63.55 12.97
1240
1611
08411532
2611-8.2 1834
5666
1128
16:57 2023
14
GrophG88 Eviled
91136
6187-761
06.25
12
6621
2.73
19-41
4.6
06.36
550
21
30.71
15.7
25.52 500
21:00+200
Explosner Unilubb
07.25 2:11
78.92
53500
08:96
3338
09.81
701
68.59
2318
43:43
3362
Grah Mande.
66445
72.09 4.16
63.95
4073
54.484955
282
03.03
07.26
15.04
23.7
9482
11.0
4221
Grigh Marcha Unilbah
16.36 9117
76.92
110
71.38
613801111
57
11:00110
1.8 APPLICABILITY ANALYSIS OF DYNAMIC BALANCED PROTOTYPE
To further validate the effectiveness of our core module DBP, we conducted experiments by incorpo- rating DBP into GIN, GCN, and GraphSage on the REDDIT-B and IMDB-MULTI datasets under extreme class-imbalance scenarios. The results, as shown Table 24 indicate that only adding DBP to these vanilla GNNs significantly improves performance, demonstrating its effectiveness in mitigating imbalance. This also confirms that our two-stage design, prototype-percepting and prototype-bat balance, enables balanced interactions between each graph and its most similar prototypes, encouraging the prototypes to capture critical representations and thereby enhancing learning for tail classes.
37

---

Published as a conference paper at ICLR 2026
Table 24: Analysis Only DBP Performance on class imbalanced datasets with extreme imbalance degree.
REDDIT-B
IMDB-MULTI
Macro-Fi
Micro-FI
Micro-FI
GIN
33,19
4:38
6.6
GCN
71.49 1199:13:
1.40
29.27 3.17
8386
GruphSage
2221
29.62
L.50
12.26 2.72
14.01 2332
GraphSagg + DBP
64.29 64.29:191 1.92
78.28
5337
30.69 : 1.08
44.26 0384
1.9 GRAPH STRUCTURE VISUALIZATION
We further analyze the effectiveness of prototypes. Taking the highly class-imbalanced COLLAB dataset an example, extract the atfention matrix between prototypes and graph instances for prototype feature extraction (refer to Eq. 6). Based on this matrix, for a given anchor prototype we retrieve the three minority-clan graph instances with the highest affinity to that prototype, as shown in Figure 15 (c) We further examine the graph structures of these example graphs and find they all contain the same local topology (as shown (d)), a topology that important each graph (because each node has high degree). This indicates that prototypes can capture discriminative feature for minority classes, thereby improving minority-class classification performance.
(s) Graph
(b) Graph
(e) Graph
(d) Co-occuriing substructures
Figure 15: Visualization analysis of prototype effects.
LIMITATION DISCUSSION AND FUTURE WORK
Broaden Dataset Evaluation Across Diverse Scenarios Our current datasets primarily consist of homogencous graphs, which do not fully capture the diversity and complexity of real-world networks. In the future, we will evaluate the performance on heterogeneous graphs that contain multiple types of nodes and edges, enabling more comprehensive assessment across varied scenarios.
Compare Imbalanced Graph Learning Algorithms Across Various Tasks We aim to extend the comparison of imbalanced graph learning algorithms beyond classification to include tasks such as few-shot learning, dynamic graph learning, and anomaly detection. By broadening the range of tasks, gain deeper insights into the progress of the field and better understand how different algorithms perform across diverse applications.
FAZI
46.63
25.40
77.14
Figure 16: Macro-F1 class-imbalance and topological-imbalance datasets with extreme imbalance degree.
38


Reviews:
[{"review_id": "6vo4JVXTUy", "reviewer": "Reviewer_qbQQ", "summary": "The paper introduces UniImb, a unified framework tackling both class and topological imbalance in graph classification. The core novelty lies in the Dynamic Balanced Prototype (DBP) module that learns balanced prototype representations to reduce bias toward majority graphs. UniImb enhances graph representations and achieves up to 41.6% and 20.6% gains on class- and topology-imbalanced tasks, respectively.", "questions": "(1) Since the main contribution lies in DBP, it might be more effective to focus solely on this technique rather than stacking it with multiple others, such as TopoEncoding. Emphasizing DBP’s independent efficacy would strengthen the justification of its contribution—particularly if experimental evidence can directly show that DBP alone mitigates class imbalance.\n\n(2) It would also be valuable to provide an intuitive and visual illustration of how DBP alleviates both quantitative and topological imbalance. For instance, does the learned DBP reflect specific topological patterns that are desirable for minority-class test graphs but absent in the minority-class training data? Such visualization could make the mechanism and impact of DBP more tangible and convincing.\n\n(3) My main concern, as noted in the weakness section, lies in the real-world practicality of the study. While the analyzed graphs are important, they do not capture realistic structural properties, such as the 3D conformations of proteins and molecules, which are crucial in real-world applications. This limitation raises concerns about the method’s applicability to domains where complex structural dependencies and class imbalance are pervasive.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "(1) The experimental results are comprehensive.\n\n(2) The imbalance graph classification is a very important problem, and the imbalance phenomenon is very widely encountered across many real-world applications, such as drug discovery.\n\n(3) The paper also introduces another air-pollution graph classification dataset, which might be of interest to the following research.", "weaknesses": "(1) The main weakness is the novelty of the study. Although imbalance graph classification is indeed a very important problem, the more realistic setting, as the field advances, should focus on graphs aligned more with reality, such as graphs with 3D structures.\n\n(2) The proposed method stacks many different techniques together, which makes the core contribution hardly justified. Although the ablation study has demonstrated that the most critical part is DBP, it might be more interesting to intuitively demonstrate how DBP can benefit the imbalance graph classification. Specifically, what new signals have been learned by introducing DBP techniques?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2VDVrAmDaB", "reviewer": "Reviewer_LB4s", "summary": "This paper addresses a comprehensive graph classification scenario involving both class imbalance and topological imbalance, and proposes a method named UniImb. Specifically, UniImb integrates multiple carefully designed techniques and introduces a prototype learning strategy to capture representative features, while a balancing constraint ensures its effectiveness in such imbalanced settings. Extensive experiments demonstrate that the proposed model achives better results.", "questions": "Refer to the above box.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed method to extracting prototype features is skillfully designed to address the challenge of graph imbalance. \n2. The experimental evaluation is convincing, demonstrating consistently performance gains over existing methods.\n3. Mathematical notation is clearly defined and used consistently throughout the manuscript, which is overall well-structured and clearly written.", "weaknesses": "1. Adding an algorithm pseudocode box can more clearly show the technical details.\n2. Why can't existing models comprehensively handle both types of imbalance? Although the authors provide an explanation in the introduction, a more intuitive clarification would be helpful.\n3. The primary experimental results (Tables 1 and 2) evaluate these two problems in isolation. While Section 5.4 does introduce an \"intertwined\" scenario. This makes it difficult to assess the framework's synergistic advantages. The current evidence in the manuscript demonstrates that Unilmb performs well on two separate tasks, rather than a truly unified solution whose main strength lies in handling their complex co-occurrence.\n4. The paper introduces a Personalized Graph Perturbation Strategy, claiming it customizes augmentation for small graphs or minority-class graphs . However, the paper fails to explain or validate why a perturbation strategy based on average degree would be suitable for addressing class imbalance. \n3. The AirGraph dataset, which exhibits strong practical application potential, is newly introduced; thus, a more detailed exposition of this experiment is warranted in the main body.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z6VyIkmGk1", "reviewer": "Reviewer_1cXq", "summary": "The paper presents a method for improvements in imbalance learning on graphs. Specifically, the problem is the graph classification setting that assumes the data can have both class and topological imbalance. The method contains multiple components, but the key innovations are the learned augmentations and dynamic prototypes. Across a large set of datasets the authors show significant improvements over prior work in this domain.", "questions": "1. Was there any issue with running the TopoImb baseline? The results there look unreasonably low. \n\n2. Is it possible to compare against more recent methods, e.g., [1] or others that have shown performance gains over ImbGNN?\n\n3. Although focused on graph-level tasks, providing a bit more related work on other imbalance problems could be beneficial (e.g., representative works in node, edge, and potentially subgraph), but the discussion in relation to the vision domain seems less relevant. \n\n4. Can you elaborate more on specifically why a contribution is unifying class and topological imbalance? Note that the prior work ImbGNN also discusses this innovation when building the connection between them. \n\n5. As compared to cherry picking the best result to make claims like \"UniImb delivers up to 41.62% relative improvement on ...\" providing averages across all datasets would be a preferred and less biased approach for the reader to consume the summarized information (note that already the average improvement is quite significant and need not just report the \"up to\" results here to show the empirical success/benefits of the proposed method).\n\n6. The work \"DPGNN: Dual-Perception Graph Neural Network for Representation Learning\" is cited as providing a prototype based approach, but does not appear to present such method.  \n\n7. Most of the datasets are binary classification and have overlapping domains. Are there datasets with more classes that could be explored (potentially diverse domains)? How would this prototype-based framework handle many classes? You may consider looking here [2]. \n\nBased on the response to the comments and questions above, I am open to revising my score. \n\n[1] SamGoG: A Sampling-Based Graph-of-Graphs Framework for Imbalanced Graph Classification\n[2] Graph Prototypical Networks for Few-shot Learning on Attributed Networks", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. The motivation for not applying graph augmentations/perturbations uniformly makes sense, which leads to their proposed personalized strategy to learn them based on the graph characteristics. \n\nS2. The dynamic balanced prototype strategy is interesting, having solid theoretical support, and appears to provide great empirical gains. \n\nS3. The level of detail in conducing such comprehensive experimental results is very appreciated and provides the community with rich results to further understand the present method beyond a typical work.", "weaknesses": "W1. Generally the problem of imbalance classification on graphs is interesting, but at this point in time there have been numerous works in this direction. Thus, although there is empirical performance gains, the problem itself is well established and the methodology is somewhat lacking key novelty by combining so many techniques in an attempt to obtain the highest performing method (as compared to just focusing on the core innovations of the work). \n\nFor further less critical weaknesses, please see the comments and questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6X0FqBoX79", "reviewer": "Reviewer_5aWh", "summary": "In this submission, a unified framework, UniImb, for handling both class imbalance and topological imbalance in graph classification tasks is prposed. The core innovation lies in the introduction of a dynamic balanced prototype module, coupled with a prototype load-balancing regularizer, which ensures equitable influence from tail graphs. Theoretically, the authors provide an information-theoretic justification for the prototype balancing mechanism based on the Information Bottleneck principle. Extensive experiments demonstrate that UniImb achieves competitive performance.", "questions": "Q1. What does \"Basis\" mean in Table 1?\n\nQ2. What is the final form of the model's loss? How is the classification loss combined with Equation 11? \n\nQ3. Why do the authors claim that \"Methods focused on topology imbalance adapting to small graphs but typically ignore imbalanced class distributions\"? Why can't existing models handle these two types of imbalance effectively?\n\nQ4. Imbalance learning baselines such as upsampling have only been evaluated on GNNs; can their effectiveness be extended or improved when applied to graph Transformers?\n\nMinor errors:\n\n(1) The red boxes in the article's appendix can be deleted; they seem to serve no purpose.\n\n(2) The table caption content is overly redundant.\n\n(3) Bold text in table captions should be used with caution.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "S1. The authors are the first to unify the modeling of class imbalance and topological imbalance in graph classification, filling a notable gap in existing research. \n\nS2. The proposed dynamic prototype method is elegantly designed and well-grounded theoretically. And the proposed UniImb framework is highly versatile and plug-and-play, which can be seamlessly integrated into various GNN backbones. \n\nS3. Another key strength of this work is its exceptionally comprehensive and thorough experimental evaluation, encompassing 13 datasets, over 20 baselines, and the introduction of a new large-scale graph classification benchmark. Across all settings, UniImb delivers significant performance improvements.", "weaknesses": "W1. The model contains several key hyperparameters. Although the authors analyzed the sensitivity of each individually in the original paper, a sensitivity analysis of combinations of two hyperparameters is also necessary.\n\nW2. While the overall framework is interesting, most individual components seems like incremental.\n\nW3. All datasets are from common benchmarks. It would be stronger to include an industrial-scale or biomedical use case with real-world imbalance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)