ΒΆPaper Feed

Issue 23 Β· Pick 06 AI / ML βœ“ read

The Right Measure for Physics-Constrained Generation: A Co-Area Correction for Posterior-Consistent PDE Inverse Problems

Jian Xu, Yanning Wu, Delu Zeng, John Paisley, Qibin Zhao

TL;DR. A wave of recent papers solves PDE inverse problems by taking a diffusion or flow-matching prior, forcing samples to exactly satisfy the physics (by projection or guidance), and reporting the results as a calibrated Bayesian posterior. This paper shows that recipe samples the wrong distribution β€” not by an implementation bug, but because conditioning on an exact constraint is conditioning on a measure-zero set, and the physically correct way to do that carries a Jacobian factor [\det(JJ^\top)]^{-1/2} that projection and guidance silently drop. The factor has been known in molecular dynamics since 1974 (the Fixman correction); the contribution here is showing it is exactly what's missing from physics-constrained generative inference, quantifying the bias (up to ~20–40Γ— the sampling-noise floor, with posterior standard deviations off by 96%), and giving a corrected sampler (CoCoS) plus an amortized fast version (CoCo-Flow).

The setup: physics as a hard constraint

The standard pipeline for generative PDE inverse problems looks like this. You have an unknown field \boldsymbol\theta (say, a log-permeability field expressed in a truncated basis), a differentiable PDE solver \mathcal{G}, and sparse noisy observations \mathbf{y} = \mathcal{H}(\mathcal{G}(\boldsymbol\theta^\star)) + \boldsymbol\eta. A pretrained generative model gives you a prior \pi over \boldsymbol\theta. To make samples "respect the physics," you enforce a hard constraint \mathbf{c}(\mathbf{x}) = \mathbf{0} β€” the PDE residual, boundary conditions, or data consistency β€” most commonly via minimal-displacement projection: move each generated sample to the nearest point on the constraint manifold \mathcal{M} = \{\mathbf{x} : \mathbf{c}(\mathbf{x}) = \mathbf{0}\}, implemented with Gauss–Newton steps. That's what PCFM and its relatives do. Guidance-based methods instead steer the sampling dynamics with the residual gradient. Either way, the constrained samples get reported as posterior draws with error bars.

The question this paper asks is simple and uncomfortable: does enforcing the constraint actually give you the posterior?

The measure problem, and why it's genuinely subtle

The manifold \mathcal{M} is lower-dimensional, so it has probability zero under any continuous prior. "Condition on \mathbf{c} = \mathbf{0}" is therefore not a well-posed operation by itself β€” this is the classical Borel–Kolmogorov paradox. The answer depends on how you shrink a positive-measure set down to the manifold, and there are two natural choices that give different answers:

  • The residual limit: condition on \{\|\mathbf{c}(\mathbf{x})\| \le \varepsilon\} and let \varepsilon \to 0. This is the physically meaningful one, because it's the limit of the honest model "the PDE holds up to small residual noise, \mathbf{c}(\mathbf{x}) \sim \mathcal{N}(\mathbf{0}, \gamma^2 \mathbf{I})" as \gamma \to 0.
  • The Euclidean limit: condition on the geometric tube \{\mathbf{x} : \mathrm{dist}(\mathbf{x}, \mathcal{M}) \le \varepsilon\}. This is what nearest-point projection implicitly targets β€” it collapses each Euclidean normal fiber onto the manifold.

Here's the intuition for why they differ. The residual band \|\mathbf{c}\| \le \varepsilon is not a tube of uniform thickness. Where the constraint is sensitive (the gradient J = \nabla\mathbf{c} is large), the band is thin: a tiny move off the manifold blows up the residual. Where the constraint is flat (small J), the band is fat: you can wander far and still nearly satisfy the physics. So when you condition the prior on the residual band and shrink it, flat regions of the manifold soak up more probability mass than sensitive ones β€” by exactly the local band thickness, which is [\det(JJ^\top)]^{-1/2} for a codimension-m constraint. The Euclidean tube, by contrast, has uniform thickness everywhere, so it weights the manifold by the bare prior.

