Physics-Informed Guidance

Enforcing physical constraints during the reverse diffusion process.

Standard diffusion models may generate vector fields that look visually correct but violate fundamental physical laws. Our approach integrates physics constraints directly into the sampling process using gradient guidance on the estimated clean data.

The Guidance Principle

Modify the noisy sample $x_t$ at each step by moving it against the gradient of a physical loss function $\mathcal{L}_{physics}$:

$$ x_t \leftarrow x_t - \eta_t \cdot \nabla_{x_t} \mathcal{L}_{physics}(\hat{x}_0(x_t)) $$

Tweedie's Formula

We cannot apply physics laws to noisy data $x_t$. Instead, we use Tweedie's formula to estimate the clean data $\hat{x}_0$:

$$ \hat{x}_0 = \frac{x_t - \sqrt{1-\bar{\alpha}_t}\epsilon_\theta(x_t)}{\sqrt{\bar{\alpha}_t}} $$

Implementation Strategies

We analyzed two methods for computing the gradient $\nabla_{x_t} \mathcal{L}_{physics}(\hat{x}_0)$. While conceptually different, they yield identical results under standard single-step approximations.

Correctness Autograd Method

Uses PyTorch's automatic differentiation to handle the chain rule through the Tweedie formula. Ensures theoretical correctness but may be slower. 

x_t = x_t.detach().requires_grad_(True)

# Predict clean sample
x0_pred = (x_t - sqrt_1_alpha * eps) / alpha

# Compute loss AND gradient automatically
loss = physics_loss(x0_pred)
grad = torch.autograd.grad(loss, x_t)[0]

x_t = x_t - eta * grad

Recommended Delta Noise Method

Computes gradient on $\hat{x}_0$ and manually scales it. This provides fast execution speed with minimal accuracy loss, making it the preferred method for real-time applications.

x0_pred = x0_pred.detach().requires_grad_(True)

loss = physics_loss(x0_pred)
grad_x0 = torch.autograd.grad(loss, x0_pred)[0]

# Manual chain rule approximation
grad = grad_x0 / alpha

x_t = x_t - eta * grad

Terminology Reference

Symbol Description Role in Guidance
$x_t$ Noisy sample at timestep $t$ The variable being updated (guided) towards adherence to physics laws.
$\hat{x}_0$ Estimated clean sample Derived via Tweedie's formula; used to compute physical loss.
$\epsilon_\theta$ Predicted noise Output of the neural network, estimates the noise added to $x_0$.
$\eta_t$ Guidance scale (step size) Controls how strongly the physics gradient modifies $x_t$.
$\bar{\alpha}_t$ Cumulative noise schedule term Defines the signal-to-noise ratio at timestep $t$.
$\nabla$ Gradient operator Direction in which to move $x_t$ to minimize the physics loss.

Key Physics Constraints

1

Divergence (Incompressibility)

Ensures the fluid is incompressible (mass neither created nor destroyed).

$$ \mathcal{L}_{div} = \|\nabla \cdot \mathbf{v}\|^2 $$
2

Curl (Irrotationality)

Ensures the flow is free of vortices (irrotational).

$$ \mathcal{L}_{curl} = \|\nabla \times \mathbf{v}\|^2 $$