Adaptive Sampling
& Training Strategies
for Physics-Informed
Neural Networks

Why Uniform Sampling Breaks Down

Every introductory PINN implementation uses uniform collocation: sample $N_f$ points $(x_j, t_j)$ randomly from the domain according to a uniform distribution, and minimize the PDE residual at those points. For simple, smooth problems, a mild diffusion equation, a slowly-varying ODE, this works adequately. But for the class of problems where physics is most interesting, uniform sampling fails systematically.

Consider the viscous Burgers equation:

with $\nu = 0.01/\pi$ (the benchmark case from Raissi et al. 2019 and from PINNs literature at large). With this small viscosity, the solution develops a steep shock-like gradient near $x \approx 0$ at $t \approx 0.5$–$0.8$. The spatial derivative $|\partial u / \partial x|$ near the shock can be $\mathcal{O}(100)$, while in the smooth regions it is $\mathcal{O}(1)$.

With 10,000 uniform collocation points, approximately 100–200 fall in the shock region (width $\approx \mathcal{O}(\nu) \approx 0.003$). The PDE residual is enormous there and negligible elsewhere. The network spends most of its capacity fitting the 9,800 trivial points and fails to resolve the shock. The result: apparent convergence (the total loss goes down), but catastrophic pointwise error near the shock.

RAR: Sample Where the Error Is

Residual-Based Adaptive Refinement (RAR), introduced by Lu et al. (2021) in the DeepXDE paper [1], addresses spatial stiffness by iteratively concentrating collocation points in regions where the current network violates the PDE most severely. The algorithm is elegantly simple:

The key theoretical motivation for RAR is that the PDE residual is an error indicator: it measures how much the current network approximation fails to satisfy the governing equation. Under mild regularity assumptions, the PDE residual bounds the pointwise approximation error (this is the analog of a posteriori error estimators in finite element methods). Concentrating collocation points where the residual is large therefore steers training toward the regions of highest error, a form of active learning for scientific ML.

def rar_refinement(model, x_f, t_f, n_add=200, n_eval=10000):
    """Add n_add points where PDE residual is largest."""
    # Dense evaluation mesh
    x_eval = torch.linspace(-1, 1, int(np.sqrt(n_eval))).unsqueeze(1)
    t_eval = torch.linspace(0, 1, int(np.sqrt(n_eval))).unsqueeze(1)
    X, T = torch.meshgrid(x_eval.squeeze(), t_eval.squeeze(), indexing='ij')
    x_grid = X.reshape(-1, 1).requires_grad_(True)
    t_grid = T.reshape(-1, 1).requires_grad_(True)

    # Compute PDE residual on the grid
    u = model(x_grid, t_grid)
    u_t = torch.autograd.grad(u.sum(), t_grid, create_graph=False)[0]
    u_x = torch.autograd.grad(u.sum(), x_grid, create_graph=True)[0]
    u_xx = torch.autograd.grad(u_x.sum(), x_grid, create_graph=False)[0]
    residual = (u_t + u * u_x - NU * u_xx).detach()

    # Select top-n_add points by |residual|^2
    err = (residual**2).squeeze()
    _, top_idx = torch.topk(err, n_add)

    x_new = x_grid[top_idx].detach()
    t_new = t_grid[top_idx].detach()

    # Augment collocation set
    x_f_new = torch.cat([x_f, x_new], dim=0)
    t_f_new = torch.cat([t_f, t_new], dim=0)
    return x_f_new, t_f_new

Learnable Point Weights: the Min-Max Formulation

Self-Adaptive PINNs (SA-PINNs), proposed by McClenny and Braga-Neto (2023) [3], take a different approach. Instead of adding or removing points, they assign a learnable weight $\lambda_i$ to each collocation point and train these weights simultaneously with the network parameters, but in the opposite direction.

The standard PINN loss minimization becomes a min-max problem:

where $r_i(\theta)$ is the PDE residual at collocation point $i$, and $\lambda_i \geq 0$ are the adaptive weights. The network parameters $\theta$ minimize the loss (as usual). The weights $\lambda_i$ maximize it, they are trained to assign higher weight to points where the residual is currently larger. This creates a natural feedback loop: points the network is getting wrong automatically receive more training emphasis.

