Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
Authors: Weihuang Zheng
Jiashuo Liu
Jiaxing Li
Jiayun Wu
Peng Cui
Youyong Kong
Venue: ICML 2025 poster (2025)
Keywords: graph neural networks
out of distribution
sample reweighting
Primary area: deep_learning->graph_neural_networks

Abstract: Graph Neural Networks (GNNs) are widely used for node classification tasks but often fail to generalize when training and test nodes come from different distributions, limiting their practicality. To address this challenge, recent approaches have adopted invariant learning and sample reweighting techniques from the out-of-distribution (OOD) generalization field. However, invariant learning-based methods face difficulties when applied to graph data, as they rely on the impractical assumption of obtaining real environment labels and strict invariance, which may not hold in real-world graph structures. Moreover, current sample reweighting methods tend to overlook topological information, potentially leading to suboptimal results. In this work, we introduce the Topology-Aware Dynamic Reweighting (TAR) framework to address distribution shifts by leveraging the inherent graph structure. TAR dynamically adjusts sample weights through gradient flow on the graph edges during training. Instead of relying on strict invariance assumptions, we theoretically prove that our method is able to provide distributional robustness, thereby enhancing the out-of-distribution generalization performance on graph data. Our framework's superiority is demonstrated through standard testing on extensive node classification OOD datasets, exhibiting marked improvements over existing methods.

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
Weihuang Zheng"
Jiashuo Liu 2
Jiaxing Li"
Jiayun Wu? Peng Cui Youyong Kong'
Abstract
become cornerstone models for node classification. How- ever, these GNN models typically assume that the training and test graph data are drawn from the same distribution, which does not always hold in practice. In real-world graph data, sample selection bias (Fan et al., 2022; He et al., 2020) and the methods used for graph construction (Qiao et al., 2018; Zhou et al., 2023) often brings distribution shifts between training nodes and test nodes. For instance, In WebKB (Pei al., 2020) datasets, web pages (nodes) and categories (labels) are heavily affected by the university they originate from, leading to distribution shifts among nodes drawn from different universities. Therefore, address- ing distribution shifts is crucial to improving the practical effectiveness of GNNs for node classification.
Graph Neural Networks (GNNs) are widely used for node classification tasks but often fail gen- eralize when training and test nodes come from different distributions, limiting their practicality. To address this challenge, recent approaches have adopted invariant learning and sample reweight- ing techniques from the out-of-distriutiion (OOD) generalization field. However, invariant learning- based methods face difficulties when applied graph data, they rely on the impractical as- sumption of obtaining real environment labels and strict invariance, which may not hold in real- world graph structures. Moreover, current sample reweighting methods tend to overlook topologi- cal information, potentially leading to suboptimal results. In this work, we introduce the Topology- Aware Dynamic Reweighting (TAR) framework to address distribution shifts by leveraging the inherent graph structure. TAR dynamically ad- justs sample weights through gradient flow on the graph edges during training. Instead of relying on strict invariance assumptions, we theoretically prove that our method is able to provide distribu- tional robustness, thereby enhancing the out-of- distribution generalization performance on graph data. Our framework's superiority is demon- strated through standard testing on extensive node classification OOD datasets, exhibiting marked improvements over existing methods.
To address the distribution shift problem on the node clas- sification task, recent works (Wu al., 2021; Xia et al., 2024; Wu et al., 2022: Liu et al., 2023) adopt the idea of invariant learning. Invariant learning (Arjovsky et al., 2019; Liu et al., 2021a) stems from the causal inference litera- ture, and now becomes one of the key approaches to solve distribution shifts problem on graphs. The core concept of invariant learning is to identify invariant features with stable prediction mechanisms across different environments (distri- butions), thereby mitigating performance degradation under distribution shifts. Despite its success, methods based on invariant learning are built upon strong invariance assump- tions that lack further validation for their actual validity (Liu et al., 2024). In different environments, invariant features may not exist, the environment can influence the mapping relationship between features and labels. For example, in social networks, users with the same profile (age, interests, etc.) may exhibit different behaviors depending on their location. In region A (where personalized ads are allowed), they may frequently click on ads, while in region B (where privacy regulations limit ad tracking), their click-through rate may significantly decrease.
1. Introduction
Node classification tasks have widespread applications in real life, such as advertising recommendation (Jiang al., 2023), social network anomaly detection (Tang et al., 2022), and more. Recently, Graph Neural Networks (GNNs) have
Besides, sample reweighting methods, particularly those based on distributionally robust optimization (DRO), have also been used to address distribution shifts in graph data (Hu & Hong, 2013; Wang et al., 2024; Gui et al., 2022). These methods reweight the probability distribution of train- ing samples within a predefined distribution set to identify the worst-case distribution under which the model performs the poorest. By optimizing against this worst-case distribu-
Equal contribution 'School of Computer Science and Engineer- ing, Southeast University DDpaatmentt of Computer Science and Technology, Tsinghua University. Correspondence to: Youyong Kong <kongyouyong @ scu.edu.en>.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
tion, these methods theoretically ensure improved robust- ness to potential distribution shifts. For example, KL.-DRO (Hu & Hong. 2013) improves robustness by assigning higher sample weights to high-risk samples. DR-GNN (Wang et al., 2024) builds on the idea of KL-DRO to address distribution shift issues in recommender systems. Furthermore, (Gui et al., 2022) demonstrated the effectiveness of Group DRO (Sagawa et al., 2019) for node classification under distribu- tion shifts by assigning greater losses to the group of nodes with the highest loss in the graph. However, these methods directly adopt general -omari reweighting strategies with- out considering the topological structure of graphs, which may lead to suboptimal results.
2. Preliminaries
Notations. X E X denotes the covariates, Y E y de- notes the target, Y) and Y)) represent the joint source distribution and the target distribution, ab- breviated with P. and P. respectively. The prediction model is denoted by so(.) : X + V. for which we use graph neural networks (GNN) throughout this paper. [N] - (1, N denotes the set of integers from I to N. A weighted finite graph is denoted by Go = EV, W), where V = {u uN) ) the node set, E is the edge set, and W = (wij)(. are the edge weights. (i) denotes the set of adjacent nodes for the i-th node.
Problem setting. In this paper, we focus on the problem of distribution shift in node classification. Specifically, our goal is to learn a predictor ʃă that generalizes well under target distribution. Formally, we aim to find fā that satisfy:
In this work; we focus on the problem of distributional shift in node classification tasks. Instead of adopting invariant learning that relies strict invariance assumptions, as in most current approaches, we adopt sample reweighting method. Specifically, we propose the Topology-Awar ware Dy- namic Reweighting (TAR) framework to remedy the issues above. During training. TAR assigns more densities to high-risk nodes, thereby prompting the prediction model to prioritize these nodes. Additionally, we theoretically prove that this method can identify the local worst distribution, and by optimizing this distribution, the model can achieve improved robustness. Our main contributions are as follows:
S0 = arg min E(XY)-,, [X(o(X),))), fo
where e(..-) is a predefined loss function. As the target distribution is not available during training, previous ap- proaches typically optimize the empirical loss on the train- ing source distribution (x),))=,, (x()o(X), Y)] as a surro- gate. However, the distribution shifts between IP, and P. can lead to the failure of the predictor.
Method:
Our TAR framework involves a minimax pro- cedure, where the inner maximization problem learns sample probability densities (also referred to as sample weights) under the entropy and topology constraints, and the outer minimization problem optimizes the GNN model under the learned distribution. For the reweighting scheme (inner problem), as illustrated in Figure 1, we perform gradient flow along the graph edges assign greater sample weight nodes with higher loss. In this way, we incorporate the topologi- cal structure information into the learning of sample probability densities. Furthermore, we propose lever- aging graph extrapolation to expand the given data distribution, further enhancing the model's 's resilience to potential distribution shifts.
Sample Reweighting. Formally, sample reweighting is formulated as:
f0 = arg min max E(x,Y)= ()) PES(E,)
where S(P,) indicates the potential distribution set around the source distribution. By adjusting the probability den- sities q(z, y) (i.c. reweight samples) to simulate the worst potential shifts in the inner maximization, this method can ensure the predictor's 's robustness when encountering poten- tial unseen distributions. Typically, S(P) is constrained by d(P, P.) < d, where d(-,,) -) is distance metric measuring the difference between two distributions, d is a predefined constant. This constraint ensures that the reweighting avoids unrealistic distributions, focusing only on plausible shifts.
Discrete geometric Wasserstein distance. We review some key concepts and introduce the discrete geometric Wasserstein distance (Chow et al., 2017), where we adopt the notations used in (Chow et al., 2017; Liu et al., 2022).
Theory:
In Section 3.2, we theoretically prove that our gradient flow procedure is equivalent to finding the local worst-case distribution, which enhances the distributional robustness of our GNN model. We also characterize the error rate introduced by our gradient flow as CT, (Tin is the number of steps).
The (empirical) probability set supported on all nodes of Go is denoted as:
{
N Σ 1=1
Experimental results on standard OOD node classification datasets demonstrate the effective- ness of the TAR framework, showing its superiority in addressing distributional shift problems.
Experiments:
(M)) P(Go) = - 1 P-
0, for [N]
which contains all empirical distributions on the node set V, and the interior of P(Go) is denoted as PooG.
2

