NTK Analysis and Spectral Bias in PINNs¶
This notebook studies one mixed-signal 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 forcing contains one broad mode and one oscillatory correction. That makes it a clean test for spectral bias in a true PINN solve.
In [1]:
Copied!
from pathlib import Path
# Paths are relative to the folder that holds this notebook.
ROOT = Path.cwd()
LOCAL_FIG_DIR = ROOT / "figs"
CHAPTER_FIG_DIR = LOCAL_FIG_DIR
import os
from dataclasses import dataclass
from pathlib import Path
import tempfile
MPLCONFIGDIR = Path(tempfile.gettempdir()) / "sciml-matplotlib"
MPLCONFIGDIR.mkdir(exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
STYLE_PATH = ROOT / "sciml_style.mplstyle" # optional; the book style needs a LaTeX install
if STYLE_PATH.exists():
plt.style.use(STYLE_PATH)
torch.set_num_threads(1)
PINN_FIG_BASENAME = "mixed-frequency-poisson"
AMPLITUDE_FIG_BASENAME = "mixed-frequency-amplitudes"
NTK_SPECTRA_FIG_BASENAME = "mixed-frequency-ntk-spectra"
NTK_BLOCKS_FIG_BASENAME = "mixed-frequency-ntk-blocks"
NTK_DRIFT_FIG_BASENAME = "mixed-frequency-ntk-drift"
@dataclass(frozen=True)
class MixedPoissonProblem:
"""Poisson problem whose exact solution is a mixed signal."""
low_amp: float = 1.0
low_freq: int = 1
high_amp: float = 0.5
high_freq: int = 3
xmin: float = 0.0
xmax: float = 1.0
def exact(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.low_amp * np.sin(self.low_freq * np.pi * x) + self.high_amp * np.sin(
self.high_freq * np.pi * x
)
def low_component(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.low_amp * np.sin(self.low_freq * np.pi * x)
def high_component(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.high_amp * np.sin(self.high_freq * np.pi * x)
def forcing(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
low = self.low_amp * (self.low_freq * np.pi) ** 2 * np.sin(self.low_freq * np.pi * x)
high = self.high_amp * (self.high_freq * np.pi) ** 2 * np.sin(
self.high_freq * np.pi * x
)
return low + high
class HardBCMLP(nn.Module):
"""MLP with an output transform that enforces Dirichlet data exactly."""
def __init__(self, xmin: float, xmax: float, width: int = 64, depth: int = 3):
super().__init__()
self.xmin = float(xmin)
self.xmax = float(xmax)
layers: list[nn.Module] = [nn.Linear(1, width), nn.Tanh()]
for _ in range(depth - 1):
layers += [nn.Linear(width, width), nn.Tanh()]
layers += [nn.Linear(width, 1)]
self.net = nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> torch.Tensor:
s = (x - self.xmin) / (self.xmax - self.xmin)
return s * (1.0 - s) * self.net(s)
def project_amplitude(u: np.ndarray, x: np.ndarray, freq: int) -> float:
basis = np.sin(freq * np.pi * x)
num = np.trapezoid(u * basis, x)
den = np.trapezoid(basis * basis, x)
return float(num / den)
def evaluate_residual(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
xg = x.clone().detach().requires_grad_(True)
u = model(xg)
du = torch.autograd.grad(u, xg, torch.ones_like(u), create_graph=True)[0]
d2u = torch.autograd.grad(du, xg, torch.ones_like(du), create_graph=True)[0]
return -d2u
def evaluate_solution(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
return model(x)
def evaluate_training_state(
model: nn.Module,
problem: MixedPoissonProblem,
x_collocation: torch.Tensor,
f_collocation: torch.Tensor,
x_test_tensor: torch.Tensor,
x_test: np.ndarray,
u_exact: np.ndarray,
) -> dict[str, float]:
residual = evaluate_residual(model, x_collocation)
loss = (residual - f_collocation).pow(2).mean()
with torch.no_grad():
pred = model(x_test_tensor).cpu().numpy().ravel()
a_low = project_amplitude(pred, x_test, problem.low_freq)
a_high = project_amplitude(pred, x_test, problem.high_freq)
return {
"loss": float(loss.item()),
"l2": float(np.sqrt(np.mean((pred - u_exact) ** 2))),
"a_low": a_low,
"a_high": a_high,
"low_ratio": a_low / problem.low_amp,
"high_ratio": a_high / problem.high_amp,
}
def train_mixed_poisson(
problem: MixedPoissonProblem,
*,
num_steps: int = 1000,
lr: float = 5e-4,
n_collocation: int = 160,
width: int = 64,
depth: int = 3,
seed: int = 42,
log_every: int = 10,
verbose: bool = False,
) -> dict[str, object]:
torch.manual_seed(seed)
np.random.seed(seed)
model = HardBCMLP(problem.xmin, problem.xmax, width=width, depth=depth)
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
x_collocation = torch.linspace(
problem.xmin + 1e-3, problem.xmax - 1e-3, n_collocation
).reshape(-1, 1)
f_collocation = torch.tensor(problem.forcing(x_collocation.numpy()), dtype=torch.float32)
x_test = np.linspace(problem.xmin, problem.xmax, 600)
x_test_tensor = torch.tensor(x_test.reshape(-1, 1), dtype=torch.float32)
u_exact = problem.exact(x_test)
history: list[dict[str, float]] = []
snapshots: dict[int, np.ndarray] = {}
state_snapshots: dict[int, dict[str, torch.Tensor]] = {}
snapshot_steps = {0, 100, 500, 1000}
for step in range(num_steps + 1):
if step % log_every == 0 or step in snapshot_steps:
row = evaluate_training_state(
model,
problem,
x_collocation,
f_collocation,
x_test_tensor,
x_test,
u_exact,
)
row["step"] = float(step)
history.append(row)
snapshots[step] = model(x_test_tensor).detach().cpu().numpy().ravel()
if step in snapshot_steps:
state_snapshots[step] = {
name: tensor.detach().cpu().clone()
for name, tensor in model.state_dict().items()
}
if verbose:
print(
f"step={step:5d} loss={row['loss']:.3e} l2={row['l2']:.3e} "
f"low={row['low_ratio']:.3f} high={row['high_ratio']:.3f}",
flush=True,
)
if step == num_steps:
break
optimizer.zero_grad()
residual = evaluate_residual(model, x_collocation)
loss = (residual - f_collocation).pow(2).mean()
loss.backward()
optimizer.step()
return {
"model": model,
"history": history,
"snapshots": snapshots,
"state_snapshots": state_snapshots,
"x_test": x_test,
"u_exact": u_exact,
"problem": problem,
"width": width,
"depth": depth,
}
def state_vector(state: dict[str, torch.Tensor]) -> torch.Tensor:
return torch.cat([tensor.reshape(-1) for tensor in state.values()])
def jacobian_matrix(model: nn.Module, x: torch.Tensor, operator: str) -> np.ndarray:
params = list(model.parameters())
rows: list[torch.Tensor] = []
for i in range(x.shape[0]):
model.zero_grad(set_to_none=True)
xi = x[i : i + 1].clone().detach().requires_grad_(True)
if operator == "solution":
scalar = evaluate_solution(model, xi).reshape(())
elif operator == "residual":
scalar = evaluate_residual(model, xi).reshape(())
else:
raise ValueError(f"unknown operator: {operator}")
scalar.backward()
grads: list[torch.Tensor] = []
for param in params:
if param.grad is None:
grads.append(torch.zeros_like(param).reshape(-1))
else:
grads.append(param.grad.reshape(-1))
rows.append(torch.cat(grads).detach().cpu())
return torch.stack(rows, dim=0).numpy()
def analyze_ntk(
result: dict[str, object],
*,
ntk_points: int = 40,
) -> dict[str, object]:
problem: MixedPoissonProblem = result["problem"] # type: ignore[assignment]
width = int(result["width"])
depth = int(result["depth"])
state_snapshots: dict[int, dict[str, torch.Tensor]] = result["state_snapshots"] # type: ignore[assignment]
snapshot_steps = sorted(state_snapshots)
x_ntk = torch.linspace(problem.xmin + 1e-3, problem.xmax - 1e-3, ntk_points).reshape(-1, 1)
blocks: dict[int, dict[str, np.ndarray]] = {}
spectra: dict[int, dict[str, np.ndarray]] = {}
ntk_change: dict[int, float] = {}
weight_change: dict[int, float] = {}
theta0 = state_vector(state_snapshots[snapshot_steps[0]])
K0: np.ndarray | None = None
for step in snapshot_steps:
model = HardBCMLP(problem.xmin, problem.xmax, width=width, depth=depth)
model.load_state_dict(state_snapshots[step])
J_u = jacobian_matrix(model, x_ntk, "solution")
J_r = jacobian_matrix(model, x_ntk, "residual")
K_uu = J_u @ J_u.T
K_ur = J_u @ J_r.T
K_rr = J_r @ J_r.T
K = np.block([[K_uu, K_ur], [K_ur.T, K_rr]])
blocks[step] = {"K_uu": K_uu, "K_ur": K_ur, "K_rr": K_rr, "K": K}
spectra[step] = {
"K": np.sort(np.linalg.eigvalsh(K))[::-1],
"K_uu": np.sort(np.linalg.eigvalsh(K_uu))[::-1],
"K_rr": np.sort(np.linalg.eigvalsh(K_rr))[::-1],
}
if K0 is None:
K0 = K
ntk_change[step] = float(np.linalg.norm(K - K0) / np.linalg.norm(K0))
theta = state_vector(state_snapshots[step])
weight_change[step] = float(torch.linalg.norm(theta - theta0) / torch.linalg.norm(theta0))
return {
"steps": snapshot_steps,
"x_ntk": x_ntk.cpu().numpy().ravel(),
"blocks": blocks,
"spectra": spectra,
"ntk_change": ntk_change,
"weight_change": weight_change,
}
def plot_solution_snapshots(
result: dict[str, object],
panel_steps: tuple[int, int, int] = (100, 500, 1000),
) -> plt.Figure:
problem: MixedPoissonProblem = result["problem"] # type: ignore[assignment]
snapshots: dict[int, np.ndarray] = result["snapshots"] # type: ignore[assignment]
x: np.ndarray = result["x_test"] # type: ignore[assignment]
u_exact: np.ndarray = result["u_exact"] # type: ignore[assignment]
low = problem.low_component(x)
fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.1), sharey=True)
for ax, step in zip(axes, panel_steps):
pred = snapshots[step]
ax.plot(
x,
u_exact,
color="tab:green",
lw=1.9,
label=rf"$u(x)=\sin(\pi x)+0.5\sin({problem.high_freq}\pi x)$",
)
ax.plot(x, low, color="black", ls="--", lw=1.4, label=r"$\sin(\pi x)$")
ax.plot(
x,
pred,
color="tab:red",
lw=1.4,
marker="o",
ms=2.0,
markevery=18,
label="PINN",
)
ax.set_title(f"Step {step}")
ax.set_xlabel("$x$")
ax.set_xlim(problem.xmin, problem.xmax)
axes[0].set_ylabel("$u(x)$")
handles, labels = axes[-1].get_legend_handles_labels()
axes[-1].legend(handles[:3], labels[:3], fontsize=8, loc="upper right")
return fig
def plot_amplitude_history(result: dict[str, object]) -> plt.Figure:
history: list[dict[str, float]] = result["history"] # type: ignore[assignment]
steps = [row["step"] for row in history]
low = [row["low_ratio"] for row in history]
high = [row["high_ratio"] for row in history]
fig, ax = plt.subplots(figsize=(7.2, 4.2))
ax.plot(steps, low, lw=1.8, label="Low-mode amplitude")
ax.plot(steps, high, lw=1.8, label="High-mode amplitude")
ax.axhline(1.0, color="black", ls="--", lw=1.0)
ax.set_xlabel("Gradient descent steps")
ax.set_ylabel("Recovered amplitude / exact amplitude")
ax.set_ylim(0.0, max(2.2, max(low) * 1.05))
ax.legend(loc="upper right")
return fig
def plot_ntk_spectra(analysis: dict[str, object]) -> plt.Figure:
steps: list[int] = analysis["steps"] # type: ignore[assignment]
spectra: dict[int, dict[str, np.ndarray]] = analysis["spectra"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 3, figsize=(14.2, 4.2), constrained_layout=True)
for step in steps:
axes[0].semilogy(spectra[step]["K"], lw=1.2, label=f"Step {step}")
axes[1].semilogy(spectra[step]["K_uu"], lw=1.2, label=f"Step {step}")
axes[2].semilogy(spectra[step]["K_rr"], lw=1.2, label=f"Step {step}")
axes[0].set_title(r"Full block kernel $K$")
axes[1].set_title(r"Solution block $K_{uu}$")
axes[2].set_title(r"Residual block $K_{rr}$")
for ax in axes:
ax.set_xlabel("Mode index")
ax.set_ylabel("Eigenvalue")
axes[0].legend(loc="upper right", fontsize=8)
return fig
def plot_ntk_blocks(analysis: dict[str, object], step: int = 100) -> plt.Figure:
blocks: dict[int, dict[str, np.ndarray]] = analysis["blocks"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 3, figsize=(13.8, 4.0), constrained_layout=True)
labels = [("K_uu", r"$K_{uu}$"), ("K_ur", r"$K_{ur}$"), ("K_rr", r"$K_{rr}$")]
for ax, (key, title) in zip(axes, labels):
image = ax.imshow(blocks[step][key], cmap="magma", aspect="auto")
ax.set_title(f"{title} at step {step}")
ax.set_xlabel("Point index")
ax.set_ylabel("Point index")
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
return fig
def plot_ntk_drift(analysis: dict[str, object]) -> plt.Figure:
steps: list[int] = analysis["steps"] # type: ignore[assignment]
ntk_change: dict[int, float] = analysis["ntk_change"] # type: ignore[assignment]
weight_change: dict[int, float] = analysis["weight_change"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.1), constrained_layout=True)
axes[0].plot(steps, [ntk_change[step] for step in steps], lw=1.8)
axes[0].set_title(r"$\|K(t)-K(0)\|/\|K(0)\|$")
axes[0].set_xlabel("Gradient descent steps")
axes[0].set_ylabel("Relative NTK change")
axes[1].plot(steps, [weight_change[step] for step in steps], lw=1.8)
axes[1].set_title(r"$\|\theta(t)-\theta(0)\|/\|\theta(0)\|$")
axes[1].set_xlabel("Gradient descent steps")
axes[1].set_ylabel("Relative parameter change")
return fig
def save_figure(fig: plt.Figure, basename: str, *, copy_to_chapter: bool = False) -> Path:
LOCAL_FIG_DIR.mkdir(parents=True, exist_ok=True)
CHAPTER_FIG_DIR.mkdir(parents=True, exist_ok=True)
png = LOCAL_FIG_DIR / f"{basename}.png"
pdf = LOCAL_FIG_DIR / f"{basename}.pdf"
fig.savefig(png, dpi=220, bbox_inches="tight")
fig.savefig(pdf, bbox_inches="tight")
if copy_to_chapter:
fig.savefig(CHAPTER_FIG_DIR / f"{basename}.png", dpi=220, bbox_inches="tight")
fig.savefig(CHAPTER_FIG_DIR / f"{basename}.pdf", bbox_inches="tight")
plt.close(fig)
return png
def history_lookup(result: dict[str, object], step: int) -> dict[str, float]:
history: list[dict[str, float]] = result["history"] # type: ignore[assignment]
return next(row for row in history if int(row["step"]) == step)
from pathlib import Path
# Paths are relative to the folder that holds this notebook.
ROOT = Path.cwd()
LOCAL_FIG_DIR = ROOT / "figs"
CHAPTER_FIG_DIR = LOCAL_FIG_DIR
import os
from dataclasses import dataclass
from pathlib import Path
import tempfile
MPLCONFIGDIR = Path(tempfile.gettempdir()) / "sciml-matplotlib"
MPLCONFIGDIR.mkdir(exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
STYLE_PATH = ROOT / "sciml_style.mplstyle" # optional; the book style needs a LaTeX install
if STYLE_PATH.exists():
plt.style.use(STYLE_PATH)
torch.set_num_threads(1)
PINN_FIG_BASENAME = "mixed-frequency-poisson"
AMPLITUDE_FIG_BASENAME = "mixed-frequency-amplitudes"
NTK_SPECTRA_FIG_BASENAME = "mixed-frequency-ntk-spectra"
NTK_BLOCKS_FIG_BASENAME = "mixed-frequency-ntk-blocks"
NTK_DRIFT_FIG_BASENAME = "mixed-frequency-ntk-drift"
@dataclass(frozen=True)
class MixedPoissonProblem:
"""Poisson problem whose exact solution is a mixed signal."""
low_amp: float = 1.0
low_freq: int = 1
high_amp: float = 0.5
high_freq: int = 3
xmin: float = 0.0
xmax: float = 1.0
def exact(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.low_amp * np.sin(self.low_freq * np.pi * x) + self.high_amp * np.sin(
self.high_freq * np.pi * x
)
def low_component(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.low_amp * np.sin(self.low_freq * np.pi * x)
def high_component(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
return self.high_amp * np.sin(self.high_freq * np.pi * x)
def forcing(self, x: np.ndarray) -> np.ndarray:
x = np.asarray(x)
low = self.low_amp * (self.low_freq * np.pi) ** 2 * np.sin(self.low_freq * np.pi * x)
high = self.high_amp * (self.high_freq * np.pi) ** 2 * np.sin(
self.high_freq * np.pi * x
)
return low + high
class HardBCMLP(nn.Module):
"""MLP with an output transform that enforces Dirichlet data exactly."""
def __init__(self, xmin: float, xmax: float, width: int = 64, depth: int = 3):
super().__init__()
self.xmin = float(xmin)
self.xmax = float(xmax)
layers: list[nn.Module] = [nn.Linear(1, width), nn.Tanh()]
for _ in range(depth - 1):
layers += [nn.Linear(width, width), nn.Tanh()]
layers += [nn.Linear(width, 1)]
self.net = nn.Sequential(*layers)
def forward(self, x: torch.Tensor) -> torch.Tensor:
s = (x - self.xmin) / (self.xmax - self.xmin)
return s * (1.0 - s) * self.net(s)
def project_amplitude(u: np.ndarray, x: np.ndarray, freq: int) -> float:
basis = np.sin(freq * np.pi * x)
num = np.trapezoid(u * basis, x)
den = np.trapezoid(basis * basis, x)
return float(num / den)
def evaluate_residual(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
xg = x.clone().detach().requires_grad_(True)
u = model(xg)
du = torch.autograd.grad(u, xg, torch.ones_like(u), create_graph=True)[0]
d2u = torch.autograd.grad(du, xg, torch.ones_like(du), create_graph=True)[0]
return -d2u
def evaluate_solution(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
return model(x)
def evaluate_training_state(
model: nn.Module,
problem: MixedPoissonProblem,
x_collocation: torch.Tensor,
f_collocation: torch.Tensor,
x_test_tensor: torch.Tensor,
x_test: np.ndarray,
u_exact: np.ndarray,
) -> dict[str, float]:
residual = evaluate_residual(model, x_collocation)
loss = (residual - f_collocation).pow(2).mean()
with torch.no_grad():
pred = model(x_test_tensor).cpu().numpy().ravel()
a_low = project_amplitude(pred, x_test, problem.low_freq)
a_high = project_amplitude(pred, x_test, problem.high_freq)
return {
"loss": float(loss.item()),
"l2": float(np.sqrt(np.mean((pred - u_exact) ** 2))),
"a_low": a_low,
"a_high": a_high,
"low_ratio": a_low / problem.low_amp,
"high_ratio": a_high / problem.high_amp,
}
def train_mixed_poisson(
problem: MixedPoissonProblem,
*,
num_steps: int = 1000,
lr: float = 5e-4,
n_collocation: int = 160,
width: int = 64,
depth: int = 3,
seed: int = 42,
log_every: int = 10,
verbose: bool = False,
) -> dict[str, object]:
torch.manual_seed(seed)
np.random.seed(seed)
model = HardBCMLP(problem.xmin, problem.xmax, width=width, depth=depth)
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
x_collocation = torch.linspace(
problem.xmin + 1e-3, problem.xmax - 1e-3, n_collocation
).reshape(-1, 1)
f_collocation = torch.tensor(problem.forcing(x_collocation.numpy()), dtype=torch.float32)
x_test = np.linspace(problem.xmin, problem.xmax, 600)
x_test_tensor = torch.tensor(x_test.reshape(-1, 1), dtype=torch.float32)
u_exact = problem.exact(x_test)
history: list[dict[str, float]] = []
snapshots: dict[int, np.ndarray] = {}
state_snapshots: dict[int, dict[str, torch.Tensor]] = {}
snapshot_steps = {0, 100, 500, 1000}
for step in range(num_steps + 1):
if step % log_every == 0 or step in snapshot_steps:
row = evaluate_training_state(
model,
problem,
x_collocation,
f_collocation,
x_test_tensor,
x_test,
u_exact,
)
row["step"] = float(step)
history.append(row)
snapshots[step] = model(x_test_tensor).detach().cpu().numpy().ravel()
if step in snapshot_steps:
state_snapshots[step] = {
name: tensor.detach().cpu().clone()
for name, tensor in model.state_dict().items()
}
if verbose:
print(
f"step={step:5d} loss={row['loss']:.3e} l2={row['l2']:.3e} "
f"low={row['low_ratio']:.3f} high={row['high_ratio']:.3f}",
flush=True,
)
if step == num_steps:
break
optimizer.zero_grad()
residual = evaluate_residual(model, x_collocation)
loss = (residual - f_collocation).pow(2).mean()
loss.backward()
optimizer.step()
return {
"model": model,
"history": history,
"snapshots": snapshots,
"state_snapshots": state_snapshots,
"x_test": x_test,
"u_exact": u_exact,
"problem": problem,
"width": width,
"depth": depth,
}
def state_vector(state: dict[str, torch.Tensor]) -> torch.Tensor:
return torch.cat([tensor.reshape(-1) for tensor in state.values()])
def jacobian_matrix(model: nn.Module, x: torch.Tensor, operator: str) -> np.ndarray:
params = list(model.parameters())
rows: list[torch.Tensor] = []
for i in range(x.shape[0]):
model.zero_grad(set_to_none=True)
xi = x[i : i + 1].clone().detach().requires_grad_(True)
if operator == "solution":
scalar = evaluate_solution(model, xi).reshape(())
elif operator == "residual":
scalar = evaluate_residual(model, xi).reshape(())
else:
raise ValueError(f"unknown operator: {operator}")
scalar.backward()
grads: list[torch.Tensor] = []
for param in params:
if param.grad is None:
grads.append(torch.zeros_like(param).reshape(-1))
else:
grads.append(param.grad.reshape(-1))
rows.append(torch.cat(grads).detach().cpu())
return torch.stack(rows, dim=0).numpy()
def analyze_ntk(
result: dict[str, object],
*,
ntk_points: int = 40,
) -> dict[str, object]:
problem: MixedPoissonProblem = result["problem"] # type: ignore[assignment]
width = int(result["width"])
depth = int(result["depth"])
state_snapshots: dict[int, dict[str, torch.Tensor]] = result["state_snapshots"] # type: ignore[assignment]
snapshot_steps = sorted(state_snapshots)
x_ntk = torch.linspace(problem.xmin + 1e-3, problem.xmax - 1e-3, ntk_points).reshape(-1, 1)
blocks: dict[int, dict[str, np.ndarray]] = {}
spectra: dict[int, dict[str, np.ndarray]] = {}
ntk_change: dict[int, float] = {}
weight_change: dict[int, float] = {}
theta0 = state_vector(state_snapshots[snapshot_steps[0]])
K0: np.ndarray | None = None
for step in snapshot_steps:
model = HardBCMLP(problem.xmin, problem.xmax, width=width, depth=depth)
model.load_state_dict(state_snapshots[step])
J_u = jacobian_matrix(model, x_ntk, "solution")
J_r = jacobian_matrix(model, x_ntk, "residual")
K_uu = J_u @ J_u.T
K_ur = J_u @ J_r.T
K_rr = J_r @ J_r.T
K = np.block([[K_uu, K_ur], [K_ur.T, K_rr]])
blocks[step] = {"K_uu": K_uu, "K_ur": K_ur, "K_rr": K_rr, "K": K}
spectra[step] = {
"K": np.sort(np.linalg.eigvalsh(K))[::-1],
"K_uu": np.sort(np.linalg.eigvalsh(K_uu))[::-1],
"K_rr": np.sort(np.linalg.eigvalsh(K_rr))[::-1],
}
if K0 is None:
K0 = K
ntk_change[step] = float(np.linalg.norm(K - K0) / np.linalg.norm(K0))
theta = state_vector(state_snapshots[step])
weight_change[step] = float(torch.linalg.norm(theta - theta0) / torch.linalg.norm(theta0))
return {
"steps": snapshot_steps,
"x_ntk": x_ntk.cpu().numpy().ravel(),
"blocks": blocks,
"spectra": spectra,
"ntk_change": ntk_change,
"weight_change": weight_change,
}
def plot_solution_snapshots(
result: dict[str, object],
panel_steps: tuple[int, int, int] = (100, 500, 1000),
) -> plt.Figure:
problem: MixedPoissonProblem = result["problem"] # type: ignore[assignment]
snapshots: dict[int, np.ndarray] = result["snapshots"] # type: ignore[assignment]
x: np.ndarray = result["x_test"] # type: ignore[assignment]
u_exact: np.ndarray = result["u_exact"] # type: ignore[assignment]
low = problem.low_component(x)
fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.1), sharey=True)
for ax, step in zip(axes, panel_steps):
pred = snapshots[step]
ax.plot(
x,
u_exact,
color="tab:green",
lw=1.9,
label=rf"$u(x)=\sin(\pi x)+0.5\sin({problem.high_freq}\pi x)$",
)
ax.plot(x, low, color="black", ls="--", lw=1.4, label=r"$\sin(\pi x)$")
ax.plot(
x,
pred,
color="tab:red",
lw=1.4,
marker="o",
ms=2.0,
markevery=18,
label="PINN",
)
ax.set_title(f"Step {step}")
ax.set_xlabel("$x$")
ax.set_xlim(problem.xmin, problem.xmax)
axes[0].set_ylabel("$u(x)$")
handles, labels = axes[-1].get_legend_handles_labels()
axes[-1].legend(handles[:3], labels[:3], fontsize=8, loc="upper right")
return fig
def plot_amplitude_history(result: dict[str, object]) -> plt.Figure:
history: list[dict[str, float]] = result["history"] # type: ignore[assignment]
steps = [row["step"] for row in history]
low = [row["low_ratio"] for row in history]
high = [row["high_ratio"] for row in history]
fig, ax = plt.subplots(figsize=(7.2, 4.2))
ax.plot(steps, low, lw=1.8, label="Low-mode amplitude")
ax.plot(steps, high, lw=1.8, label="High-mode amplitude")
ax.axhline(1.0, color="black", ls="--", lw=1.0)
ax.set_xlabel("Gradient descent steps")
ax.set_ylabel("Recovered amplitude / exact amplitude")
ax.set_ylim(0.0, max(2.2, max(low) * 1.05))
ax.legend(loc="upper right")
return fig
def plot_ntk_spectra(analysis: dict[str, object]) -> plt.Figure:
steps: list[int] = analysis["steps"] # type: ignore[assignment]
spectra: dict[int, dict[str, np.ndarray]] = analysis["spectra"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 3, figsize=(14.2, 4.2), constrained_layout=True)
for step in steps:
axes[0].semilogy(spectra[step]["K"], lw=1.2, label=f"Step {step}")
axes[1].semilogy(spectra[step]["K_uu"], lw=1.2, label=f"Step {step}")
axes[2].semilogy(spectra[step]["K_rr"], lw=1.2, label=f"Step {step}")
axes[0].set_title(r"Full block kernel $K$")
axes[1].set_title(r"Solution block $K_{uu}$")
axes[2].set_title(r"Residual block $K_{rr}$")
for ax in axes:
ax.set_xlabel("Mode index")
ax.set_ylabel("Eigenvalue")
axes[0].legend(loc="upper right", fontsize=8)
return fig
def plot_ntk_blocks(analysis: dict[str, object], step: int = 100) -> plt.Figure:
blocks: dict[int, dict[str, np.ndarray]] = analysis["blocks"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 3, figsize=(13.8, 4.0), constrained_layout=True)
labels = [("K_uu", r"$K_{uu}$"), ("K_ur", r"$K_{ur}$"), ("K_rr", r"$K_{rr}$")]
for ax, (key, title) in zip(axes, labels):
image = ax.imshow(blocks[step][key], cmap="magma", aspect="auto")
ax.set_title(f"{title} at step {step}")
ax.set_xlabel("Point index")
ax.set_ylabel("Point index")
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
return fig
def plot_ntk_drift(analysis: dict[str, object]) -> plt.Figure:
steps: list[int] = analysis["steps"] # type: ignore[assignment]
ntk_change: dict[int, float] = analysis["ntk_change"] # type: ignore[assignment]
weight_change: dict[int, float] = analysis["weight_change"] # type: ignore[assignment]
fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.1), constrained_layout=True)
axes[0].plot(steps, [ntk_change[step] for step in steps], lw=1.8)
axes[0].set_title(r"$\|K(t)-K(0)\|/\|K(0)\|$")
axes[0].set_xlabel("Gradient descent steps")
axes[0].set_ylabel("Relative NTK change")
axes[1].plot(steps, [weight_change[step] for step in steps], lw=1.8)
axes[1].set_title(r"$\|\theta(t)-\theta(0)\|/\|\theta(0)\|$")
axes[1].set_xlabel("Gradient descent steps")
axes[1].set_ylabel("Relative parameter change")
return fig
def save_figure(fig: plt.Figure, basename: str, *, copy_to_chapter: bool = False) -> Path:
LOCAL_FIG_DIR.mkdir(parents=True, exist_ok=True)
CHAPTER_FIG_DIR.mkdir(parents=True, exist_ok=True)
png = LOCAL_FIG_DIR / f"{basename}.png"
pdf = LOCAL_FIG_DIR / f"{basename}.pdf"
fig.savefig(png, dpi=220, bbox_inches="tight")
fig.savefig(pdf, bbox_inches="tight")
if copy_to_chapter:
fig.savefig(CHAPTER_FIG_DIR / f"{basename}.png", dpi=220, bbox_inches="tight")
fig.savefig(CHAPTER_FIG_DIR / f"{basename}.pdf", bbox_inches="tight")
plt.close(fig)
return png
def history_lookup(result: dict[str, object], step: int) -> dict[str, float]:
history: list[dict[str, float]] = result["history"] # type: ignore[assignment]
return next(row for row in history if int(row["step"]) == step)
Matplotlib is building the font cache; this may take a moment.
Train the mixed-signal Poisson PINN and analyze the NTK¶
The run uses full-batch gradient descent so that the NTK linearization is the natural reference model.
In [2]:
Copied!
problem = MixedPoissonProblem()
result = train_mixed_poisson(problem, verbose=False)
analysis = analyze_ntk(result)
for step in (100, 500, 1000):
row = history_lookup(result, step)
print(
f"step={step:4d} loss={row['loss']:.3e} "
f"low={row['low_ratio']:.3f} high={row['high_ratio']:.3f}"
)
problem = MixedPoissonProblem()
result = train_mixed_poisson(problem, verbose=False)
analysis = analyze_ntk(result)
for step in (100, 500, 1000):
row = history_lookup(result, step)
print(
f"step={step:4d} loss={row['loss']:.3e} "
f"low={row['low_ratio']:.3f} high={row['high_ratio']:.3f}"
)
step= 100 loss=7.815e+02 low=1.971 high=0.149 step= 500 loss=4.577e+01 low=0.945 high=0.908 step=1000 loss=9.066e+00 low=0.981 high=0.995
Solution snapshots and amplitude history¶
The first plot shows what gradient descent recovers at several stages. The second plot tracks the low- and high-mode amplitudes directly.
In [3]:
Copied!
plot_solution_snapshots(result)
plot_amplitude_history(result);
plot_solution_snapshots(result)
plot_amplitude_history(result);
Block NTK spectra and block structure¶
The solution values and residual values have different Jacobians. Their block kernels show which directions are learned quickly and which remain slow.
In [4]:
Copied!
plot_ntk_spectra(analysis)
plot_ntk_blocks(analysis, step=100);
plot_ntk_spectra(analysis)
plot_ntk_blocks(analysis, step=100);
How far does the kernel move?¶
The frozen-kernel picture is only local. These curves show how far the NTK and the parameter vector move from initialization.
In [5]:
Copied!
plot_ntk_drift(analysis);
plot_ntk_drift(analysis);
