Work

Mitigating Imbalance and Distribution Shift in Node Classification via Norm Consistency

Year: 2026

Type: article

Abstract: Graph neural networks (GNNs) for semi-supervised node classification suffer from two pervasive yet seemingly unrelated pathologies: class imbalance, where minority classes are dominated during message passing, and distribution shift, where test nodes follow a different distribution from training nodes. Across the settings studied in this work, we observe a recur ring geometric signature: minority-class and out-of-distribution (OOD) nodes tend to develop smaller ℓ2 representation norms, which can weaken their contribution during aggregation and reduce absolute classification margins. Building on this insight, we propose NodeReg, a lightweight, model-agnostic regularizer that penalizes per-node deviation from a stop-gradient global mean norm via a smooth-ℓ1 loss. NodeReg adds O(n) overhead, requires no synthetic data or environment inference, and plugs into any GNN pipeline. Theoretically, we provide (i) a conditional population-risk bound under the stylized SNM-SBM feature noise model and extend it to a shifted distribution through a second-order loss-moment condition, and (ii) a conditional analysis showing how radial balance is compatible with Fisher discriminative and neural-collapse geometry when intra-class collapse and angular ETF structure are supplied separately. On five public benchmarks, NodeReg lifts the average F1 score by 5.4% under extreme class imbalance (ratio 0.1) and the average accuracy by 1.0% under feature, temporal, and sub graph distribution shifts, consistently outperforming the strongest tailored baselines for either scenario.

Source: IEEE Transactions on Knowledge and Data Engineering

DOI: 10.1109/tkde.2026.3748943