The NTK spectrum as a proxy-hacking diagnostic
A follow-on to Reward hacking in RLHF and loss imbalance in PINNs are the same optimization pathology. By Robby Sneiderman. ~1100 words.
The earlier note argued that PINN residual-loss imbalance and RLHF reward hacking are the same abstract problem: proxy optimization under an implicit trust region. The fix in both fields turns out to be a ratio-balancing rule — NTK-trace ratio in PINNs, KL budget in RLHF — and algorithmic refinements transfer between the two subfields (causal curriculum, hard-constraint decoding / BC, etc.). The argument was structural. It did not commit to any empirical prediction beyond "they look the same on paper."
The prediction below does commit. It is the kind of thing that is either wrong, in which case I learn something, or right, in which case it connects two literatures that currently pretend not to see each other.
1. The signature
Recall the mechanistic content of Wang-Teng-Perdikaris 2021 (arXiv:2001.04536). They showed that during PINN training, the data-fitting loss neural-tangent kernel K_uu and the PDE-residual loss NTK K_rr have eigenvalue spectra separated by several orders of magnitude, almost always with K_rr spread much wider. The optimizer sees a single combined loss but is effectively driven by the kernel whose spectrum has higher mean, and the other term stalls. The interactive NTK-spectrum-live demo on the PINN page reproduces this in real time.
If the isomorphism with RLHF is mechanical — not just an analogy — the same signature should show up during PPO training of a language model. Specifically:
Prediction. LetK_rbe the empirical NTK of the reward-model gradient signal through the policy (the token-level gradient that PPO applies at each step), andK_klthe NTK of the KL-penalty gradient againstpi_sft. Then during a run where reward hacking begins, the condition number ofK_rshould diverge by one or more orders of magnitude, whileK_kl's stays bounded. The divergence should precede the degradation of eval-set quality by a nontrivial number of update steps.
In words: the spectral gap is the proxy-hacking signature, and it should be detectable from gradient statistics alone, without running the eval set.
2. Why this is tractable to measure
Full NTK matrices are infeasible at language-model scale — the gradient vector is hundreds of billions of parameters long — but the traceis an expectation that can be estimated with Hutchinson's trick: tr(K) ≈ (1/M) Σ_m ⟨∇_θ f(x_m), ∇_θ f(x_m)⟩ averaged over M random prompts. The top eigenvalue can be estimated via a few power-iteration steps on the Gram matrix of a small prompt subset. Both are standard additions to the PPO training loop and contribute negligible compute overhead relative to the reward-model forward pass.
The analogous PINN measurement is exactly what the live NTK panel does: per-sample Gram matrices over a small set of collocation points, eigendecomposed via Jacobi rotation. It is cheap enough to run every 30 training steps in a browser tab.
3. What the prediction actually forbids
Predictions are only informative if they can be wrong. The ways this one can fail:
- Constant spectra under hacking. If reward hacking occurs without any change in NTK eigenvalue spread, the structural isomorphism is wrong; RLHF pathology is either pre-training-inherited or located elsewhere in the loss geometry (maybe the entropy term, maybe value-function contamination).
- Spectra diverge but hacking doesn't follow.If the reward NTK condition number blows up and eval quality is unaffected, the spectral gap is a necessary-but-not-sufficient property. That would be interesting in the other direction — it would mean the proxy-hacking story is richer than the PINN story, and RLHF has something PINNs don't.
- Both diverge together but at the wrong time.If spectra diverge only after eval quality has already degraded, the signature is a lagging indicator, not a leading one. Still useful for attribution but not for interventions.
The case the prediction genuinely commits to is the middle case not on the list above: eigenvalue spread widens several training steps before quality degrades, the timing of the widening is reproducible across seeds, and it is attenuated when an NTK-balanced KL-penalty rescaling is applied online.
4. Connection to the empirical artifact in this repo
The PINN side of the analogy is mostly already exhibited here. The in-browser NTK spectrum demo shows the data-NTK and PDE-residual-NTK eigenvalue spreads in real time during training, with an optional NTK-balanced toggle that visibly closes the gap. A companion ablation documents that a naive grad-norm proxy for NTK balancing does not close the gap at 2,000-iteration scale — and that parameter count does not rescue it either. This is the PINN side of the "proxy optimization is hard" claim, empirically grounded.
The RLHF side is what the prediction above would test. I do not have access to a large-scale PPO training run at frontier-model scale, so I cannot run this experiment. The instruments needed to run it do already exist in every open RLHF training implementation, as side outputs of the gradient computation. If someone with access to a run does measure it, the result either corroborates the isomorphism or falsifies it cleanly. Either way is useful.
5. Why this shape of prediction matters
Most cross-subfield observations in machine learning are safe — they say "these two things look similar" and cannot be wrong because similarity is in the eye of the beholder. The NTK spectral-spread prediction is less safe because it commits to a quantitative statement that is testable with existing instruments. It is the kind of claim a mechanistic interpretability researcher would recognise as "measure this one thing during training and tell me the answer," which is how mechanistic claims are supposed to be structured.
If the prediction holds, the practical consequence is that online NTK-balanced rescaling of the KL penalty would give you a principled β schedule without any eval-set in the loop. The PINN community has already explored this direction; the RLHF community would get a seven-year head start on the tooling by copying it.
If it fails, the failure localizes the disanalogy. The two subfields look structurally identical; if they behave differently under the same spectral stress test, whatever distinguishes them is worth naming. Either outcome is more informative than the current state of the literature, which pretends they do not share a mathematical surface.
Primary sources
- Wang, Teng, Perdikaris. When and why PINNs fail to train: a neural tangent kernel perspective. arXiv:2001.04536 (2021).
- Ouyang et al. Training language models to follow instructions with human feedback. arXiv:2203.02155 (2022).
- Gao, Schulman, Hilton. Scaling laws for reward model overoptimization. arXiv:2210.10760 (2022). The RLHF analog of NTK-imbalance observations.
- Hutchinson. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Communications in Statistics (1989). The trace estimator that makes NTK measurements tractable at scale.