Residual band β€–c(x)β€– ≀ Ξ΅ thin band: β€–βˆ‡cβ€– large fat band: β€–βˆ‡cβ€– small β†’ posterior ∝ Ο€ Β· [det(JJα΅€)]^(βˆ’1/2) Euclidean tube dist(x, M) ≀ Ξ΅ uniform thickness everywhere β†’ posterior ∝ Ο€ (what projection targets)
Two ways to shrink onto the constraint manifold. The residual band (left) is fat where the constraint is insensitive and thin where it's steep, so the limiting posterior is tilted by the co-area factor. The geometric tube (right) β€” the implicit target of nearest-point projection β€” has uniform width and drops that factor. They agree only if $\det(JJ^\top)$ is constant on the manifold, which essentially never happens for a PDE.

The formal statement runs through the co-area formula (Federer): the ambient Lebesgue integral disintegrates over the level sets of \mathbf{c} with the factor [\det \mathrm{G}]^{-1/2}, where \mathrm{G} = JJ^\top is the m \times m constraint Gram matrix. Theorem 1 of the paper shows the soft posterior p_\gamma \propto \pi \,\ell(\mathbf{y}\mid\cdot)\, e^{-\|\mathbf{c}\|^2/2\gamma^2} converges as \gamma \to 0 to

p^\star(\mathbf{x}) \;\propto\; \pi(\mathbf{x})\,\ell(\mathbf{y}\mid\mathbf{x})\,\big[\det(JJ^\top)\big]^{-1/2}, \qquad \mathbf{x} \in \mathcal{M},

where \ell is the data likelihood. Proposition 1 shows the Euclidean tube limit gives p^{\mathrm{E}} \propto \pi with no Jacobian factor, and Proposition 2 shows minimal-displacement projection's pushforward agrees with p^{\mathrm{E}} to leading order. The gap between what you should sample and what you do sample is exactly the Fixman factor β€” the same correction that appears in constrained molecular dynamics when rigid bond constraints replace stiff springs. The paper also gives it a nice variational reading (Prop. 3): p^\star is the \gamma \to 0 limit of the minimizers of \mathrm{KL}(q \| \pi_y) + \frac{1}{2\gamma^2}\mathbb{E}_q\|\mathbf{c}\|^2, whereas projection minimizes Euclidean transport cost β€” a different objective. There's an appealing analogy to the natural gradient: the residual-noise model measures discrepancy in \mathbf{c}-space, endowing parameter space with the pullback metric J^\top J; projecting in the ambient Euclidean metric is conditioning in the wrong metric.

Crucially, the bias scales with how much \det \mathrm{G} varies over the posterior. If the constraint sensitivity is homogeneous, the two limits coincide (Corollary 1) and projection is fine. For PDEs β€” where sensitivity is dictated by the operator and where your sensors happen to sit β€” it is generically very heterogeneous: the paper measures 12–33Γ— variation on 1D Darcy problems and over five orders of magnitude on a 2D Darcy problem.

CoCoS: sampling the right measure

The fix imports machinery from manifold MCMC (Zappa, Holmes-Cerfon & Goodman 2018). Write p^\star \propto e^{-V} on \mathcal{M} with

V(\mathbf{x}) = -\log\pi(\mathbf{x}) - \log\ell(\mathbf{y}\mid\mathbf{x}) + \tfrac{1}{2}\log\det\big(J(\mathbf{x})J(\mathbf{x})^\top\big).

