Stabilizing Physics-Informed Consistency Models via Structure-Preserving Training

TL;DR

We address ill-posed PDE problems by learning a joint distribution over coefficients and states with a physics-informed consistency framework. The method, sCM-PINN, uses a two-stage training protocol and a structure-preserving channel partition to stabilize optimization, coupled with a two-step residual constraint and a projection-based forward solver for fast, physically consistent inference. Empirical results on Darcy, Poisson, and Helmholtz show strong forward accuracy (relative $H^1$) with far fewer function evaluations than diffusion-based baselines, and unconditional samples exhibit lower PDE residuals with minimal compute. The approach promises real-time, physically grounded scientific simulations and scalable extensions to more complex, time-dependent problems.

Abstract

We propose a physics-informed consistency modeling framework for solving partial differential equations (PDEs) via fast, few-step generative inference. We identify a key stability challenge in physics-constrained consistency training, where PDE residuals can drive the model toward trivial or degenerate solutions, degrading the learned data distribution. To address this, we introduce a structure-preserving two-stage training strategy that decouples distribution learning from physics enforcement by freezing the coefficient decoder during physics-informed fine-tuning. We further propose a two-step residual objective that enforces physical consistency on refined, structurally valid generative trajectories rather than noisy single-step predictions. The resulting framework enables stable, high-fidelity inference for both unconditional generation and forward problems. We demonstrate that forward solutions can be obtained via a projection-based zero-shot inpainting procedure, achieving consistent accuracy of diffusion baselines with orders of magnitude reduction in computational cost.

Figures (4)

• Figure 1: Necessity of Two-Stage Training. We visualize samples generated by physics-constrained consistency training on 2D toy manifolds (Circle, Ellipse, Double Ellipse). (Left) Directly enforcing constraints during consistency training leads to unstable dynamics and severe mode collapse, covering only a small fraction of the target manifold. (Right) Our two-stage protocol first learns a stable, high-coverage distribution (Stage 1), then refines samples to satisfy geometric constraints (Stage 2), achieving both diversity and physical validity.
• Figure 2: Architecture of the Structure-Preserving Consistency Model. To mitigate coefficient mode collapse during physics-informed fine-tuning, we employ a split-decoder architecture. The Frozen Backbone (blue, dashed) consists of the shared encoder and the Stage 1 decoder, whose parameters are fixed to preserve the learned coefficient distribution $\mathbf{a}$. The Active Branch (red, solid) is a trainable decoder that generates the solution field $\mathbf{u}$ conditioned on the fixed latent representation. The final output is the joint state $\mathbf{x} = (\mathbf{a}, \mathbf{u})$.
• Figure 3: Unconditional Sampling Quality (Darcy Flow). Visual comparison of generated coefficient fields $\mathbf{a}$ (left), solution fields $\mathbf{u}$ (middle), and absolute PDE residuals $|\mathcal{R}|$ (right). The baseline sCM (middle row) exhibits visible artifacts, whereas sCM-PINN (bottom row) produces sharper coefficient interfaces and solution fields with lower and more uniform residuals, comparable to those obtained by the 63-step DiffusionPDE baseline (top row) while requiring only 2 function evaluations. Note the reduced residual variance in the sCM-PINN samples relative to the diffusion baseline.
• Figure 4: Analysis of Coefficient Mode Collapse.(Top) Pixel-wise histogram showing the bimodal ground-truth distribution (gray). Joint training (red) collapses the distribution to a single mode at $a \approx 12$. (Bottom) Representative spatial samples. The joint-training baseline (middle row) produces only high-permeability features, whereas the frozen decoder strategy (bottom row) recovers the full coefficient structure.

Theorems & Definitions (1)

• definition 1: Two-Step Consistency Sampling Operator