LabStA: Label- and Structure-Aware Mini-Batch Sampling for Node Classification with Graph Transformers

Scalability is a major challenge in graph representation learning, as training on large-scale graphs may require processing millions of nodes that do not fit within limited GPU memory. For Graph Neural Networks (GNNs), mini-batch training provides a practical alternative to full-graph training, typically using connectivity-oriented sampling methods that preserve local connectivity for message passing. However, mini-batch construction for Graph Transformers (GTs) remains underexplored, and many prior works rely on simple random node sampling, which can lead to suboptimal performance. In this work, we motivate two key objectives for mini-batch sampling in GT-based node classification: reducing class-distribution discrepancy and increasing edge retention rate. Based on these objectives, we propose LabStA, a Label- and Structure-Aware mini-batch sampling method that augments graph partitions with informative virtual nodes to jointly improve class-distribution alignment and edge preservation. Experiments on diverse node classification benchmarks show that LabStA outperforms existing sampling baselines on most dataset–backbone pairs, especially on heterophilous graphs.