Concretely, the weight update rule uses gradient ascent:

where $T_\lambda$ is a temperature parameter controlling the sharpness of the weight distribution, and the softmax normalization ensures the total weight sums to $N_f$ (so the effective total loss magnitude is preserved). In practice, $\lambda_i$ is implemented as an additional nn.Parameter in the model, with a separate optimizer that takes ascent steps.

NTK Analysis and Loss Weighting

Wang, Teng, and Perdikaris (2021) [2] provided the first principled explanation for a well-known empirical observation: PINNs trained with equal weights on the physics loss and boundary/initial condition losses often fail to satisfy the boundary conditions accurately. The explanation comes from Neural Tangent Kernel (NTK) theory.

Under the NTK regime (infinite-width limit), the training dynamics of the network can be described by the eigendecomposition of the NTK matrix $\mathbf{K}$. The convergence rate for each loss component is governed by the eigenvalue spectrum of the corresponding block of $\mathbf{K}$. If the NTK eigenvalues for the physics loss ($\kappa_f$) and boundary loss ($\kappa_b$) satisfy $\kappa_f \ll \kappa_b$, then the network fits the boundary conditions rapidly while the physics loss barely moves, an imbalance that forces the solution to satisfy the boundaries correctly but violate the PDE interior.

The practical remedy is to weight the losses inversely proportional to their NTK spectral norms at each training step:

where $\mathbf{K}_k$ is the NTK block corresponding to loss component $k$. Computing $\text{tr}(\mathbf{K}_k)$ exactly requires a backward pass per sample, which is expensive. The practical approximation used in the notebook is to compute the gradient magnitude ratio:

and apply an exponential moving average to smooth the weight updates across training steps.

class AdaptiveLossWeighter:
    def __init__(self, alpha=0.9):
        self.alpha = alpha  # EMA smoothing
        self.weights = {'bc': 1.0, 'ic': 1.0, 'physics': 1.0}

    def update(self, model, losses):
        grads = {}
        for key, loss in losses.items():
            loss.backward(retain_graph=True)
            grads[key] = torch.cat([
                p.grad.abs().flatten() for p in model.parameters()
                if p.grad is not None
            ])
            model.zero_grad()

        max_grad = max(g.max().item() for g in grads.values())
        for key, g in grads.items():
            new_w = max_grad / (g.mean().item() + 1e-8)
            # Exponential moving average
            self.weights[key] = (
                self.alpha * self.weights[key] + (1 - self.alpha) * new_w
            )
        return self.weights

Respecting the Arrow of Time

Wang, Sankaran, and Perdikaris (2022) [4] identified a fundamental failure mode specific to time-dependent PINNs: the network simultaneously receives gradient signals from all time points, including $t = T$, before it has correctly learned the solution at $t = 0$. This violates the causal structure of time-dependent PDEs, you cannot correctly satisfy the equation at $t = 0.8$ if you have not correctly satisfied it at $t = 0.2$ first, because the initial condition propagates forward through the equation.

The causal training algorithm addresses this by introducing time-dependent weights that exponentially decay the contribution of late-time residuals:

where $\epsilon > 0$ is a causality parameter. If the residual at earlier times $t_k < t_j$ is large, then $w_j \approx 0$, and the contribution of point $j$ to the gradient is suppressed, the network is told "don't worry about $t_j$ yet, you haven't got $t_k$ right." As earlier time points are correctly learned and their residuals approach zero, $w_j \to 1$, and the gradient from $t_j$ begins to contribute fully.

Expanding Time Windows

Curriculum learning for time-dependent PINNs takes an even more explicit approach to causality: train the network only on the interval $[0, T_1]$ to convergence, then extend to $[0, T_2]$ with $T_2 > T_1$, and continue. The trained solution on the first window serves as a high-quality initial condition (in the sense of network weights) for the next window, a form of warm-starting that exploits the temporal structure of the problem.

For the damped harmonic oscillator $\ddot{x} + 2\gamma\dot{x} + \omega_0^2 x = 0$ with $\gamma = 0.5$, $\omega_0 = 5$ (approximately 8 oscillation periods in $[0, 10]$), a PINN trained uniformly on $[0, 10]$ from scratch typically fails after period 3–4 due to accumulated phase error. Curriculum learning with windows $T_k = 2k$ for $k = 1, \ldots, 5$ recovers the correct solution through all 8 periods:

