PINN Conditioning and Training Dynamics
PINNs fail in characteristic ways - spectral bias and unbalanced gradients; the neural-tangent-kernel view explains why and guides fixes.
Why PINNs are hard to train
PINN losses are notoriously stiff. Two pathologies dominate: spectral bias, where the network learns low-frequency structure long before sharp gradients like the edge pedestal; and gradient imbalance, where the PDE and boundary terms pull the shared weights at very different magnitudes. Understanding these prevents shipping a confidently wrong equilibrium into the twin.
Neural tangent kernel (NTK) view:
d(residual)/dt ~ -K * residual
K : NTK matrix; its eigenvalues set convergence rates per mode
Large spread in eig(K) -> stiff training, slow high-freq modes
Spectral bias: high-frequency error decays slowly
Cure directions: feature scaling, Fourier features, weightingDiagnostics
The stack monitors the eigenvalue spread of the NTK (or a cheap proxy: per-term gradient-norm ratios) during training. A widening spread signals the boundary and residual terms decoupling, which is when adaptive weighting and learning-rate rebalancing are applied. Loss curves alone hide these failures, so gradient diagnostics are logged.
# gradient-imbalance diagnostic
gn_pde = grad_norm(L_pde, theta)
gn_bc = grad_norm(L_bc, theta)
ratio = gn_pde / (gn_bc + eps)
if ratio > hi or ratio < lo:
rebalance_weights() # restore comparable gradient scalesArchitectural remedies
Fourier feature embeddings inject high-frequency basis functions so the network escapes spectral bias and can resolve steep edge gradients. Input normalization keeps R, Z, psi order-one. Adaptive activation and residual connections improve the loss landscape. These are applied deliberately for the breeder's steep-pedestal and the burner's sharp axial potential gradients.
- Fourier features: overcome spectral bias for sharp gradients.
- Input/output normalization: keep loss terms comparably scaled.
- Adaptive weighting: hold PDE and boundary gradients in balance.
- NTK/gradient diagnostics: detect stiffness before it corrupts output.
A PINN that will not condition is not forced into service; the stack falls back to the FEM solver or a reduced-order surrogate rather than trust an ill-trained network in the control twin.