Freeze, Then Select: Structured Field Adapters and Stability-Validated Weak Selection for PDE Discovery from Sparse Observations

Overview§

This paper targets PDE discovery from sparse, noisy spatiotemporal observations. The core difficulty lies in jointly reconstructing a continuous field and selecting the correct differential terms. The authors analyze coupled neural PDE discovery and identify three optimization behaviors: (1) the exact support persists to the end, (2) appears only transiently, or (3) never emerges. To separate equation selection from neural optimization, they propose a freeze-then-select pipeline.

Methodology§

The method has two stages:

1. Structured Field Adapter (Freeze)§

The field $u(x,t)$ is factorized as a sum of learned spatial features and temporal coefficients:

$$ u(x,t) = \sum_{k=1}^{K} \phi_k(x) \, c_k(t)$$

where $\phi_k(x)$ is a neural network output and $c_k(t)$ is represented by cubic B-spline coefficients. This architectural bias promotes smoothness in time and spatial locality. Training is purely data-driven (minimizing reconstruction MSE on sparse observations) and does not involve any PDE residual. Once trained, the neural field is frozen.

2. Stability-Validated Weak Selection (SVWS)§

With the field frozen, one can compute spatial and temporal derivatives analytically or via automatic differentiation. The weak form for a candidate PDE term library $\{\xi_i(u)\}$ is constructed by multiplying with test functions $v_j$ and integrating over spacetime. Let $\mathcal{T}_i = \int \xi_i(u)\, v_j \, dx dt$ and $\mathcal{U} = \int u_t\, v_j\,dx dt$. The coefficients $\lambda$ solve:

$$ \sum_i \lambda_i \mathcal{T}_i + \mathcal{U} = 0 $$

SVWS builds several independent weak-form systems by using different test functions and data subdomains. It identifies recurrent terms (those appearing consistently across systems), refits candidate supports to the full system, and then selects the final equation using a held-out weak-form system, thereby adding stability against noise and ill-conditioning. The same principle is applied to non-polynomial expression libraries generated via genetic programming, enabling discovery of e.g. nonlinear diffusion laws.

Implementation Sketch§

The following PyTorch snippet illustrates the structured field adapter and the freeze-then-select workflow:

import torch, torch.nn as nn

class CubicSplineBasis(nn.Module):
    # simplified B-spline basis evaluation
    def forward(self, t):
        # returns basis matrix shape (N, n_basis)
        return basis_fn(t)

class StructuredFieldAdapter(nn.Module):
    def __init__(self, in_dim, out_dim, n_basis):
        super().__init__()
        self.spatial = nn.Sequential(nn.Linear(in_dim, 64), nn.GELU(), nn.Linear(64, out_dim))
        self.spline = CubicSplineBasis()
        self.coeff = nn.Parameter(torch.randn(n_basis, out_dim))

    def forward(self, x, t):
        phi = self.spatial(x)          # (N, out_dim)
        c = self.spline(t) @ self.coeff  # (N, out_dim)
        return (phi * c).sum(dim=-1, keepdim=True)

# Stage 1: train on sparse observations with MSE loss
model = StructuredFieldAdapter(2, 16, 20)
# ... training loop ...

# Stage 2: freeze and build weak-form library
model.eval()
for p in model.parameters():
    p.requires_grad_(False)
# compute weak-form matrices T_i and U, solve with stability-validated selection

Results§

Across all six sparse MDBench regimes, the method achieves the highest exact support recovery rate. The largest margin over classical and neural baselines appears on the chaotic Kuramoto–Sivashinsky equation, where the weak-form SVWS helps avoid spurious terms and transient supports.

Contribution§

This work provides a principled, modular pipeline that decouples field approximation from equation discovery. The stability-validated weak selection is a useful addition to the scientific machine learning toolbox.