Prior-Informed Flow Matching for Graph Reconstruction

We introduce Prior-Informed Flow Matching (PIFM), a conditional flow model for graph reconstruction. Reconstructing graphs from partial observations remains a key challenge; classical embedding methods often lack global consistency, while modern generative models struggle to incorporate structural priors. PIFM bridges this gap by integrating embedding-based priors with continuous-time flow matching. Grounded in a permutation equivariant version of the distortion-perception theory, our method first uses a prior, such as graphons or GraphSAGE/node2vec, to form an informed initial estimate of the adjacency matrix based on local information. It then applies rectified flow matching to refine this estimate, transporting it toward the true distribution of clean graphs and learning a global coupling. Experiments on different datasets demonstrate that PIFM consistently enhances classical embeddings, outperforming them and state-of-the-art generative baselines in reconstruction accuracy.

💡 Research Summary

This paper introduces Prior-Informed Flow Matching (PIFM), a novel framework designed for the fundamental problem of graph reconstruction from partial observations. The core challenge lies in bridging the gap between classical methods, which excel at local link prediction but lack global structural consistency, and modern generative models, which can produce plausible graphs but are not optimized for accurate, faithful recovery of a specific ground-truth structure from a partial view.

PIFM addresses this by reformulating the problem through the lens of the distortion-perception trade-off, adapted for graphs with permutation equivariance. The theoretical insight is that an optimal estimator can be constructed in two stages: first, computing the Minimum Mean Squared Error (MMSE) estimate (the posterior mean) using local information, and second, applying an optimal transport map to refine this estimate towards the true distribution of clean graphs, thereby enforcing global structural fidelity.

The method practically implements this two-stage theory. In the prior-informed stage, PIFM leverages existing inductive or transductive graph representation techniques to form an intelligent initial guess for the missing adjacency matrix entries. This can be achieved using dataset-level priors like graphons (learned via methods like SIGL) or instance-level predictors like GraphSAGE or node2vec, which output probabilities for each potential edge based on node embeddings. This stage provides the crucial local context.

The key innovation lies in the flow matching stage. Instead of starting from generic noise, PIFM uses the probabilistic prediction from the first stage as the source distribution. It then trains a rectified flow model to learn a continuous-time trajectory that transports this initial estimate to the target distribution of clean, binary adjacency matrices. The model learns a velocity field that dictates how to move from any intermediate state along this trajectory towards the ground truth. Critically, during this process, the model learns the global coupling between edges—the hidden correlations and structural patterns that are not apparent when predicting edges independently. The architecture is designed to be permutation equivariant, respecting the inherent symmetry of graphs.

The authors validate PIFM through comprehensive experiments on graphs with varying densities. The tasks include standard link prediction and two more challenging “blind” tasks: expansion (recovering truly missing edges) and denoising (removing spuriously added edges). Results demonstrate that PIFM consistently and significantly enhances the reconstruction performance of all base priors (GraphSAGE, node2vec, graphons). Furthermore, it outperforms state-of-the-art generative baselines adapted for constrained generation, such as DiGress (with inpainting) and PRoDiGY. This empirically confirms that PIFM successfully integrates the local accuracy of embedding methods with the global consistency and structural awareness of flow-based generative models, establishing a new effective approach for high-fidelity graph topology inference.