T_windows = [2.0, 4.0, 6.0, 8.0, 10.0]
model = OscillatorPINN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for T_end in T_windows:
    # Collocation points only in [0, T_end]
    t_f = torch.rand(N_COLLOC, 1) * T_end
    t_ic = torch.zeros(1, 1)  # IC always at t=0

    for epoch in range(EPOCHS_PER_WINDOW):
        optimizer.zero_grad()
        loss_ic = compute_ic_loss(model, t_ic)
        loss_ode = compute_ode_residual(model, t_f)
        loss = loss_ic + loss_ode
        loss.backward()
        optimizer.step()

    print(f"Window [0, {T_end:.1f}] done. Max error: {eval_error(model, T_end):.4f}")

The key implementation detail is that the initial condition loss is always evaluated at $t = 0$, regardless of the current window, this anchors the solution to the true starting point as the window expands. Without this anchor, solutions from earlier windows can drift from the true initial condition when the window is extended.

Overcoming Spectral Bias

Rahaman et al. (2019) [6] proved that neural networks with standard activation functions exhibit a "spectral bias" or "frequency principle": during gradient descent, the network learns low-frequency components of the target function first and high-frequency components last (or not at all). For PINNs applied to wave equations, high-frequency oscillations, or fine-scale spatial structure, this means the network systematically fails to represent the physical solution regardless of depth, width, or training duration.

The Fourier feature embedding (Tancik et al. 2020 [5]) addresses this by explicitly mapping the input coordinates to a high-dimensional space of sinusoidal features before passing them through the network:

where $\mathbf{B} \in \mathbb{R}^{m \times d}$ is a random matrix with entries drawn from $\mathcal{N}(0, \sigma^2)$. The network then operates on $\gamma(\mathbf{x}) \in \mathbb{R}^{2m}$ rather than $\mathbf{x} \in \mathbb{R}^d$. The effect on the NTK of the embedded network is that the resulting kernel approximates a stationary isotropic Gaussian kernel with bandwidth $\sigma$:

This means the embedded network is initialized in a regime where it can represent all spatial frequencies up to $\mathcal{O}(\sigma)$, not just the low-frequency modes. The bandwidth $\sigma$ is the key hyperparameter: too small and the embedding provides no benefit; too large and it introduces aliasing. A practical heuristic from the notebook: set $\sigma$ to the Nyquist frequency of the finest feature you expect in the solution, estimated from the PDE parameters.

class FourierEmbedding(nn.Module):
    def __init__(self, in_dim=2, n_features=128, sigma=10.0):
        super().__init__()
        # Fixed random matrix B — not trained
        B = torch.randn(n_features, in_dim) * sigma
        self.register_buffer('B', B)

    def forward(self, x):
        # x: (batch, in_dim)
        proj = (2 * np.pi * x) @ self.B.T  # (batch, n_features)
        return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
        # output: (batch, 2*n_features)

class FourierPINN(nn.Module):
    def __init__(self, sigma=10.0):
        super().__init__()
        self.embed = FourierEmbedding(in_dim=2, n_features=64, sigma=sigma)
        self.net = nn.Sequential(
            nn.Linear(128, 128), nn.Tanh(),  # 128 = 2*64
            nn.Linear(128, 128), nn.Tanh(),
            nn.Linear(128, 64),  nn.Tanh(),
            nn.Linear(64, 1)
        )

    def forward(self, x, t):
        inp = torch.cat([x, t], dim=1)
        return self.net(self.embed(inp))

When to Use What

A Practical Flowchart for Strategy Selection

Combining strategies

The strategies described here are not mutually exclusive. A common high-performance configuration for challenging PDEs is: Fourier features (to overcome spectral bias) + causal weighting (to enforce temporal causality) + NTK loss balancing (to prevent boundary condition dominance). This combination, sometimes called the "enhanced PINN" configuration in the literature, typically achieves 1–2 orders of magnitude improvement over baseline uniform PINNs on benchmark stiff problems, at a total overhead of $\sim$30–40% additional compute per epoch, a favorable trade given the convergence improvement.

Bibliography