Stabilizing Physics-Informed Consistency Models via Structure-Preserving Training

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.

💡 Research Summary

The paper tackles the instability that arises when incorporating physics constraints into Consistency Models (CMs), which are fast generative models that replace the iterative sampling of diffusion models with a single or few-step mapping from noise to data. Directly adding PDE residual losses to the CM objective often drives the model toward trivial or overly smooth solutions that minimize the residual but collapse the learned data distribution, especially in high‑dimensional PDE settings.

To resolve this, the authors propose a two‑stage training pipeline. In Stage 1, the model is trained purely on the standard continuous‑time CM loss (sCM) to learn the joint distribution of coefficients a and solutions u without any physics term. They adopt the TrigFlow parameterization, which smoothly interpolates between data and a Gaussian prior, ensuring stable learning of the manifold p(a,u).

In Stage 2, the model is fine‑tuned with a composite loss that adds a PDE residual term ‖R(u,a)‖². Crucially, the “coefficient decoder” that generates a is frozen after Stage 1, preventing the optimizer from collapsing a toward low‑complexity configurations that would trivially reduce the residual. The architecture therefore splits the decoder into two parallel branches after a shared encoder: a frozen branch that reconstructs a and a trainable branch that produces u conditioned on the frozen latent representation. The final output concatenates the frozen a with the newly generated u, guaranteeing that the physical parameters remain on the learned distribution while the solution is driven toward physical consistency.

The paper also introduces a “two‑step consistency sampling operator.” Instead of applying the physics loss to the raw one‑step prediction (which is noisy and high‑frequency), they perform an additional re‑noising and denoising step, effectively projecting the sample onto a smoother manifold before evaluating the PDE residual. This reduces gradient variance from the residual term and stabilizes training.

Experiments on synthetic 2‑D manifolds (circle, ellipse, double ellipse) and on real PDEs such as Poisson, wave, and Navier‑Stokes equations demonstrate several key outcomes:

• Training stability: Direct physics‑constrained training collapses to a small subset of the manifold, whereas the two‑stage protocol preserves high coverage.
• Accuracy: Using the relative H¹ error (which accounts for both solution values and spatial derivatives), the proposed sCM‑PINN achieves comparable or lower errors than physics‑informed diffusion models while using only 1–4 sampling steps.
• Computational efficiency: The zero‑shot projection‑based inpainting procedure solves forward problems with a single consistency mapping, yielding speed‑ups of two orders of magnitude compared to diffusion‑based samplers that require thousands of steps.
• Ablation studies: Freezing the coefficient decoder is essential; without it, the model drifts toward constant‑coefficient solutions despite low residuals. The two‑step residual operator further improves convergence and reduces mode collapse.

The authors also contribute a new evaluation metric (relative H¹ error) and release the full codebase for reproducibility. In summary, the work presents three intertwined innovations—(1) a two‑stage distribution‑then‑physics training schedule, (2) a structure‑preserving channel‑partitioned architecture that freezes the coefficient manifold, and (3) a two‑step residual enforcement on refined samples—that together stabilize physics‑informed consistency models and enable fast, high‑fidelity PDE solution generation and forward inference.