Prior-Informed Flow Matching for Graph Reconstruction

TL;DR

Prior-Informed Flow Matching (PIFM) tackles reconstructing graphs from partial observations by fusing embedding-based priors with a permutation-equivariant, continuous-time flow. The method first computes a MMSE-like posterior mean from local information using inductive (graphon, GraphSAGE) or transductive (node2vec) priors, then learns a rectified flow to transport this initialization toward the true graph distribution, enforcing global consistency. Empirical results on several benchmark datasets (IMDB-B, PROTEINS, ENZYMES, CORA) show that PIFM consistently improves over both pure priors and existing flow-based baselines, across link prediction, expansion, and denoising tasks. The approach highlights the value of explicitly incorporating global structure into graph reconstruction, offering improved fidelity and a scalable, principled pathway to graph inpainting.

Abstract

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.

Figures (30)

• Figure 1: Overview of the Prior-Informed Flow Matching (PIFM) graph reconstruction framework. Starting from a partially observed adjacency matrix ${\mathbf A}_1^\xi = \xi \odot {\mathbf A}$, where $\xi$ denotes a mask, we form an initialization ${\mathbf A}_0$ by combining the observed entries with prior predictions $f_{\text{prior}}({\mathbf A}_1^\xi)$ obtained with an element-wise predictor. In dark red we denote the true edges that are masked, while in light red those masked positions where there is no edge between nodes. A rectified flow then interpolates linearly from ${\mathbf A}_0$ to the ground-truth graph ${\mathbf A}_1 = {\mathbf A}$, learning global structural information from a coupling of all the edges. The intermediate states ${\mathbf A}_t$ improve on the prior-informed initialization, enabling recovery of the missing edges.
• Figure 2: Toy experiment showcasing the advantage of PIFM (in this case, for link prediction). a) Graph ${\mathcal{G}}$ with four nodes, where the hidden edges are $e_{02}$ and $e_{13}$. b) Generated samples by using node2vec and PIFM (our proposed method): clearly, our method learns a probabilistic coupling, rendering a model that generates only the two valid modes. c) Proportions of samples generated with PIFM from each mode; remarkably, the method also learns a good approximation of the probability of each mode.
• Figure 3: Evaluation of the distortion-perception trade-off. (a) The MMD score, measuring the distance to the true data distribution, decreases as $K$ increases. (b-d) This result is corroborated by key graph statistics, where the average degree, number of triangles, and clustering coefficient for graphs generated with $K=100$ more closely match the ground-truth distribution compared to those generated with $K=1$. Error bars indicate the standard deviation over $300$ samples ($10$ samples for each of the $30$ test graphs).
• Figure 4: ROC as a function of the noise $\sigma_s$ in $p({\mathbf A}_0)$. The impact of noise level $\sigma_s$ on model performance, measured by the best ROC AUC score. Results are shown for two different drop rates: 0.1 (blue) and 0.5 (orange). A small amount of noise improves performance for both configurations, after which increasing noise leads to performance degradation.
• Figure 5: An analysis of the ROC AUC score as a function of the number of processing steps (K) for a drop rate of 10%. This experiment was conducted with a fixed drop rate of 0.1, while varying the noise level, $\sigma_s$. The results show that the optimal number of steps is small, typically under 10.

Theorems & Definitions (1)

• proof