Decoding Hidden Layer Dynamics: Linear Probes as a Drop-in Upgrade for Neural ODE-SINDy Equation Discovery
Abstract
Sparse Identification of Nonlinear Dynamics (SINDy) is a popular method for equation discovery, which is often combined together with surrogate modelling for noisy trajectory data, such as Neural Ordinary Differential Equation (Neural ODE). However, equation discovery from the surrogate's output field is shown to be inefficient because of its late generalization. Therefore, instead of reading the field off the output layer, we read it off a suitably chosen hidden layer of the same network, using a linear probe. In order to automatically determine the checkpoint where the probe should be applied, an interpretability-based label-free law is also utilized. Considering four benchmark systems: simple harmonic oscillator, Duffing oscillator, Van der Pol and Lorenz systems, the proposed method is shown to significantly improve discovery accuracy as well as training speed, at a matched training budget with noisy observation data. We also explain why this approach works. The extractable hidden layer compresses toward the system's phase-space dimension, its decodability is largely architectural rather than learned, and a layer-wise analysis and probe-target ablation pin down which layer to read. Because these pipelines train the surrogate regardless, applying the probe makes it a drop-in upgrade to the already existing Neural-ODE discovery pipelines, for improving their efficiency.