Graph reconstruction from partial observations often comes with structural side information, such as degree bounds, triangle counts, or an edge-density band. Prior-Informed Flow Matching (PIFM) reconstructs graphs by transporting a local prior toward the graph distribution, but it provides no mechanism to incorporate this side information. We put forth Constrained Primal-Dual PIFM (CPD-PIFM), which augments the sampler with Lagrange multipliers that evolve along each trajectory. The multipliers respond to constraint violations at a predicted endpoint and guide subsequent sampling steps without retraining. We prove that the sampler inherits PIFM's permutation equivariance and bound its expected terminal slack by a term that decays as the inverse square root of the number of steps, plus two approximation terms. On three link-prediction benchmarks and nine combinations of datasets and constraints, CPD-PIFM raises feasibility by 11-26 percentage points and remains competitive with fixed guidance without selecting a separate multiplier for each constraint.
Figures & tables
Figure 1: (a) PROTEINS with competing edge-floor and degree-cap constraints. Curves show feasibility versus K for PIFM and for CPD-PIFM with increasing and decaying schedules for λk . (b)–(c) ENZYMES with tri≥10 . Each panel shows triangle counts on full reconstructions from 200 paired samples of one partially observed graph; only the Gaussian source noise changes. In (b), P(feasible) changes from 0.105 to 1.00 , and the mean (triangle markers) changes from 7.0 to 22.7 . In (c), feasibility stays at 1.00 , and the mean changes from 22.7 to 23.0 .
Feas. % ↑
AUC ↑
MMD ↓
Constraint
PIFM
CPD
PIFM
CPD
PIFM
CPD
ENZYMES
deg. cap
83.2±7.4
99.4±1.4
.550
.554
.205
.204
tri. floor
74.2±4.0
99.4±1.4
.552
.579
.222
.318
mixture
47.1±5.9
71.0±6.0
.560
.576
.243
.199
competing
43.9±10.4
69.7±8.7
.551
.561
.225
.172
Table 1: Effect of CPD guidance by constraint at K=32 , averaged over five test seeds. S.d. is shown for feasibility only; the s.d. of the paired AUC and MMD changes is at most 0.010 and 0.07 .
Fixed η=0.1
Fixed η=1.6
CPD-PIFM
Constraint
Feas
AUC
MMD
Feas
AUC
MMD
Feas
AUC
MMD
ENZYMES
deg. cap
93.5
.569
.210
100
.572
.597
100
.567
.195
tri. floor
74.2
.578
.195
100
.618
1.007
100
.590
.385
mixture
58.1
.603
.204
80.6
.681
.199
74.2
.599
.168
competing
58.1
.567
.213
74.2
.610
.081
80.6
.568
.184
Table 2: Fixed guidance and CPD-PIFM at K=32 , single seed. Both fixed multipliers η are shared across all settings; PROTEINS edge-band gradients cancel under shared η , so fixed guidance equals PIFM. Feasibility is reported as a percentage; the highest in each setting is in bold.
We consider the problem of learning to generate data while satisfying a prescribed sample-wise constraint with a user-specified probability of constraint satisfaction. Starting with a population-level constrained negative log-likelihood minimization and using a flow-matching (FM) parameterization, under regularity assumptions and sufficient expressivity of the velocity field family, we derive guarantees on the probability of constraint satisfaction for the corresponding primal-dual optimization algorithm. We provide a practical implementation of this approach, called Primal-Dual Flow Matching (PDFM), which has direct control over constraint satisfaction rate and requires only binary membership oracle feedback for constraint satisfaction, avoiding the need for differentiable distances, projections, or special structure of the constraint set, such as convexity. We show that PDFM improves constraint satisfaction over vanilla FM and achieves competitive performance against constraint-aware baselines while maintaining distributional fidelity. Our code is available at https://github.com/ZhengyanHuan/PDFM.
Zhengyan Huan, Peter Y. Lu, Shuchin Aeron
Department of Electrical and Computer Engineering, Tufts University · NSF AI Institute for Artificial Intelligence and Fundamental Interactions
Reconstructing PDE solutions from sparse observations is a core challenge in scientific computing. We present FM4PDE, a flow-matching generative framework that learns the joint distribution of PDE coefficients (or initial states) and solutions (or final states), enabling both forward simulation and inverse recovery with limited paired data. At inference, sampling is guided by a composite loss that enforces agreement with sparse measurements and reduces the PDE residual; we support deterministic, stochastic, and hybrid samplers. We provide error guarantees for these guided procedures. For the deterministic optimizer, a coercivity condition ensures trajectory boundedness and a phase-wise contraction yields logarithmic complexity in the target accuracy. For the stochastic sampler, we introduce adaptive guidance and assume dissipativity of the velocity field to obtain uniform moment bounds independent of the noise-floor parameter. This leads to polynomial-time error bounds, and a matching lower bound shows constant guidance induces an unavoidable positive bias, motivating adaptivity. A hybrid deterministic-stochastic analysis is also provided. Experiments on static and time-dependent benchmark PDEs demonstrate competitive accuracy and faster inference than diffusion-based generative models.
Xifeng Zhang, Jin Zhao
School of Mathematical Science Capital Normal University Beijing, 100048, China · Academy for Multidisciplinary Studies Capital Normal University2026 Beijing, 100048, China
We propose ConicFlow, which frames constrained sampling from a flow matching model at inference time as an optimal control problem that balances deviation from the pretrained trajectory against reward, subject to hard terminal constraints. We discretize this problem and adapt receding-horizon control, optimizing only a short window of the trajectory at each step for computational tractability. We address the resulting sequence of receding-horizon subproblems with a receding-horizon primal-dual algorithm. Its dual multipliers capture the difficulty of satisfying each constraint and are refined throughout sampling, steering generation toward feasible final samples. We further propose a terminal-estimate variant that avoids accumulating local truncation errors by evaluating the reward and constraints at the predicted terminal point rather than at the end of the current horizon. Under standard regularity assumptions, we establish finite-horizon guarantees on suboptimality. Under convex-affine structure, we derive a sharper bound that scales linearly with omitted horizon length, and we show that the terminal-estimate variant is provably tighter when the flow dynamics change slowly. Experiments in both low- and high-dimensional settings with a recent large-scale flow matching model show that ConicFlow achieves high reward and strong constraint satisfaction while deviating less from the pretrained flow than recent baselines.
Shervin Khalafi, Kasra Fallah, Sung Cho +3
University of Pennsylvania · Columbia University · University of Tennessee, Knoxville