Each CoCoS step: (1) propose an isotropic Gaussian move in the tangent space of \mathcal{M} at the current point; (2) project the proposal back onto \mathcal{M} with a small Newton solve along the normal directions; (3) Metropolis accept/reject using V, with a reverse-projection check that guarantees reversibility. Theorem 2 shows p^\star is the invariant law. Two practical points worth noting: the Fixman term enters only through values of \log\det\mathrm{G} (no gradients of the log-det needed), and J is already computed by any projection-based method, so the marginal cost over PCFM is an m \times m log-determinant per step. A subtle robustness point: the Newton solve can be regularized (\mathrm{G} + \epsilon\mathbf{I}) without biasing the target, because regularization only affects the proposal while the acceptance ratio uses the exact potential β€” a structural advantage over projection methods, where any such regularization directly contaminates the reported samples.

One clarifying result deserves emphasis (Prop. 4): amortized inference trained on forward-simulated joint pairs (\boldsymbol\theta, \mathbf{y}) β€” simulate \boldsymbol\theta \sim \pi, run the solver, add noise, fit a conditional flow β€” is automatically measure-correct. The joint draws respect the residual-tube geometry by construction; you never condition on the manifold. The bias afflicts specifically the popular regime where you post-process a fixed pretrained prior at test time. This cleanly partitions the literature into methods that are fine (simulation-based amortized flow matching) and methods that aren't (projection/guidance on a frozen prior), and the paper exploits it: CoCo-Flow distills exact CoCoS samples (or forward-simulated pairs) into a conditional flow-matching student, so the co-area cost is paid once at training and test-time inference is a single ODE solve.

The evidence

The cleverest methodological choice is the arbiter. MCMC "gold standards" fail here: HMC on the soft posterior at small \gamma stops mixing (Appendix B shows ESS collapsing 737 β†’ 21 and \hat{R} blowing past 1.2 as \gamma shrinks), so it can masquerade as agreement or disagreement depending on tolerances. Instead the paper uses i.i.d. rejection: draw from the prior, keep samples with \|\mathbf{c}\| < \varepsilon, verify convergence as the band shrinks. Brutally expensive (up to 5.4 \times 10^8 raw draws for 18k accepted at the tightest band) but unimpeachable β€” and it converges to p^\star, the residual limit, by construction.

On the controlled d{=}4 benchmark (single quadratic constraint, analytic gradients, sensitivity varying 4.2Γ—):

Controlled d=4 benchmark: distance to i.i.d. ground truthavg W1 to arbiter (Γ— noise floor)051015202.5CoCoS (Fixman)9PCFM projection16PCFM + scalar reweight21No Fixman (Hausdorff)Table 1. Noise floor = arbiter's own two-sample W1 distance (0.004).

Every point here matters. CoCoS matches the arbiter to within sampling noise, confirming the residual limit really is the co-area posterior. Dropping the Fixman term is off by 21Γ—. Projection is off by 9Γ—. And β€” the result that closes the obvious escape hatch β€” reweighting projected samples by a scalar [\det\mathrm{G}]^{-1/2} makes things worse (16Γ—), because the projection pushforward is a curvature-dependent distortion of the prior, not a pure density rescaling. You can't patch projection after the fact; you need a sampler that targets the right measure.

On the 1D Darcy inverse problem (m{=}3 pressure sensors, codimension-3 manifold, d{=}8 and d{=}16 basis coefficients), the ordering is invariant across all instances, and CoCoS actually gets cleaner at higher dimension (1.2Γ— the floor at d{=}16) while PCFM sits at ~16–36Γ—. Most damning for the UQ claims is Table 3, which compares one representative method from each paradigm:

Uncertainty quality, d=8 Darcy (3-instance mean)relative error of posterior std (%)0204060801001CoCoS26Soft penalty (Ξ³=0.02)56Guided Langevin96PCFM projectionTable 3. Nominal-90% interval coverage: CoCoS 0.89, PCFM 0.95, guidance 0.93 (ideal 0.90). PCFM's covariance is off by 122% in Frobenius norm (Fig. 2), inventing correlations the posterior doesn't have.

Projection gets the posterior spread wrong by 96% and distorts the covariance structure by 122% β€” it manufactures spurious correlations. This is a calibration failure, not a cosmetic shift: if you're using these samples for scientific error bars, the error bars are wrong.

