From Binary to Continuous: Stochastic Re-Weighting for Robust Graph Explanation

Graph Neural Networks (GNNs) have achieved remarkable performance in a wide range of graph-related learning tasks. However, explaining their predictions remains a challenging problem, especially due to the mismatch between the graphs used during training and those encountered during explanation. Most existing methods optimize soft edge masks on weighted graphs to highlight important substructures, but these graphs differ from the unweighted graphs on which GNNs are trained. This distributional shift leads to unreliable gradients and degraded explanation quality, especially when generating small, sparse subgraphs. To address this issue, we propose a novel iterative explanation framework which improves explanation robustness by aligning the model's training data distribution with the weighted graph distribution appeared during explanation. Our method alternates between two phases: explanation subgraph identification and model adaptation. It begins with a relatively large explanation subgraph where soft mask optimization is reliable. Based on this subgraph, we assign importance-aware edge weights to explanatory and non-explanatory edges, and retrain the GNN on these weighted graphs. This process is repeated with progressively smaller subgraphs, forming an iterative refinement procedure. We evaluate our method on multiple benchmark datasets using different GNN backbones and explanation methods. Experimental results show that our method consistently improves explanation quality and can be flexibly integrated with different architectures.