Gradient-Free Generation for Hard-Constrained Systems
Chaoran Cheng, Boran Han, Danielle C. Maddix, Abdul Fatir Ansari, Andrew Stuart, Michael W. Mahoney, Yuyang Wang
TL;DR
This paper tackles the challenge of hard-constrained generation for PDE systems by proposing ECI sampling, a gradient-free framework that steers a pre-trained unconstrained flow-matching model toward exact constraint satisfaction in a zero-shot setting. The method interleaves extrapolation, correction, and interpolation at each sampling step to transfer constraint information from the final prediction back through the iterative flow, while controlling stochasticity via re-sampling. Empirically, ECI achieves state-of-the-art or competitive performance across multiple PDEs (Stokes, Heat, Darcy, Navier–Stokes) in generative and regression tasks, delivering exact constraint satisfaction, improved distributional fidelity (e.g., Fréchet Poseidon distance), and substantial efficiency gains over gradient-based baselines. This gradient-free approach enables robust zero-shot constrained generation and uncertainty-aware predictions, with potential applicability to broader SciML tasks beyond PDEs.
Abstract
Generative models that satisfy hard constraints are critical in many scientific and engineering applications, where physical laws or system requirements must be strictly respected. Many existing constrained generative models, especially those developed for computer vision, rely heavily on gradient information, which is often sparse or computationally expensive in some fields, e.g., partial differential equations (PDEs). In this work, we introduce a novel framework for adapting pre-trained, unconstrained flow-matching models to satisfy constraints exactly in a zero-shot manner without requiring expensive gradient computations or fine-tuning. Our framework, ECI sampling, alternates between extrapolation (E), correction (C), and interpolation (I) stages during each iterative sampling step of flow matching sampling to ensure accurate integration of constraint information while preserving the validity of the generation. We demonstrate the effectiveness of our approach across various PDE systems, showing that ECI-guided generation strictly adheres to physical constraints and accurately captures complex distribution shifts induced by these constraints. Empirical results demonstrate that our framework consistently outperforms baseline approaches in various zero-shot constrained generation tasks and also achieves competitive results in the regression tasks without additional fine-tuning.