---

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
g()
{
Reweight
Gruph Go
Graph 22
Graph G₁
Graph G₃
X
X
X
Sample Weight
Gradient Flow
① Loss Value
Figure 1. Illustration of TAR. Given graph go. first conduct graph extrapolation generate G1. Next, we compute the classification loss, initially assigning equal sample weights to each node. Subsequently, we perform topology-awand reweighting. Specifically, nodes with higher loss values attract more sample weight through gradient flow from neighboring nodes. Finally, nodes with higher losses assigned greater weights, while nodes with identical loss values exhibit different weights due to variations in their topological positions.
velocity field (eu))eev E NNN graph Go is skew-symmetric matrix on the edge set E:
where the infimum is taken over all velocity fields Go- and (py(p) is a pre-defined interpolation function between pi and py. Note that p(() is a continuously differentiable curve p(t) [0, 1] - Po (Go), which characterizes the probability densities at time
Vig if (i, E,
Vig
otherwise.
Given the probability function p PG))) and a velocity field the flu function is defined as the product pe E NNN;
pu = Eu (P))
where (p) is predefined "cross- sectional area", typically interpolated with the associated nodes' densities Py. To ensure the positiveness of during optimization, we adopt the upwind interpolation from statistical mechanics (Hsu, 1981): (p) = I(uyy > O)p, + I[v.y- I O)p, throughout this paper, which relies on the corresponding velocity field. Intuitively, this characterizes the "flux" of sample density from node i Based on this, the divergence vector of pe on graph Go is defined as:
Remark 2.2. In contrast with the conventional Wasserstein distance defined within Euclidean space, the geometric Wasserstein distance necessitates that the transportation of probability density along the geodesic determined the graph structure Go. In particular, the constraint de divc, (pv) 0 imposes the condition that the change in probability density remains continuous with respect to Go-
3. Method
(iva,(wu) =
Σ JEV:G.JEEB
V ERN
Motivated by the discrete geometric Wasserstein distance in Definition 1. propose the Topology-Aware Dynamic Reweighting (TAR) algorithm to deal with node classifica- tion tasks under distribution shifts. As shown in Figure 1, our proposed TAR consists of two key stages. In the first stage, we employ graph extrapolation to expand the orig- inal data distribution, thereby generating more diverse and comprehensive
which is supposed to lie in the tangent space of (Go). Intuitively, the i-th element in divg, (pv) sums over all the in-fluxes and out-fluxes along edges to a certain target node i, with each source edge transporting a probability density M (p).
Now define the discrete geometric Wasserstein distance:
of node samples. Building upon this expanded graph data, the second stage introduces topology-awer dynamic sample reweighting mechanism, which aims to simulate the potential worst-case distribution. By training the model to perform well on this constructed distribution, our method enhances the model's 's robustness against unseen distribution shifts.
Definition 2.1 (Discrete Geometric Wasserstein Distance). Given finite graph Go. for any pair of distributions p", p'pp' E P. (Go), the discrete geometric Wasserstein dis- tance is defined as:
GW3 (v" p²) inf [/1 Σ E3
Consider source data D, = {(r,, N and the corre- sponding graph structure Go (V, (V,,WW) Denote the empirical marginal distribution as TP the overall objective
E
dp dt + divg. (pe) = p(0) = pp.p(1)))))))m)pp)))(())) =p'
s.t.
3

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
of our TAR algorithm is formulated as:
Algorithm I Topology-Aware Dynamic Reweighting (TAR)
Scheme
min eEB (Go)
max
qi (
Input: Labeled training nodes D = {(r,, b) learn- ing rate 7, gradient flow iterations Tin entropy term B,
graph structure GD = (V, E, W).
Weighted loss
(1)
Initialization: Sample probability densities initialized as (I/N, ................................ 1/N)". Model parameters initialized as g())
3
qi log -
X GWG ())..
for = 0 to Epochs do
1. G = g(Go), detailed in 3.3.
Entropy penalty
Topology penalty
2. Simulate gradient flow for Tim time steps according
where 9 is the hyper- paremetrr,,,er, and the objective function in general is minimax optimization over model parame- ters 0 and sample probability densities Note that for the parameter A, we set as À = 2F in our optimization (for details, please refer to Section 1. During training, the inner maximization assigns more densities to high-risk sam- ples, thereby prompting the prediction model to prioritize these points. This approach aims for a uniformly robust performance across all samples the graph and helps miti- gate potential distribution shifts. Moreover, to mitigate the risk of overemphasizing unrealistic distributions (e.g., noisy nodes accumulating excessive densities), we introduce en- tropy and topology penalties as regularization terms. These penalties integrate topology information for smooth sample weight assignments along the graph structure.
to Equation 3 and 4 to learn an approximate worst-case probability weight q"tn
3. g( + g(0 V. 9c)
end for
involves a minimax optimization, where we iteratively per- form gradient ascents (on q) for the inner maximization and descents (on 0) for the outer minimization. The pseudo-code of our algorithm is shown in Algorithm 1.
Inner maximization problem.
For easy notion, we define
N ((0, = Σ ((((o(z))) v.) 3 Σ
q. log q:-
Illustrations. Here we make some remarks on our objec-
Generally, the goal of the inner maximization problem in Equation is to maximize L(0,, q) and to minimize the topol- ogy penalty GW2, (f...) w.r.t. sample densities q. Instead of directly computing the topology penalty, we solve the in- ner maximization via gradient ascents on g in the geometric Wasserstein space (Po(Go), GWa.). where the topology penalty gW& Go (Ps,4) 9) is approximated by the length of the gradient flow trajectory in the metric space.
tive function:
(a) Entropy penalty: (- E- qi log (:) represents the en- tropy of empirical probability distribution y. As illustrated in Section 3.2, this term serves as a non-linear graph Lapla- cian operator that encourages sample weights to be smooth along the manifold, avoiding extreme sample weights in the weighted distribution.
(b) Topology penalty: GW2, (P...4 represents the optimal Gi
transport distance between the source distribution P,, and the weighted distribution q. measured along the graph struc- ture. This term explicitly integrates topology information to enforce minimal changes in sample densities along the manifold. As detailed in Section 3.1, this term transfers the optimization of sample densities from Euclidean space to geometric Wasserstein space. Here, densities are con- strained to change exclusively along the graph structure. This enforcement encourages local smoothness of sample densities relative to the manifold, which helps to mitigate against potential noisy samples and edges.
As stated in Definition 2.1, the continuous gradient flow is denoted by q : [0, 1] Pa (Go), and q(t) represents the sample density at time t E [0, 1]. In order to derive empirical optimization approaches, we introduce the time- discretized gradient flow, denoted by q" [O, + (Go), and the superscript r is the value of time step (here we introduce this superscript because different time steps refer to different úime-diccrttzzed gradient flow function). For the approximate optimization, we leverage this time-discretized gradient flow (with time step 7) of -L(0, q) in the geometric Wasserstein space (Po(Go), (.GG)))Wa.)) as:
3.1. Towelegy-Aware Dynamic Reweighting
qt + +- arg max C(0.g) (GP)(Go)
1 2T
GWe, (1),
The main challenge of Problem 1 lies in the computation of discrete geometric Wasserstein distance GWZ, (P.,, g), which itself involves an complicated optimization problem and does not have an analytical form. In this section, we pro- pose to leverage Wasserstein gradient flow to approximately solve the inner maximization problem. The whole algorithm
which aims to obtain the "local" maximum of L( ④) around (t) at time and restricts the topology distance gW2 Go (q (t), g). We derive the analytical form of Equa- tion as T + O. For the case of notion, the sample density of the i-th node at time f, originally denoted by qf (t), is
4

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
Geometric Wasserstein Space (Po(Go), gW
Theorem 1
local distributional robustness
P
gw? 6o
(1)
q'
gradient flow
Theorem approximation emoor rate
(2(2)
local distribution ball
q' (T)
Figure 2. Illustration of the gradient flow in the geometric Wasserstein space (G.) where each point denotes probability distribution in (Co). and the distance is measure by the discrete geometric Wasserstein distance. The black cirele denotes the local distribution set around a distribution, and the blue arrow represents the one-step gradient flow. q" (T) denotes the approximated inner maximizer obtained by our algorithm, and q" denotes the ground inner maximizer (defined in Theorem 3.3). In Theorem 3.2, we demonstrate that the one-step gradient flow is equivalent to distributionally robust optimization around local uncertainty set, and in Theorem 3.3, we characterize the approximation error rate between q (T) and q"
abbreviated as q:()), and then Equation becomes:
(d) Due to the random sampling of labeled nodes during training for node classification tasks, it means that for cer- tain nodes we cannot compute the loss, which disrupts the connectivity and hinders the calculation of Equation 3, we intuitively set the loss for these unlabeled nodes to the mean loss of the labeled nodes, and this approach has proven to be adequate. For other potential solutions, please refer to the Appendix E.
(t) dt
Σ Wiyei [(vi 0)g, + +1(v s 0)g:
)((j)E
Bij = 1y-1,- - + 3(log 9 lo) for (i.j) E E
(3) where E the edge set of graph Go. wiy is the edge weight between node i and j. I(-) the indicator functtion, and f represents the prediction error on the i-th node. Intuitively, Vi can be viewed the transferring velocity of the sample density from node to node i.
Based on Equation 3, we can solve the inner maximization problem via gradient ascent as:
q: (0) 1/N, (t 1) 9:(1) + (q: for E [N].
(4)
Let A = 27 Equation 2 exactly aligns with the goal of our inner maximization problem in Problem 1. Specifi- cally, the original topology penalty calculates the distance GWC. (P..." (t)) between P, and q" (t), and our gradient flow approximates it via Et= GW2 (q (i- 1), o (i)) (see blue curves Figure 2). In Theorem 3.3, we characterize the error rate of this approximation.
(5)
In addition, we demonstrate the equivalence between Equa- tion and distributional robustness in Theorem 3.2, justify- ing how our proposed TAR can provide robustness against distribution shifts. And in Theorem 3.3, we characterize the error rate of our approximation e -CT, which allows relatively accurate approximation with finite Tin steps.
Remark 3.1. Here we make some remarks on Equation (a) The gradient of the i-th node's probability density de- pends its neighbors in graph Go. This corresponds with our motivation that the reweighting scheme should incorpo- rate topology information. Furthermore, since the transfer between neighbors, the probability density p remains lo- cally smooth w.n.t the graph structure (or manifold), which avoids overemphasis on some noisy samples.
Outer minimization problem. For the outer minimiza- tion problem, we perform gradient descent on model param- eters 0. According to the overall objective in Equation I, the loss function is simply a weighted average:
N ((++1) <- (()) 7 Vo Σ (Tin)) ((Jo(ra))))) (6)
(b) Combined with our topology penalty, the entropy penalty acts non-linear graph Laplacian operator to further the smoothness of probability densities along the manifold. (c) The gradient flow in Equation is implemented by mes- sage propagation, which scales linearly with sample size and enjoys parallelization by GPU. For a detailed complex- ity analysis, please refer to Appendix
where 7 is the learning rate, and q: (Tin) denotes the proba- bility density of the i-th node (after Ti steps gradient flow).
3.2. Theoretical Analysis
In this section, we investigate in- our proposed op- timization algorithm. As illustrated in Figure 2, we first
5

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
Table 1. The performance OOD benchmark datasets under Covariate Shift, using GCN as the backhone. We report the average accuracy (except for Twitch, where we use ROC-AUC as the evaluation metric) and standard deviations over 5 runs. The best results are shown in bold, and the second best results are shown in underline. OOM denotes out of memory.
Dataset
CBAS
WebKB
Twitch
Cora
Arxiv
Require domain information
Domain
color
university
language
word
degree
time
degree
Base
ERM
78.29:3.10
19.37=6.01
50.04::..53
64.72:1.55
55.33=0.90
71.48=0.30
57.59=00336
No
IRM
78.00+3.73
21.5949.30
50.97+3.62
65.70=00:4
55.34=0.88
71.36=0.13
57.34=00115
Yes
Invariant Learning
VREx
78.57:12.02
1.99
51.26.14.82
65.02.10.45
55.44:0..7
71.58:00.28
57.60:0.27
Yes
Domain Generalization
Coral
78.001:3.59 30.1615.67
53.4610.41
64.85:1:.33
55.21
71.48:0.40
57.23=0.16
Yes
DANN
77.711.3.13
33.49.19.61
50.47:1..14
64.72:10.39
55.32:10.92
11.68:00119
57,25:10,23
No
SRGNN
16,86++1.56
22.86+110.79
49,71+1181
64.63+-0.3 34
$5.22++1005
70.78=0.18
57,53+00.24
Yes
EERM
75.1443.29
33.97=11 12 42
00M
64.88-0.38
55,30=0.89
OOM
0OM
No
Graph OOD
FLOOD
83,1444.21
33.97=3.09
15.14
64.91=0:45
54.78=0.82
71.75=0.37
58.90=0.22
No
CIT
80.86122.96
28. 28.89:19.09
00M
64.3410.72
54.74=0.82
OOM
0OM
No
KL-DRO
17.7121.63
12.06211696
53.76:-4.19
64.8520.51
55. 15.1420.87
71.58=0.24
57.59=0.300
No
Sample Reweighting
GrouDDRO
77.144-2.67
25.24449.57
50.48++2.04
64.95+00.59
55.05+0083
71.46+0.33
57.25+0.23
Yes
Ours
TAR
8..431...09
37.46
1.18
65.64100.7
16.62
0.22
1..6
No
prove that each step of the gradient flow exactly finds the worst-case distribution within local uncertainty set (see black circle in Figure 2).
from the original source distribution. Denote the ground- truth worst-case distribution with the same distance
q" = arg
max
C(0,4),
Theorem 3.2 (Distributional robustness). For any y >> 0, > O and given f,, denote the solution of Equation 2 as
@:OW? Go (P::
Then we have:
9- = =arg max C(0, g) 79 (EPo(Go)
C(0.9(Tim) [(O,
> e-e
(8)
[(0, - P.))
Let = GW (p, we have
where C O is a constant and its value depends on the loss function f, hyper-parameter B, and sample size N. The proof can be found in Appendix G.
max gEPo(Go)
((B,9)- ~9W (p.
one-step gradient flow at time
Remark 3.4. We make some remarks here:
C(O, 9)-
(1) Since the goal of our reweighting is to maximize [(0, q) w.r.t. q. we utilize the increase of L to characterize how "approximate" is our optimization. In Equation 8, the de- nominator of the left-hand side represents the maximal in- crease, and the numerator is the increase attained through our approximation. As the ratio approaches 1.0, our approx- imation becomes increasingly precise.
g:GW &o (e
worst- case distribution within local distribution set
The proof can be found in Appendix G.
Theorem 3.2 shows that, for the inner maximization, our proposed gradient flow is equivalent to finding the worst- case distribution within a small distribution set. Therefore, the weighted average loss function in Equation captures the worst-case distribution that may occur in testing, which shares the similar idea with distributionally robust optimiza- tion (Duchi & Namkoong, 2018; Blanchet et al., 2019; Liu et al., 2022). This demonstrates the strength of our proposed TAR framework in dealing with potential distribution shifts.
(2) Our theoretical results show that the error rate is e -CT,,, which shrinks fast as the number of time step Tn increases. This further demonstrates that our optimization able to find good approximations in finite (usually small) number of gradient flow steps.
3.3. Graph Extrapolation
In the previous section, we proved that gradient flow is equivalent distributionally robust optimization around a local uncertainty set. Given that the source distribution is fixed, the scope of the uncertainty set is inherently limited. To address this limitation, we propose a Graph Extrapola- tion (GE) approach g(-) to expand the source distribution, thereby enlarging the uncertainty set and further enhancing robustness against potential distribution shifts. Specifically,
Then based on the results in (Chow et al., 2017, Theorem 5) and (Liu et al., 2022, Theorem 3.2), we move on to analyze the error rate of our approximation in Theorem 3.3.
Theorem 3.3 (Approximation error rate). Given the GNN parameter 6, denote the approximate sample densities in Equation 3 after Tin steps of gradient flow as ((Tn), and GW2 (Ps, (Ti))) is the geometric Wasserstein distance
6

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
adopt graph data augmentation techniques, DropE- dge (Rong al.) and MaskFeature (You et al., 2020), to expand the source data distribution i.e. generate new graph G1 according to G1 g(Gu). During training, we first con- duct GE to the input graph to construct an augmented graph at the beginning of each epoch, then conduct TAR on the augmented graph to train the model.
FRM
90-
'AR
Accurary (%%
Accurary (75)
30
[4
ERM
TAR
20
78
20
4. Experiment
(a) Tin on WebKB
(b) Thn CBAS
We conduct experiments on five widely used node classifi- cation datasets from GOOD benchmark (Gui et al., 2022)
ERM
FRM
ROC- AUC (96)
Accurrscy (75)
72-
to validate the effectiveness of TAR in improving out-of- distribution (OOD) generalization.
4.1. Datasets and Baselines
50-
71
10
20
10
20
Datasets. We use five node classification datasets under both concept shift and covariate shift (the detailed definition of these two shift are provided in Appendix B): : WebKB (Pei al., 2020), CBAS (Ying et al., 2019), Twitch (Rozember- ezki & Sarkar, 2020), Cora (Bojchevski & Günnemann, 2017) and Arxiv (Hu et al., 2020). We followed the GOOD benchmark (Gui et al., 2022) for data splitting, standard widely adopted in prior research (Sui et al., 2023; Liu et al., 2023; Guo et al., 2024). The statistics of these datasets and detailed information are provided in Appendix B.
(c) Tin on Twitch
(d) Ti on Arxiv-time
Figure 3. The effects of Th (the number of gradient flow) of our proposed TAR algorithm, we provide the results under covariate shift.
Graph OOD methods generally achieve better results, as they obtain the second-best performance in 7 out of 14 set- tings. We attribute this to their ability account for graph topology when addressing distribution shifts, leading to improved outcomes. This highlights the necessity of incor- porating graph topology information when tackling node classification tasks. Additionally, VREx also performs well, achieving the second-best results in out of 14 settings. However, this method relies on domain information (i.e., environment labels), which is often difficult obtain in real- world applications, thereby limiting its applicability. Group DRO achieves the second- best results in the remaining two settings, demonstrating the robustness of sample reweight- ing methods. However, since does not consider graph topology, its performance improvement remains limited.
Baselines. We compare our proposed TAR method with several commonly used baseline methods using GCN (Kipf & Welling, 2017) as the backbone (we also provide results using recently proposed sota graph transformer polynormer (Deng et al., 2024) as backbone, please refer to Appendix D). These include Empirical Risk Minimization (ERM) and two general invariant learning methods: IRM (Arjovsky et al., 2019) and VREx (Krueger al., 2021). Addition- ally, we compare with two sample reweighting methods, KL-DRO (Hu & Hong, 2013) and Group DRO (Sagawa et al., 2019). Furthermore, we evaluate our method against two common domain generalization approaches: DANN (Ganin et al., 2016) and Deep Coral (Sun Saenko, 2016),
As shown in Tables and our proposed method, TAR, consistently outperforms other basclines across all datasets under both concept shift and covariate shift. This highlights the effectiveness of TAR in addressing different types of distribution shifts. In particular, for the CBAS dataset, TAR achieves performance improvements of 4.29% and 3.00% under covariate shift and concept shift, respectively, com- pared the best-performing baseline. Similarly, for the Twitch dataset, TAR demonstrates improvements of 2.68% and 1 1.45%, respectively, over the best-performing baseline. These results underscore the superiority of our method and emphasize the importance of incorporating topological in- formation when addressing distribution shift challenges in node classification tasks.
as well as four graph-specific domain generalization base- lines: EERM (Wu et al., 2021). SRGNN (Zhu al., 2021),
FLOOD (Liu al., 2023), and CIT (Xia al., 2024). Among these methods, only KL-DRO, DANN, EERM, FLOOO,, and CTT do not require domain labels to address distribution shifts. Similarly, our proposed TAR method does not require domain labels. Please refer to Appendix C for more implementation details.
4.2. Performance Comparison
Table and 2 summarize the results of our method and other baselines on four datasets under both covariate shift and concept shift. We have the following observations:
7

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
Table 2. The performance on OOD benchmark datasets under Concept Shift, using GCN as the backhone. We report the average accuracy (except for Twitch, where we use ROC-AUC as the evaluation metric) and standard deviations over 5 runs. The best results are shown in bold, and the second best results are shown in underline. OOM denotes out of memory.
Dataset
CBAS
WebKB
Twitch
Cora
Arxiv
Require domain information
Domain
color
university
language
word
degree
time
degree
Base
ERM
82.43=2.56
27.16:+0.93
51.59=3.63
64.03=0028
60.30:2046
65.6410.27
54.81
No
IRM
82.00=1888
26.06:10.73
49.78:3.27
63.93:0.35
60.26:10.49
65.5410.34
56.72:10.30
Yes
Invariant Learning
VREx
22.86:2242 26.61:11.42
5..75:11337
64.03:10.29
60.53-10.39
65.92:10.14
56.68:10.33
Yes
Coral
81.57=1.59 28.07:1:2.64
51.80:3.32
64.04:1.31
60.30:10:45
65.79:10.50
55.141:0.24
Yes
Domain Generalization
DANN
1.557=1.52
29.36++33.1
51.67++3.50
63,9 03.96
60. 60.23+00551
65,67=00.22
55.34++0.5
No
SRGNN
22.14-22.12
26,42-+1.78
51.58+3.64
63.96+0.334
60.27+0.35
65 65.644+0.34
55.08++0,25
Yes
EERM
65.71=0.90
29.91=00.50
0OM
60,21 63.42:20.35
00M
0OM
No
Graph OOD
FLOOD
84.29 1.82
2..62.16.29
5.992-13.92
64.01:10.26
60.31:10.44
66.66
1..59
No
CTT
83.71+0.60
28.9 28.994-2.11
00M
63.77+0.36
60.05+0.599
0OM
00M
No
KL-DRO
81.14:1..48
19754
11
64.03-1.322
60.52:10.22
65.51
54.70::0.35
No
Sample Reweighting
GroupDRO
82.71=1.78
19.1741776
52.2444.05
64. 10.10
60.43:0.40
65.93:0.24
56.24:10.4
Yes
Ours
TAR
87.29=0.78
30.83++1.90
57.20=3.97
64.73=0.23
61.73=:0.22
66.08:10.23
1..26=#1.0
No
Table 3. Results ablation atudy
Dataset
CBAS
Twitch
Shift
concept
covariate
concept
covariate
TAR
87.29:10.78
87.43-1.009
57.20
57.82-12.13
w/o SR
83,57:11.52
84.2911.55
49.78-3.27
55.14111.78
w/o GF
83.
79.43+1.63
3..39
54.54+4.77
majer nodecing weight-8.99
ERM
82,43:122.56
78.29:1.10
3.659
50.04:12.53
sodeiongeeihht unlabeled unde
4.3. Hyper-Paranen Analysis
Figure 4. Learned node weight of TAR on CBAS under concept shift. The size of each node proportional to its sample weight. TAR assigned greater node weight to these minor nodes.
We analyze the impact of the key parameters proposed in Algorithm 1: Tin its value determines the number of iter- ations over which the sample weights propagate along the graph edges. A larger Ti allows weights to transfer over a broader range, whereas a smaller Tin limits weight propaga- tion to a few neighboring hops. As illustrated in Figure 3, performance reaches the best at Tim 1 for the WebKB dataset and at Th 3 for the CBAS dataset. However, for the Twitch and Axiv datasets, performance continues to improve until Ti = 20. We attribute this to the larger size of these two datasets. Larger graphs may require more iterations to effectively propagate sample weights.
4.5. Visualization
We conducted visualization experiments on the CBAS dataset under concept shift to provide an intuitive expla- nation for the effectiveness of our method. In this dataset, there is a spurious correlation between most nodes (major nodes) and their labels, while a small number of nodes (mi- nor nodes) do not exhibit this spurious correlation. If the model primarily learns these spurious correlations from the major nodes during training, its performance will degrade when tested on a distribution that does not contain such correlations. As shown in Figure 4, TAR assigns higher training weights to these minority nodes that do not exhibit the spurious correlation, encouraging the model to focus more on these nodes and avoid learning the spurious cor- relation. As a result, TAR achieves 4.86% improvement over ERM.
4.4. Ablation Study
We conducted an ablation study to investigate the effects of different modules in TAR, including the sample reweight- ing procedure (SR) and Graph Extrapolation module (GE). As shown in Table 3, we compared the TAR with its two variants "w/o SR" and "w/o GE" based on whether the SR and GE were enabled. Both modules contribute to the final result. Specifically, on the Twitch dataset, when SR is not used and only GE is applied, the performance is even lower than that of ERM. This highlights the importance of TAR's adjustment of sample weights in enhancing robustness.
5. Conclusion
Through this work, we innovatively propose the Topology- Aware Dynamic Reweighting (TAR) framework to address
8

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
the distribution shift problem in node classification tasks. TAR utilizes minimax approach to enhance the general- ization ability of GNN models, incorporating topological structure information through gradient flows along graph edges. We further conduct theoretical analysis to reveal the ability of TAR to enhance the distributional robustness of GNNs. Experimental results confirm effectiveness of TAR node classification under distribution shifts. Our TAR opens new direction for addressing the distribution shift problem for node classification tasks.
Duchi, J. C. and Namkoong, H. Learning models with uniform performance via distributionally robust opti- mization. CoRR, abs/ 2018. URL http:
(la.xi..ogg/abs//BB/l.... 08750.
Fan, S., Wang, X., Shi, C., Kuang, K., Liu, N., and Wang, B. Debiased graph neural networks with agnostic label selection bias. IEEE Transactions on Neural Networks and Learning Systems, 2022.
Fey, M. and Lenssen, J. E. Fast graph representation learning with pytorch geometric. arXiv preprint arKiv: (903:02200 2019.
Impact Statement
This paper presents work whose goal to advance the field of Graph Out of Distribution Generalization. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Ganin, Y., Ustinova, E.,, Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., March, M., and Lempitsky, V. Domain- adversarial training of neural networks. Journal of ma- chine learning research, 1 17(59):1-33 2016.
Acknowledgements
Gui, S., Li, X., Wang, L., and Ji, S. Good: A graph out-of- distribution benchmark. Advances in Neural Information Processing Systems, 35:2059- 2073, 2022.
This work is supported by grant 62471133 National Natural Science Foundation of China. This work is also supported by grant 2242024K40020 Central University Basic Research Fund of China, and the Big Data Computing Center of Southeast University. Peng Cui is supported by NSFC (No. 62425206, 62141607) and Beijing Municipal Science and Technology Project (No.
Guo, K., Wen, H., Jin, W., Guo, Y., Tang. J., and Chang, Y. Investigating
generalization of gnns: An architecture perspective. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 932 022-943, 2024.
He, Y., Cui, P., Ma, J., Zou, H., Wang, X., Yang, H., and Yu, P. Learning stable graphs from multiple environ- ments with selection bias. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2194 2202, 2020.
References
Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv: 1907002877, 2019.
Blanchet, J. H., Kang, Y. and M., K. R. Robust wasserstein profile inference and applications to machine learning. J. Appl. Probab., 56(3):830-857, 2019, doi: 10.1017/JPR.20194. URL nattps://dee htttpp://ddi.orgg/// org/ 1017/jpr. 2019.
Hsu, C.-J. Numerical heat transfer and fluid flow. Nuclear Science and Engineering. 78(2):19--197, 1981 doi: 10. 10.13182/NSE81-220112 URL httttt:///doi.org
Hu, W., Fey, M., Zimik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., and Leskovec, J. Open graph benchmark: Datasets for machine learning graphs. Advances in neural information processing systems, 33:22118 28-32333, 2020.
Bojchevski, A. and Günnemann, S. Deep gaussian em- bedding of graphs: Unsupervised inductive learning via ranking, arXiv preprint arXiv:1707.03812 2017.
Chen, R. and Paschalidis, I. C. A robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research, 19 (13):1-48) 2018
Hu, Z. and Hong, L. J. Kullback-leibler divergence con- strained distributionally robust optimization. Available at Optimization Online, 1(2):9, 2013.
Chow, S.-N., Li, W., and Zhou, H. Entropy dissipation of fokker-planck equations on graphs. arXiv preprint 1701041701.04841, 2017.
Jiang, Y., Huang, C., and Huang. L. Adaptive graph con- trastive learning for recommendation. In Proceedings of the 29th ACM SIGKDD conference on knowledge discov- ery and data mining, pp. 4252-4261, 2023.
Deng, C., Yue, Z., and Zhang, Z. Polynormer: Polynomial- expressive graph transformer in linear time. arXiv preprint arXiv:2403.0/2 2024.
Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. In ICLR, 2017.

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Bi- nas, J., Zhang, D., Le Priol, R., and Courville, A. Out- of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pp. 5815-5826. PMLR, 2021.
Rozemberczki, B. and Sarkar, R. Characteristic functions graphs: Birds of feather, from statistical descriptors to parametric models. In Proceedings of the 29th ACM international conference on information & knowledge management, pp. 1325-1334, 2020.
Li, H., Zhang. Z., Wang. X., and Zhu, W. Learning invariant graph representations for out-of-distribution generaliza- tion. Advances in Neural Information Processing Systems, 35:1 15:11828-11 11841, 2022.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case gener- alization. arXiv preprint arXiv: orxiv:/91 1.007731,, 2019.
Liu, J., Hu, Z., Cui, P., Li, B., and Shen, 7. Heteroge- ncous risk minimization. In International Conference on Machine Learning (ICML), pp. 6804-6814. PMLR, 2021a.
Sinha, A., Namkoong, H., Volpi, R., and Duchi, J. Cer- tifying some distributional robustness with principled adversarial training. arXiv preprint 170vv177001105// 2017.
Liu, J., Shen, Z., He. Y., Zhang, X., Xu, R., Yu, H., and Cui, P. Towards out-of-distribution generalization: A survey. arXiv preprini arXiv:2108. 100008136224 13624, 2021
Sui, Y. Wu, Q. Wu, Cui, Q.. Li, L., ZHOU, J., Wang, X., and He, X. Unleashing the power of graph data augmentation covariate distribution shift. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Liu, J., Wu, J., Li, B., and Cui, P. Distributionally robust op- timization with data geometry. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Sun, B. and Saenko, K. Deep coral: Correlation alignment for deep domain adaptation. In Computer Vision-ECCV 2016 Workshops: Amstendam, The Netherlands, October 8-10 and 15-16, 2016, Proceedings, Part III 14, pp. 443- 450. Springer, 2016.
Liu, J., Wang. T., Cui, P., and Namkoong. On the need for a language describing distribution shifts: Illustrations tabular datasets. Advances in Neural Information Processing Systems, 36, 2024.
Tang. J., Li, J., Gao, Z., and Li, J. Rethinking graph neural networks for anomaly detection. In International Con- ference Machine Learning, pp. 21076-21089. PMLR, 2022.
Liu, Y., Ao, X., Feng, F., Ma, Y., Li, K., Chua, T.-.., and He, Q. Flood: A flexible invariant learning framework for out- of-distribution generalization on graphs. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1548-15552 2023.
Wang, B., Chen, J., Li, C., Zhou, S., Shi, Q. Gao, Y., Feng, Y., Chen, C., and Wang, C. Distributionally robust graph- based recommendation system. In Proceedings of the ACM on Web Conference 2024, pp. 3777-3788, 2024.
Wu, Q., Zhang, H., Yan, J., and Wipf, D. Handling distri- bution shifts on graphs: An invariance perspective. In International Conference on Learning Representations, 2021.
Paszke, A., Gross, S., Massa, F,, Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Wu. Y.-.., Wang. X., Zhang, A., He, X., and Chua, T.- S. Discovering invariant rationales for graph neural net- works. arXiv preprint arXiv:2201.12872 2022.
Pei, H., Wei, B., Chang, K. C.-C., Lei, Y., and Yang, B. Geom-gen: Geometric graph convolutional networks. arXiv preprint ar arXiv:2002.05287, 2020.
Xia, D., Wang, X., Liu, N., and Shi, C. Learning invariant representations of graph neural networks via cluster gen- eralization. Advances in Neural Information Processing Systems, 36, 2024.
Qiao, L., Zhang, L., Chen, S., and Shen, D. Data-driven graph construction and graph learning: A review. Neuro- computing, 312:336-351, 2018.
Rong, Y., Huang, W., Xu, T., and Huang, J. Dropedge: Towards deep graph convolutional networks on node clas- sification. In International Conference on Learning Rep- resentations.
Yang, N., Zeng K., Wu, Q. Jia, X., and Yan, J. Learning substructure invariance for -ustrf-distriuution molecular representations. Advances in Neural Information Process- ing Systems, 35:12964-122 12978, 2022.
10

---

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
Ying. Z., Bourgeois, D., You, J., Zitnik, M., and Leskovec, Gnnexplainer: Generating explanations for graph neural networks. Advances in neural information processing systems, 32, 2019.
You, Y., Chen, T., Sui, Y., Chen, T., Wang, Z., and Shen, Y. Graph contrastive learning with augmentations. Advances in neural information processing systems, 33:5812-5823,, 2020.
Zhou, Y., Kutyniok, G., and Ribeiro, B. Ood link predic- tion generalization capabilities of message-passing gnns in larger test graphs. Advances in Neural Information Processing Systems, 35:20257-20272. 2022.
Zhou, Z.. Zhou, S., Mao, B., Zhou, X., Chen, J., Tan, Q., Zha, D., Wang, C., Feng, Y., and Chen, C. Opengsl: A comprehensive benchmark for graph structure learning. arXiv preprint arXiv:2306. / 10280, 2023.
Zhu, Q. Ponomareva, N., Han, J., and Perozzi, B. Shift- robust gnns: Overcoming the limitations of localized graph training data. Advances in Neural Information Processing Systems, 34:27965-27977, 2021.

---

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
A. Related Works
Out of Distribution Generalization. Out-of-Distriuttion (OOD) generalization aims to address the challenge of ensuring model robustness and generalization when faced with data that differ from the training distribution. Numerous studies have been dedicated addressing the problem of OOD generalization, leading to the development of various methods for tackling OOD issues(Liu et al., 2021b). By accurately identifying the causal relationships between features and their corresponding labels, causal learning methods are expected to perform well even when the data distribution changes, the underlying causal structure is often assumed to remain invariant across different environments or domains. Shifting the focus from strict causality to invariance, invariant learning aims to develop a representation or model that remains consistent across various environments. Invariant Risk Minimization (ArMMAArjyvsky et al., 2019) and Variance Risk Extrapolation (VREx) (Krueger et al., 2021) are two prominent methods specifically designed to address these challenges. IRM focuses on learning invariant features by ensuring that the optimal classifier remains the same across different environments, whereas VREx aims to minimize the variance of risks across environments, ensuring stable performance under distributional shifts. Another line of research focused addressing OOD generalization problems involves distributionally robust optimization methods. These model-agnostic techniques come with strong theoretical guarantees and achieve OOD generalization by incorporating distributional robustness into the training process. This ensures that the model's performance remains stable across different data distributions. KL-DRO (Duchi & Namkoong, 2018) minimizes the KL divergence between training and potential test distributions. WDRO (Chen & Paschalidis, 2018; Sinha et al., 2017) leverages the Wasserstein distance to ensure robustness to distributional changes. Group DRO (Sagawa et al., 2019) aims to provide consistent performance across different subgroups by minimizing the worst-case risk among them. While invariant learning has been extensively applied in graph tasks(Wu et al., 2021; Xia et al., 2024; Li et al., 2022; Sui et al., 2023; Wu et al., 2022), there is relatively less application of distributionally robust optimization methods in graph tasks. Applying these methods to graphs requires addressing the unique structural properties of graphs. posing challenges that are specific to graph data.
Graph Invariant Learning. Recently, graph invariant learning has shown enormous success in addressing graph out-of- distribution problems(Wu et al., 2022; 2021; Sui et al., 2023; Zhou et al., 2022; Li et al., 2022; Xia et al., 2024). Graph invariant learning aims to exploit the invariant relationships between graph features(which can be divided into topological structures and node features) and labels across distribution shifts, while filtering out the variant spurious correlations caused by the environment. Recently, many methods have been proposed for graph-level tasks. GIL (Li et al., 2022) captures the invariant relationships between predictive graph structural information and labels in mixture of latent environments. DIR (Wu et al., 2022) selects subset of causal rationales and conducts data augmentation to create multiple distributions to improve generalization. MoleOOD (Yang et al., 2022) enhances the robustness of molecule learning and infers the environment in a fully data-driven manner. AIA (Sui et al., 2023) generates new environments while preserving the original stable features during the augmentation process with adversarial strategies. Compared to research on graph-level ood, less attention has been paid learning node-level representations under distribution shifts from the invariant learning perspective (Wu et al., 2021; Liu et al., 2023; Xia et al., 2024). EERM (Wu et al., 2021) leverages multiple context explorers that are adversarially trained to maximize the variance of risks from multiple virtual environments to learn node invariant predictor. Instead, FLOOD (Liu et al., 2023) applies data augmentation to construct multiple environments and maximize the variance of risks from multiple virtual environments to learn a node invariant predictor. CIT (Xia et al., 2024) generates nodes across different clusters, significantly enhances the diversity of the nodes, and helps GNNs learn the invariant representations. However, this line of invariant learning typically focuses on specific types of invariance (e.g., subgraph invariance in graphs), which may not cover all possible shifts. Besides, due to the lack of environmental information in the real world, generating new samples might introduce bias noise.
B. Datasets
We define the distribution shift problem in the node classification task following (Gui et al., 2022) below. The joint data distribution can be decomposed as IP((Y, X) = P(Y)X)I((X) The main causes of distribution shifts can be separated into two types of shifts:
Covariate shift (Y = P.(Y (X) + (x)))): This indicates that the feature distribution differs between the source distribution and the target distribution.
Concept shift (P, (Y A PP = P.((X)) This indicates that there are spurious statistical correlations in the source distribution that may not hold in the target distribution.
12

---

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
Table 4. Statistics of GOOD datasets
Dataset
#Node
#Edge
#Class
#Feat
Domain
CBAS
700
3962
4
4
Color
WebKB
617
1138
5
1703
University
Twitch
34120
892346
2
128
Language
Cora
19793
126842
70
8710
Word/Degree
Arxiv
169343
1166243
40
128
Time/Degree
In this paper, we use five OOD node classification datasets from GOOD benchmark (Gui et al., 2022), including WebKB, CBAS, Twitch, Cora, and Arxiv. The statistics of these datasets are shown in Table 4.
CBAS a synthetic dataset where the input graph is constructed by attaching 80 house-like motifs to a 300-node Barabási-Albert base graph. The task involves predicting the role each node, including whether node is the top, middle, bottom of house-like motif, or belongs to the base graph, resulting in 4-class classification task. In CBAS, different node colors are used as features, so OOD algorithms must address both covariate shifts caused by node color differences and concept shifts driven by color-label correlations.
WebKB is a university webpage network dataset. Each node corresponds to a webpage, with the words on the page serving as node features, and edges representing hyperlinks between webpages. The task is a 5-class classification problem, where the goal is to predict the category of each webpage. The dataset split based on the domain of the university, meaning the classification based on the content of the webpages and their link structures, rather than on university-specifii attributes.
Twitch is gamer network dataset where nodes represent gamers, with their games serving as node features, and edges representing friendships between gamers. The task a binary classification problem aimed at predicting whether a user streams mature content. The dataset is split based on user language, which implies that the prediction should not be influenced by the language the user uses.
Cora is citation network dataset derived from the full Cora dataset. It consists of a small-scale citation graph where nodes represent scientific publications and edges represent citation links. The task is a 70-class classification of publication types. The dataset is split based on two domain criteria: word and degree. The first domain, word diversity, is defined by the number of selected words in a publication, which is completely unrelated the label. The second domain is node degree, implying that the popularity of a paper (as indicated by its citation count) should not influence its classification.
Arxiv is a citation network adapted from OGB (Hu et al., 2020). The input is a directed graph representing the citation network among arviv papers in the field of computer science (CS). Nodes correspond to arXiv papers, and directed edges represent citation links between them. The task to predict the subject area of these CS papers, making it domain criteria: time (publication year) and node
40-class classification problem. The dataset is split based degree.
C. Implementation Details
In our experiments, we utilized GCN and a recently proposed sota graph transformer Polynormer (Deng et al., 2024) as the backbone models.
For GCN, we configured the models with 3 layers, hidden dimension of 300, dropout rate of 0.5, and a learning rate of 0.01.
For Polynormer, the layer of the local module is set 5 and the global module is set 1, with hidden dimension of 512 and learning rate of 0.001.
Throughout all experiments, we employed the Adam optimizer with a weight decay O. For all baselines, we conducted a grid search as defined by the GOOD Benchmark or their original papers and reported their best results. Note that the graph
13

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts Graph
OOD algorithm EERM (Wu al., 2021) and CIT (Xia al., 2024) encounters CUDA of memory on Arxiv and Twitch datasets due to its high memory requirement. The searching spaces for all the papeee-arsmeters of TAR are as follows.
Entropy term 3: (1, 0.1, 0.01, 0.001).
TAR inner learning rate y:: {0.1, 0.01, 0.001}.
Gradient flow iterations Tin: {1, 3, 5, 10, 20).
Graph extrapolation ratio: {0.0, 0.2, 0.4)
Software and Hardware. Our implementation is under the architecture of PyTorch (Paszke et al., 2019) and PyG (Fey & Lenssen, 2019). All of our experiments are run on one GeForce RTX 3090 with 24GB. The detailed versions of some key packages are listed below:
python: 3.8
pytorch: 1.13.1
pyg: 2.3.1
D. Comparion on SOTA Graph Transformer
Apart from the result using GCNs as the backbone in Table and 2, we also conducted comparative experiments using Polynormer (Deng et al., 2024) as the backbone on the WebKB, CBAS and Twitch datasets.
As shown in the Tabel 5, we have the following observations: Compared to Tables and 2, Polynormer with ERM outperforms GCN in 4 out of settings, demonstrating that Polynormer, recently proposed state-of-the-art Graph Transformer, exhibits superior generalization capabilities compared to GCN. Our proposed method, TAR, achieves better performance than all other baselines across all settings, particularly under concept shift the Twitch dataset, where it outperforms the second- method by significant margin of 7.4%. Experiments conducted with both GCN and Polynormer as backbones further validate the versatility our approach and its superiority over competing methods.
Table performance on three OOD benchmark datasets, using Polynormer (Deng et al., 2024) as the backbone. We report the average test accuracy (except for Twitch, where we use ROC-AUC the evaluation metric) and standard deviations over 5 runs. The best results are shown in bold, and the second best results are shown underline. OOM denotes out memory.
Dataset
CBAS
WebKB
Twitch
Require domain information
Shift
concept
covariate
concept
covariate
concept
covariate
Base
ERM
88.14+2.35
98.29+ 1.20
42.57+33.59
14.60++4.48
51.02++..23
14.06+1.41
No
IRM
87.86:13.23 97.1412.12
42.57:15.25
13.81#1444
51.07:1..5
54.08:11.33
Yes
Invariant Learning
VREx
89.43.L1.37 99.14.11.20
44.04
15.71.12.11
52.80.10.85
$4.30.11.10
Yes
Coral
87.7111885 97.71111.3
42.75.15.37
13.49.10.56
51.3311119
53.45:11.29
Yes
Domain Generalization
DANN
87.86+2.7
96.57+2.96
1..9
14.29+3.81
50.91+1.33
54.27+1..37
No
SRGNN
87.43-13.06
98.00-22.96
11.10:11.29
14.60:12.95
51.00::..24
54.23++1.30
Yes
EERM
85.71=2.81
97.14+2..67
43.49:13.92
25.71
OOM
OOM
No
Graph OOD
FLOOD
88.57-12.08
98. 08.57-11.01
41.28.11.83
14.92.14.14
51.75.1..0
54.50.1..3
No
CIT
89.57=10.81
1.22
43.12:1..54
14.76.13.30
00M
0OM
No
KL-DRO
1.30
08.86:t:1
43.59
29.37 14.64
55.06
14.46 1.42
No
Sample Reweighting
GroupDRO
1..71
99.14++1.92
43.49-13.35
14.29:11.59
53.62443.52
14.21+1.26
Yes
Ours
TAR
91.14:11.20 1.20
99.43++1..8
47.89-+1.76
30.48-+7.88
63.25++1.84
56.69-+0.88
No
14

---

Topology-Aware Dynamic Reweighting for Distribution Shifts on Graph
E. Comparative Analysis of Solutions for Disconnected Trained Nodes
During the training of node classification tasks, the labeled nodes are randomly sampled, which often results in these labeled nodes not being directly connected in the space. Consequently, the gradient defined in Equation cannot be computed. In this paper, we sets the loss of the unlabeled nodes the mean loss of the labeled nodes, we also try other solutions, as detailed below:
TAR. This method sets the loss of the unlabeled nodes to the mean loss of the labeled nodes. This approach preserves the original graph structure, allowing the gradient flow between training nodes to occur indirectly through the unlabeled nodes.
TAR-N. This method performs a breadth- first search for each labeled node during training to find K reachable labeled nodes. Then, it adds edges between the current node and these K nodes, creating a new graph, and performs the gradient flow on this new graph. This approach enables the gradient to propagate directly among the training nodes but disrupts the original graph structure.
We conduct experiments both CBAS and Cora (degree) to compare their performance. The results summarized in Table 6. TAR-N performed worse than TAR under all settings. This suggests that preserving the original graph structure might be more reliable. Additionally, TAR-N requires extra computational overhead for the breadth-first search compared to TAR.
Table The performance TAR with orginal graph or TAR with reconnccted graph on graph OOD datasets.
Dataset
CBAS
Cora
Shift
concept
covariate
concept
covariate
ERM
82.43:2.56
78.2913.10
60.30:0.46
55.33:10990
TAR
87.29
87.43
61.73.10.22
56.62.10.23
TAR-N
86.71+2.229
2..99
61.57++022
56. 07.40440.4
F. Complexity Analysis
Our approach is model-agnostic, making applicable to any graph neural network. Moreover, it does not involve any learnalle parameters, ensuring high efficiency.
Below, we analyze its time complexity: For an l-layer, d-dimensional GCN on a graph with n nodes and m edges, the time cost of TAR includes: Outer Training: O(Idn 4m), which encompasses feature transformation and neighbor aggregation. Inner Loop: The & iterations have a time complexity of Okk + km), introducing minimal additional computation since < ldf, The inner loop of TAR operates a weessage-passing style and is parameter-fie ffree, ensuring scalability comparable to GCNs for large-scale graphs.
We compared the time complexity of our method with other model-agnostic approaches, shown in Table 7. It can be observed that our method, similar to Group DRO and VREx, introduces only little overhead, demonstrating excellent scalability. In comparison to CIT, our approach achieves faster training speeds and avoids out-of-memory (OOM) issues, making it more efficient and practical for large-scale applications.
Table 7. Time cost of different OOD methods (ms/epoch).
Methods
ERM
GroupDRO
VREx
CIT
TAR (ours)
CBAS
10.1 181.8+1.0
10.3+1.2 232.9+1.4
10.8
15.0.4.33
12.0:2.2
Arxiv
231.8+0.8
0OM
235.1
15

---

Topelogy-Aware Dynamic Reweighting for Distribution Shifts on Graph
G. Proof
G.I. Proof of Theorem 1
Proof. Denote qº arg sup g(0,g) 2r gW (p. g). and we have e = GW Go (p. )*).
qE (E())
Then we prove by contradiction: we first assume
q arg
sup [(0.g) g)
g: (c.
which indicates that g().9') 2 c(), and GMG, (p. S. Therefore, we have (p.4) S Denote
R(0,4) = ((0,0) 2 gWW (p.4)
1
then have ((0, R(0,g') This leads contradiction since 4² is the supremum point of R(0,)) -).
G.2. Proof of Theorem 2
Proof. Based on the (Chow et al., 2017, Theorem 5), we have
[(q(oo)) - -(g()) S (I([(((o))) - L(q(0))), -(((())
where C' > is a constant depending on the loss function E, hyper-parameter B, and the sample size N.
Then we denote the real worst-case distribution within the e-radius discrete Geometric Wasserstein-ball as q", and we have
((q(oo)) + ((p") -(())) s (*(((g(o)) [(g)) =(g") -(()))
Therefore, we have
((g²) - C(a(t)) S -CC((((g)) - C(q(0))) -1 (1 -(("))
and
C(a²) - C(y(1)) -c- (1 -e -01) [(y((o)) Ligg) <c-
C(qx)- - (()()))
L(qr) - L(g(0))
H. Limitations and Future Work
In this paper, we propose a model-agnostic node classification OOD generalization algorithm TAR that leverages the existing topological information in node classification tasks to achieve better local robustness. However, as discussed in Appendix E, due to the random sampling labeled nodes in node classification tasks, the nodes with computed losses not directly connected, with other nodes without computed losses in between. Although explored two potential solutions in Appendix E, we consider these solutions suboptimal. In the future, we aim to design more effective method to address this issue
16


Reviews:
[{"review_id": "UAotp8Kmrb", "reviewer": "Reviewer_XE4R", "summary": "This paper proposes TAR, a framework which dynamically weights and reweights nodes within a Graph-Neural-Network (GNN) given the “risk”-level of nodes, incorporating topological structural information and providing robustness against shifts of distribution.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insights in the evaluation and the hints for revising manuscripts.\n\n> Graph extrapolation could have been mentioned earlier in the paper.\n\nThank you for your suggestion. We will add a summary in the methodology section to introduce GE earlier.\n\nOnce again, thank you for your careful review. We will thoroughly check our manuscript to avoid any grammar or spelling errors!", "claims_and_evidence": "Claims are backed by theoretical proofs when made.", "methods_and_evaluation_criteria": "Five classification datasets commonly used across OOD problems were used for validation, with a range of graph size and class/feature space size. Distribution shifts were defined. Baseline and state-of the art methods were compared against the newly proposed TAR framework, with TAR maximising the central tendency of test-accuracy without compromising the variance.", "theoretical_claims": "Proofs were concise and used where relevant.", "experimental_designs_or_analyses": "The experimental design is well suited to the problem posed.", "relation_to_broader_scientific_literature": "This paper builds on work in the field of OOD generalisation, in particular providing theoretical proofs alongside empirical validation for distributional robustness.", "essential_references_not_discussed": "I am not aware of any essential missing references.", "other_strengths_and_weaknesses": "This paper is concise where necessary and well-explains the TAR framework without unnecessary focus on preliminaries. In particular the backing of empirical work with theoretical proofs for understanding time complexity and limitations of the approach is commended.", "other_comments_or_suggestions": "Graph extrapolation could have been mentioned earlier in the paper. It is used in Figure 1, but the requirements/comments are not expanded upon until section 3.3.\nPage 5 Line 263 – unclear grammar used. “…which disrupts the connectivity and hinders the calculation of this Equation, we intuitively…”.\nPage 13 Line 665 – Spelling error. “entorpy term”.", "questions_for_authors": "No questions, thank you for the paper", "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": "opsCe5WNWb", "reviewer": "Reviewer_sMqu", "summary": "This paper proposes a Topology-Aware Dynamic Reweighting (TAR) framework to address distribution shifts in node classification tasks using Graph Neural Networks (GNNs). Addressing the limitations of existing invariant learning methods (which rely on strong invariance assumptions) and sample reweighting approaches (which ignore graph topology), TAR dynamically adjusts node weights through gradient flows along graph edges while incorporating topological constraints via discrete geometric Wasserstein distance. The framework employs a minimax optimization strategy: the inner maximization identifies local worst-case distributions through constrained gradient flows, while the outer minimization trains the GNN model on these adversarially weighted samples to enhance distributional robustness. Theoretical analysis proves that TAR achieves exponential convergence rates and provides formal distributional robustness guarantees.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your careful review and insightful feedback. Below we provide point-by-point responses to address all raised concerns.\n\n# Analysis of Graph Extrapolation\n## 1.Interaction between GE and Reweighting\nGraph Extrapolation (GE) is a crucial component of TAR. Without GE, TAR reweighting can only simulate potential worst-case distributions by adjusting sample weights of existing data, which is limited by the coverage of current graph. In contrast, GE expands the potential distribution by transforming the graph. Specifically, while OOD samples may not appear during training, GE may construct such samples, allowing TAR reweighting to assign them higher weights and simulate potential distribution shifts.\n\n## 2.Comparison with Mixup\nWe use both Drop Edge and Mask Feature as GE since they transform the graph's topology and node features, respectively. These techniques are also utilized in other graph OOD methods like FLOOD. For comparison, we replace GE with Mixup, which only interpolates node features. The results show that GE consistently outperforms Mixup.\n\n|dataset|CBAS|Twitch|WebKB|\n|-|-|-|-|\n|Mixup|74.86±1.63|52.23±1.21|31.59±10.53|\n|TAR+mixup|78.00±2.78|54.51±1.30|36.51±8.57|\n|TAR+GE (ours)|**87.43±5.09**|**57.82±2.13**|**37.46±4.84**|\n\n# Performance when there is no shift during testing\nThank you for your suggestion! Below are the comparative results and key observations:\n- On CBAS and Twitch datasets, TAR consistently outperforms baselines in both IID (no shift) and OOD (covariate shift) settings, with significantly larger gains under covariate shift.\n- On Cora dataset, TAR shows a slight, statistically insignificant drop in IID performance but achieves the best results under OOD, surpassing other methods.\n\n|dataset|CBAS|CBAS|Twitch|Twitch|Cora|Cora|\n|-|-|-|-|-|-|-|\n|shift|IID|OOD|IID|OOD|IID|OOD|\n|ERM|86.00±1.20|78.29±3.10|72.83±1.23|50.04±2.53|69.65±0.55|64.72±0.54|\n|VREx|83.14±3.26|78.57±2.02|70.84±4.15|51.26±4.82|69.63±0.31|65.02±0.45|\n|GroupDRO|86.29±1.28|77.14±2.67|71.54±3.34|50.48±2.04|**69.77±0.26**|64.95±0.59|\n|TAR|**90.29±2.93**|**87.43±5.09**|**73.73±0.44**|**57.82±2.13**|69.53±0.77|**65.64±0.37**|\n\n# Sensitivity analysis of other parameters\nThank you for your insightful suggestion!\nWe want to clrify that \"the coefficient of the topology penalty\" is not a hyperparameter, we set it to $\\frac{1}{2\\tau}$ in our optimization (page 3, line 163), and $\\tau$ is exactly the TAR inner learning rate (page 4, equation 5). Due to space constraints, additional analysis figures for other parameters are provided in this anonymous [link](https://anonymous.4open.science/r/ICML-rebuttal-16C3). Please refer to it for more details.\n\n# Training Cost Analysis\nCompared to empirical risk minimization (ERM), TAR introduces minimal overhead. Our paper (Appendix F) includes complexity analysis, and we provide additional details here.\n\nFor an $l$-layer, $d$-dimensional GCN on a graph with $n$ nodes and $m$ edges, the time cost of TAR includes: \n- **Outer Training**: $O(ld^2n + lm)$, which encompasses feature transformation and neighbor aggregation.\n- **Inner Loop**: The $k$ iterations have a time complexity of $O(kn + km)$. \n\nReal-world graphs are typically sparse, with small node degrees and larger GNN dimensions (e.g., $d=300$, max average degree $m/n=26$ in Twitch). Thus, $ld^2n \\gg lm$, and since $k \\ll ld^2$, TAR's additional computation is negligible. Besides, the inner loop operates in a message-passing style and is parameter-free, ensuring scalability for large-scale graphs. \n\nThe table below compares TAR's runtime (epoch/s) with other methods, showing it introduces slight overhead over ERM, demonstrating **excellent scalability**.\n\n|Methods|ERM|KLDRO|GroupDRO|VREx|CIT|TAR(ours)|\n|-|-|-|-|-|-|-|\n|CBAS (4k edges)|10.1±0.2|10.2±0.7|10.3±1.2|10.8±0.2|15.0±4.3|12.0±2.2|\n|Arxiv (1.2M edges)|181.8±1.0|186.9±1.1|232.9±1.4|231.8±0.8|OOM|235.1±3.1|\n\n> Q: Why does this paper additionally consider concept shift?\n\nWe would like to clarify that we follow the terminology used in GOOD benchmark, where distribution shifts are categorized into concept shift and covariate shift, with datasets specifically constructed for these cases. Our experiments strictly follow this benchmark setting, and results show that our method effectively enhances performance across both types of shifts.\n\nThank you for your question, and we will incorporate this into our final version to avoid any misunderstandings.\n\n> Q: The bi-level optimization of min and max may lead to high training costs and instability.\n\nFor TAR's training cost, we refer the reviewer to the \"Training Cost Analysis\" above. Regarding instability concerns, we provide additional loss visualizations at this [link](https://anonymous.4open.science/r/ICML-rebuttal-16C3). These results demonstrate that TAR achieves a smaller generalization loss without significant training instability.\n\nShould you have any additional concerns, please do not hesitate to let us know.", "claims_and_evidence": "The claims in the submission are ​largely supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "The core theoretical claims are ​mathematically consistent within the stated assumptions.", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "It offers a new perspective for graph node classification under distribution shifts.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths：\n1. Combines ​discrete geometric Wasserstein constraints with entropy regularization, ensuring smoothness.\n2. Provides formal guarantees for ​distributional robustness and ​exponential convergence rates.\n3. Demonstrates ​consistent improvements over most baselines across diverse datasets.\n4. Maintains computational efficiency via message-passing mechanisms, with minimal overhead compared to ERM training.\n\nWeaknesses:\n1. Limited Analysis of Graph Extrapolation (GE): (1) Compare GE to other augmentation strategies (2) Analyze how GE interacts with TAR’s reweighting mechanism.\n2. While TAR is efficient for moderate-sized graphs, the ​gradient flow iterations may scale poorly for graphs with billions of edges.\n3. The bi-level optimization of min and max may lead to high training costs and instability.", "other_comments_or_suggestions": "1. Please add the analysis of graph extrapolation.\n2. It is suggested that the authors supplement the performance of the method in this paper when there is no covariance shift or concept shift during testing.\n3. It is suggested that the authors provide sensitivity analyses for the following parameters: the coefficient of the entropy term, the TAR inner learning rate, the graph extrapolation ratio, and the coefficient of the topology penalty.", "questions_for_authors": "1. How is the sensitivity of other parameters？ \n2. Why does this paper additionally consider concept shift?", "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": "e7XUwSypLw", "reviewer": "Reviewer_RR32", "summary": "This paper introduces Topology-Aware Dynamic Reweighting (TAR), a novel framework designed to enhance node classification performance under distribution shifts by leveraging graph topology. Unlike invariant learning approaches that rely on strict assumptions about environment labels, TAR applies a dynamic sample reweighting strategy that incorporates gradient flow along graph edges to adjust sample importance during training. The framework employs a minimax optimization process: the inner maximization learns probability densities for samples, while the outer minimization trains the model on a dynamically weighted distribution. Theoretical analysis demonstrates that TAR identifies the local worst-case distribution, improving robustness against unseen distributions. Experimental evaluations on multiple OOD node classification benchmarks show that TAR consistently outperforms existing methods, particularly in settings where domain information is unavailable.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your careful review and insightful feedback. Below we provide point-by-point responses to address all raised concerns.\n\n# Performance when there is no shift during testing\nThank you for your insightful suggestions!\n\nWe compare TAR with several typical robust optimization methods, and the results are presented in the table below. Our key observations are as follows:\n\n- On the CBAS and Twitch datasets, TAR consistently outperforms baselines in both IID (no shift) and OOD (covariate shift) settings. Notably, the performance improvements under OOD are significantly more pronounced than those under IID.\n- Some methods, such as VREx, while achieving notable gains on CBAS and Twitch under the OOD setting, exhibit substantial performance degradation in the IID setting.\n- On the Cora dataset, TAR shows a slight, statistically insignificant drop in performance under IID. However, it achieves the best performance under the OOD setting, surpassing all other methods. \n\nThese results highlight TAR's robustness and effectiveness across different scenarios.\n\n|dataset|CBAS|CBAS|Twitch|Twitch|Cora|Cora|\n|-|-|-|-|-|-|-|\n|shift|IID|OOD|IID|OOD|IID|OOD|\n|ERM|86.00±1.20|78.29±3.10|72.83±1.23|50.04±2.53|69.65±0.55|64.72±0.54|\n|VREx|83.14±3.26|78.57±2.02|70.84±4.15|51.26±4.82|69.63±0.31|65.02±0.45|\n|KLDRO|86.00±1.86|77.71±1.63|73.35±0.63|53.76±4.19|69.64±0.62|64.85±0.51|\n|GroupDRO|86.29±1.28|77.14±2.67|71.54±3.34|50.48±2.04|**69.77±0.26**|64.95±0.59|\n|TAR|**90.29±2.93**|**87.43±5.09**|**73.73±0.44**|**57.82±2.13**|69.53±0.77|**65.64±0.37**|\n\n# Performance of different graph sparsity levels\nThank you for your question! As shown in the table below, the five benchmark datasets we employed exhibit diverse degree distributions, with average node degrees ranging from 1.52 to 26.15. Our method demonstrates consistently strong performance across all these varying degree distributions as you can see in Tables 1 and 2 in our paper.\n|dataset|webkb|cbas|twitch|cora|arxiv|\n|-|-|-|-|-|-|\n|avg. degree|1.52|5.66|26.15|6.41|13.67|\n\n# Performance on heterophilic graphs\n|dataset|webkb|cbas|twitch|cora|arxiv|\n|-|-|-|-|-|-|\n|avg. node homophilic|**0.14**|0.60|0.64|0.59|0.64|\n\nThank you for your insightful question! As shown in the table above, our evaluation covers datasets with varying levels of homophily. Notably, WebKB represents a standard heterophilic graph benchmark.\nWe additionally contruct an another OOD dataset chameleon, following the same priciple of GOOD benchmark [1] and using node degree as split type. As demonstrated in the results table below, TAR consistently achieves top performance on both heterophilic graphs, further validating its robustness across different homophily regimes.\n\n|dataset|WebKB|WebKB|Chameleon|Chameleon|\n|-|-|-|-|-|\n|shift|concept|covariate|concept|covariate|\n|ERM|27.16±0.93|19.37±6.01|34.65±2.56|33.13±4.53|\n|VRex| 26.61±1.42|36.83±1.99|34.81±2.51|**36.88±3.43**|\n|CIT|28.99±2.11|28.89±9.09|36.76±3.43|35.11±2.35|\n|GroupDRO|29.17±1.76|25.24±9.57|34.97±0.95|36.15±5.19|\n|TAR|**30.83±1.90**|**37.46±4.84**|**38.34±1.48**|**36.88±2.66**|\n\n# Performance on contrastive loss\nThank you for your insightful question! TAR can be applied to different loss type. We apply TAR to two graph contrastive learning methods——GRACE [2] and COSTA [3]——to verify its effectiveness on contrastive loss. Specifically, we apply TAR in the pretrain stage and keep the finetune stage the same as the raw contrastive learning method. As shown in the table below, TAR consistenly improves the performance of different graph contrastive learning under distribution shifts.\n|dataset|WebKB|WebKB|CBAS|CBAS|\n|-|-|-|-|-|\n|shift|concept|covariate|concept|covariate|\n|ERM|27.16±0.93|19.37±6.01|82.43±2.56|78.29±3.10|\n|GRACE|32.66±5.79|19.68±9.73|90.29±1.64|93.71±1.28|\n|+ TAR|36.33±5.79|26.51±3.00|91.00±0.64|94.57±2.56|\n|COSTA|32.11±4.00|16.03±5.09|82.71±1.63|76.29±3.99|\n|+ TAR|33.39±3.28|19.68±5.68|85.29±2.45|79.14±3.73||\n\n[1] S. Gui, et al., Good: Agraph out-of-distribution benchmark. Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track.\n\n[2] Y. Zhu, et al. Deep Graph Contrastive Representation Learning. ICML Workshop on Graph Representation Learning and Beyond, 2020.\n\n[3] Y. Zhang, et al. Costa: Covariance-preserving feature augmentation for graph contrastive learning. SIGKDD 2022.\n\n# Training Cost Analysis\nWe compared the time cost (epoch/s) of our method with other methods, as shown in the table below. It can be observed that our method, similar to Group DRO and KLDRO, introduces only little overhead, demonstrating **excellent scalability**. For a comprehensive complexity analysis, please see Appendix F of our paper.\n\n|Methods|ERM|KLDRO|GroupDRO|CIT|TAR(ours)|\n|-|-|-|-|-|-|\n|CBAS (4k edges)|10.1±0.2|10.2±0.7|10.3±1.2|15.0±4.3|12.0±2.2|\n|Arxiv (1.2M edges)|181.8±1.0|186.9±1.1|232.9±1.4|OOM|235.1±3.1|\n\nShould you have any additional concerns, please do not hesitate to let us know.", "claims_and_evidence": "TAR’s theoretical foundation ensures robustness and convergence：While the theory is sound, additional empirical ablations (e.g., impact of different loss functions) could strengthen the claim.", "methods_and_evaluation_criteria": "The evaluation methodology is generally appropriate for the problem.", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "The experiments are well-designed.", "relation_to_broader_scientific_literature": "The paper is well-positioned within the OOD generalization and graph representation learning literature.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths：\n1. Unlike standard sample reweighting methods that treat training instances independently, TAR explicitly considers graph topology during reweighting. The use of gradient flow along graph edges ensures that sample weights are influenced by the structural context of nodes, making the approach more aligned with the underlying data distribution. This is particularly important in graphs where neighboring nodes share similar properties.\n\n2. The inner maximization problem dynamically adjusts sample weights by optimizing against a local worst-case distribution, while the outer minimization problem trains the model using these reweighted samples. This approach is inspired by distributionally robust optimization (DRO) but extends it to a graph-based setting with topology-aware constraints. The theoretical justification for this method strengthens its credibility.\n\n3. The paper provides rigorous theoretical analysis that supports the proposed method. These results ensure that the dynamic reweighting mechanism is theoretically grounded, rather than being an empirical heuristic.\n\nWeaknesses:\n1. While TAR improves robustness under distribution shifts, it is unclear how much performance degradation occurs on in-distribution (ID) data. Many robust optimization methods sacrifice ID accuracy to gain robustness, but the paper does not explicitly analyze this trade-off. A controlled study where ID and OOD accuracy are jointly evaluated would provide deeper insights into this aspect.\n\n2. TAR’s effectiveness heavily relies on gradient flow over graph edges. However, in sparse graphs with low connectivity, the amount of information that can be propagated during sample reweighting may be severely limited. Conversely, in densely connected graphs, the method might smooth weights too aggressively. The paper does not systematically analyze how TAR performs across different graph sparsity levels or degree distributions, which could affect its generalizability.\n\n3. The method optimizes a weighted loss function using cross-entropy loss, but it is unclear how TAR would perform under alternative loss functions (e.g., margin-based losses, contrastive losses). Some OOD methods are highly sensitive to loss function choice, and the paper does not explore whether TAR is robust to such variations.", "other_comments_or_suggestions": "1. How does TAR compare to simpler DRO methods in terms of training time?\n2. How does TAR perform on heterophilic graphs, where homophily assumptions do not hold?", "questions_for_authors": "Regarding the questions, please refer to the weaknesses section in the strengths and weaknesses part.", "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:

Meta-review: The draft proposes a new alternative training proposal for Graph Neural Networks under the specific subproblem of the distribution shift between the train and test sets.
The draft gained acceptance rates by all 3 reviewers (of which at least 1 is an expert on the field and 1 is mostly agnostic). All reviewers at least confirmed that they have read the rebuttal and one of them was explicit on keeping his score.
All three reviews praised the rigorous theoretical support and some way or another also were happy with the clear and consistent evidence regarding the performance. Although on this latter I felt that inferential statistics (beyond merely descriptive statistics) was missing as the shown effect sizes are somewhat small. Among the weaknesses the contribution the reviewers appear to be more driven by "potential" issues rather than actual flaws. Among these potential issues were lack of scalability, instability, training costs and in-distribution degradation of performance among others. To all of these potential issues, the authors provided some degree of reassurance with further quantitative evidence and they did not try to duck any questions which speaks in their favor.
All in all, I can see why the reviewers agree to see this as suitable for acceptance. However given that none seem to have highlighted major actual (rather than potential) flaws, I would have expected perhaps something more solid than two 3 scores.
Decision: Accept (poster)