The mechanism check (Fig. 6 of the paper) is the most satisfying piece of evidence: binning samples by \sqrt{\det(JJ^\top)}, the density ratio of the biased samplers to the true posterior tracks the predicted \propto \sqrt{\det(JJ^\top)} law almost exactly, while CoCoS lies flat on the unbiased line. The bias isn't noise; it's the missing Jacobian, recovered quantitatively.

The theory also predicts when the bias doesn't matter, and the paper honestly tests that: on a 1D viscous Burgers problem with mild sensitivity heterogeneity, CoCoS, no-Fixman, and PCFM are nearly indistinguishable (1.8/2.2/2.3Γ— the floor). This is good scientific hygiene β€” the effect turns on and off exactly where Corollary 1 says it should.

Finally, scaling: on a 2D Darcy problem (d{=}64 KLE coefficients, 16{\times}16 grid), where \sqrt{\det\mathrm{G}} spans five orders of magnitude, the amortized CoCo-Flow trained on forward-simulated pairs recovers the posterior at held-out observations (0.9Γ— the floor in the likelihood-informed subspace) while PCFM sits at 3.1Γ—, with ~56 ms per query after one-time training β€” versus minutes per query for the rejection arbiter.

What to make of it

If this holds, the practical takeaway is sharp: "satisfying the physics" and "sampling the posterior" are different objectives, and the entire projection/guidance-on-a-frozen-prior family produces miscalibrated uncertainty whenever constraint sensitivity is heterogeneous β€” which for PDEs with sparse sensors is the generic case. The correction is not exotic: either add the \frac{1}{2}\log\det(JJ^\top) term to a proper constrained sampler (the Jacobian is already computed by these methods), or sidestep the whole issue by training on forward-simulated joint pairs, which Prop. 4 certifies as measure-correct. That second point may be the most consequential for practice: it says simulation-based amortized inference and post-hoc constraint enforcement are not interchangeable design choices, and the field's drift toward test-time constraint enforcement on pretrained priors has a hidden statistical cost.

Reasons for skepticism. The exact validation lives in low dimensions (d \le 16, codimension 3) because the i.i.d. arbiter is exponentially expensive in the band tightness; at d{=}64 the verification is indirect (a 6-dimensional likelihood-informed subspace, with ~61 prior-dominated directions diluting the raw metric β€” the paper acknowledges the full-coordinate average looks deceptively fine). The priors are standard Gaussians, not learned diffusion priors, so "does the correction remain tractable and accurate when \pi is a real pretrained generative model over fields" is untested β€” evaluating \log\pi for a diffusion prior is itself nontrivial. Codimension is tiny throughout (m{=}3); for constraints imposed at every grid point, the m \times m log-det becomes a serious cost (the paper waves at Hutchinson estimators). Near-singular \mathrm{G} in weakly identified directions makes the Fixman factor huge β€” the paper argues persuasively this is correct Bayesian behavior (up-weighting prior-dominated directions by exactly how much the residual tube widens), but it will stress numerics at scale. And the framing depends on the residual limit being the right limit: the paper's argument β€” the physics only ever holds up to discretization/model error, so the constraint is really a \gamma \to 0 Gaussian residual β€” is convincing, but a method that genuinely intends the Euclidean tube as its target isn't "wrong," it's just answering a different (and arguably unphysical) question. CoCo-Flow's evaluation is also preliminary by the authors' own admission (5.0–5.7Γ— the floor on the harder problems).

Where to spend your time: the theory section ("The Measure of a Hard Constraint") through Corollary 1 is the heart β€” it's short, and once you've internalized the two-tubes picture, everything else follows. Then Fig. 6, which shows the predicted \sqrt{\det(JJ^\top)} bias law recovered directly from samples, and Appendix B for the cautionary tale about HMC as an arbiter, which is a methodological point the broader physics-constrained generation literature should absorb regardless of whether they adopt CoCoS.