Structural Alignment Improves Graph Test-Time Adaptation

TL;DR

Graph neural networks degrade under distribution shifts that alter neighborhood connectivity. The authors propose Test-Time Structural Alignment (TSA), a GTTA method that adapts a pretrained GNN at inference by (i) neighborhood alignment via a $\boldsymbol{\gamma}$ matrix to correct conditional structure shifts, (ii) SNR-driven weighting to blend self and neighbor representations by layer, and (iii) boundary refinement using non-graph TTA signals. The approach is theoretically grounded, yielding an upper bound on the GTTA error gap that decomposes into label shift, neighborhood shift, and feature shift, and empirically TSA outperforms both non-graph TTA and existing GTTA baselines across diverse datasets and backbones. TSA is lightweight, model-agnostic, and effective under privacy and compute constraints, making practical GTTA feasible without retraining on source data.

Abstract

Graph-based learning excels at capturing interaction patterns in diverse domains like recommendation, fraud detection, and particle physics. However, its performance often degrades under distribution shifts, especially those altering network connectivity. Current methods to address these shifts typically require retraining with the source dataset, which is often infeasible due to computational or privacy limitations. We introduce Test-Time Structural Alignment (TSA), a novel algorithm for Graph Test-Time Adaptation (GTTA) that adapts a pretrained model to align graph structures during inference without the cost of retraining. Grounded in a theoretical understanding of graph data distribution shifts, TSA employs three synergistic strategies: uncertainty-aware neighborhood weighting to accommodate neighbor label distribution shifts, adaptive balancing of self-node and aggregated neighborhood representations based on their signal-to-noise ratio, and decision boundary refinement to correct residual label and feature shifts. Extensive experiments on synthetic and real-world datasets demonstrate TSA's consistent outperformance of both non-graph TTA methods and state-of-the-art GTTA baselines.

Figures (7)

• Figure 1: Example of the distribution shifts of neighborhood information due to different experimental conditions. Pileup mitigation aims to detect LC neutral particles using kNN graphs (dashed circles) that leverage nearby particles for inference. The model is trained on ${\mathcal{G}}^{\mathcal{S}}$ but needs to generalize to ${\mathcal{G}}^{{\mathcal{T}}_1}$ and ${\mathcal{G}}^{{\mathcal{T}}_2}$. The inferred nodes within the circles are the LC neutral particles, but their neighborhood node label ratios change in ${\mathcal{G}}^{{\mathcal{T}}_1}$ and ${\mathcal{G}}^{{\mathcal{T}}_2}$. Both cases represent neighborhood shift, which this work aims to address. A more formal definition of these shifts is discussed in Sec. \ref{['sec:Test Error Analysis']}.
• Figure 2: TSA utilizes neighborhood alignment to address neighborhood shift and SNR-inspired adjustment to mitigate SNR shift. It further adjusts the decision boundary to get refined predictions $\hat{y}$. The refined soft pseudo-labels are used to estimate the parameter $\boldsymbol{\gamma}$ for neighborhood alignment and to optimize $\alpha$ for combining self representations with neighborhood-aggregated representations.
• Figure 3: (a) Comparison of neighborhood alignment with $\boldsymbol{\gamma}$ from model prediction and Oracle on the CSBM graphs deshpande2018contextual. The top (or bottom) subfigures represent the assignment under neighbor shift (or neighbor shift plus label shift, respectively). Nodes are grouped by the entropy of their soft pseudo labels and the y axis shows the accuracy after assigning $\boldsymbol{\gamma}$. Ideally, a correct assignment (red) would lead to near-perfect accuracy. However, the assignment based on pseudo labels is far from optimal (blue). Figure (b) to (d) present t-SNE visualizations of node representations during GTTA. The model is trained on the source domain in (b) (CSBM with label distribution $[0.1, 0.3, 0.6]$) and evaluated on the target domain in (c) (CSBM with label distribution $[0.3, 0.3, 0.3]$). Figure (d) illustrates the qualitative result of the neighborhood alignment. Despite well-aligned neighborhood representations, a mismatched decision boundary (red) leads to class 0 misclassifications. Node colors represent the ground-truth labels. The top subfigures show the output given by the GNN encoder, while the bottom subfigures show the classifier decision boundaries.
• Figure 4: (a) and (b) compare the SNR of neighborhood-aggregated representations to that of self-node representations across node degrees and GNN layers, based on the ratio of the two quantities. A higher ratio indicates that the neighborhood-aggregated representations exhibit higher SNR. (c) Analysis of the SNR adjustment $\alpha$ across different layers and node degrees. Details are discussed in Sec. \ref{['subsec:result_analysis']}.
• Figure 5: The t-SNE visualization of a one-layer GraphSAGE under different distribution shifts: (top) the output representations given by the GNN encoder and (bottom) the node representations with the classifier decision boundaries. (a) indicates the source domain where the model is pretrained. (b), (c), and (d) stand for feature shift, conditional structure shift (CSS), and label shift. (e) represent the impact of SNR shift induced by degree changes. The color of the nodes represents the ground-truth labels. The legends in the bottom plots show the accuracy and recall scores for each class. The red contours in the bottom (d) highlight the decision boundary of the minority class.

Theorems & Definitions (9)

• Definition 3.1
• Definition 3.2
• Theorem 3.3: Error Decomposition Theorem (Informal)
• Definition 4.1
• Definition 4.2
• Lemma A.1
• proof
• Theorem A.1: Error Decomposition Theorem (Informal)
• proof