Accelerating Neural Network Training with Augmented Koopman Dynamics

Neural network training can be viewed as a discrete-time dynamical system, suggesting that its optimisation trajectory can be learned and extrapolated to reduce the cost of repeated backpropagation. Motivated by Koopman operator theory, which represents nonlinear dynamics as linear evolution in a lifted space, we propose a training acceleration method based on augmented Koopman dynamics. Instead of approximating the evolution of parameters alone, our method lifts the training state to include both model parameters and optimiser-dependent internal states, enabling the learned operator to better capture the dynamics of modern optimisers with momentum or adaptive gradient statistics. We estimate a strided, multi-step Koopman operator from training snapshots and use matrix-vector multiplications to extrapolate future parameters and optimiser states, effectively bypassing an adaptive number of backpropagation steps. This formulation suppresses high-frequency stochastic fluctuations induced by mini-batch training and reduces the storage required for operator estimation. To improve robustness, we further introduce a safeguarding mechanism that prevents Koopman-extrapolated states from degrading performance relative to the most recent backpropagation iterate. The method can be integrated with standard optimisers, including Adam, AdamW, SGD with momentum, Adadelta, and other commonly used variants. Experiments across multiple optimisers show that the proposed approach significantly reduces both the number of backpropagation steps and the wall-clock training time required to reach a target accuracy, while maintaining, and in many cases improving, final model performance.