Counterfactual Supervision for Adaptive Propagation Controls in Graph Neural Networks

Graph neural networks (GNNs) propagate information between neighboring nodes, yet the most effective propagation behavior can differ across nodes and layers. Adaptive GNNs seek to address this variation through learned controls. However, standard end-to-end training supervises these controls through gradients evaluated at their current values, thereby capturing only local loss sensitivity and failing to directly compare the downstream predictions resulting from distinct propagation behaviors. This shortcoming weakens the training signal and complicates the understanding of the learned controls. To address this issue, we introduce a method known as Counterfactual Supervision. This training strategy utilizes node-local continuous controls with two meaningful alternatives. For a selected control, we temporarily set it to each alternative, recompute the subsequent layers, and utilize the difference in prediction loss as an additional supervision signal. This approach complements end-to-end optimization while maintaining continuous controls. Evaluating both alternatives at every node can be computationally costly. Therefore, we show that alternatives at sufficiently distant nodes can be assessed together without compromising the exact loss differences for each node. This strategy enables the derivation of multiple per-node comparisons from a single pair of recomputations. We implement this framework in TriAxis-GNN, which learns distinct controls for nodes and layers related to propagation direction, balancing contributions from self and neighbors, and distinguishing between ordinary and higher-order neighborhood structures. Experiments on homophilic and heterophilic node classification benchmarks indicate that counterfactual supervision enhances adaptive propagation models. Furthermore, the proposed batching strategy considerably reduces computational costs associated with this method. Our source code is available at: https://anonymous.4open.science/r/TriAxis-GNN-7BE3/.