NTK Analysis and Spectral Bias in PINNs¶
This notebook studies one Poisson problem from beginning to end. The exact solution is $u(x)=\sin(\pi x)+0.5\sin(3\pi x)$ on $[0,1]$. The boundary values vanish at both endpoints. Differentiating twice gives $-u''(x)=\pi^2\sin(\pi x)+4.5\pi^2\sin(3\pi x)$. The signal has one broad component and one oscillatory correction. That makes it a clean test for spectral bias in a true PINN solve.
1. A mixed signal as a Poisson problem¶
The figure below shows three snapshots from full-batch gradient descent. At step 100, the low mode already has order-one amplitude while the high mode is still small. The recovered high-mode amplitude is only 0.15 of its exact value. By step 500 it reaches 0.91. By step 1000 both modes are essentially correct.
from rewrite_ntk_spectral_bias import (
MixedPoissonProblem,
analyze_ntk,
plot_amplitude_history,
plot_ntk_blocks,
plot_ntk_drift,
plot_ntk_spectra,
plot_solution_snapshots,
train_mixed_poisson,
)
problem = MixedPoissonProblem()
result = train_mixed_poisson(problem, verbose=False)
analysis = analyze_ntk(result)
2. Amplitude history¶
A direct way to see spectral bias is to track the Fourier amplitudes during training. The low mode reaches order one first. The high mode lags behind.
plot_amplitude_history(result);
3. NTK intuition¶
At the training points, a parameter update changes the solution values through the Jacobian $J_u$. The residual values have their own Jacobian $J_r$. These produce the block kernels $K_{uu}=J_uJ_u^\top$, $K_{ur}=J_uJ_r^\top$, and $K_{rr}=J_rJ_r^\top$. The spectra below show how the kernel acts on the solution side and on the residual side.
plot_ntk_spectra(analysis)
plot_ntk_blocks(analysis, step=100);
4. How far does the kernel move¶
The Poisson notebook in PINNsNTK also tracks how far the kernel and the parameters move from initialization.
That tells us whether the frozen-kernel picture is plausible over the time window we are studying.
plot_ntk_drift(analysis);
5. Reproduce the main figure¶
The call below redraws the three solution snapshots used in the chapter.
plot_solution_snapshots(result);