Revisiting Pre-Propagation GNNs: Robust Diffusion Operators and Hidden-State Re-Propagation

Pre-propagation graph neural networks (PP-GNNs) decouple node feature propagation from transformation: graph diffusion is performed once as preprocessing, and training reduces to dense per-node transformations. This design enables mini-batch training without inter-node dependencies, avoids repeated sparse matrix--matrix multiplications, and better matches modern accelerators optimized for dense compute. However, their expressivity remains unclear, and empirical results show a gap between PP-GNNs and their message-passing counterparts on commonly used graph benchmarks, especially heterophilic ones. In this paper, we propose a suite of robust graph diffusion operators for preprocessing and a few-shot hidden-state re-propagation scheme during training. Our methods improve the validation and test accuracy of PP-GNNs, enabling them to match the accuracy of message-passing GNNs while maintaining training efficiency.

Many AI systems need to learn from networks, such as social networks, citation networks, or recommendation graphs, where each item is connected to others. Powerful graph learning models often work by repeatedly passing information between connected items, but this can be slow and memory-intensive on large networks. A faster family of models, called pre-propagation graph neural networks, spreads information across the graph once before training and then trains more efficiently, but these models can lose accuracy, especially when connected items are not very similar. We focus on two reasons for this gap and address both. First, the usual way of mixing neighbor information mainly smooths things out, which can miss useful contrasts between different kinds of connected items. We replace it with more stable mixing methods based on well-understood mathematical tools. Second, the usual one-time preparation cannot use what the model learns during training, so we add a lightweight step that occasionally re-mixes the model’s learned features. In experiments, these changes improve accuracy and help pre-propagation graph neural networks match more expensive message-passing models while retaining their main training-efficiency advantages.