Learning to evolve by matching velocities

Learning to evolve by matching velocities. By Nicola Farenga and one other, from the Politecnico di Milano. Published on the thirtieth of October.

Learning how a nonlinear dynamical system evolves in continuous time usually relies on trajectory-based objectives. The learned dynamics are simulated during training, and gradients are backpropagated through the numerical scheme, at a cost that grows with the rollout horizon. The authors show that this approach has two deeper problems besides cost. First, trajectory-based objectives degenerate as the rollout horizon increases, which hinders gradient-based optimization. Second, the numerical scheme biases what is learned: it constrains the spectrum of the learned Jacobian to the scheme's absolute stability region. The solver used for training quietly shapes the model it trains.

The proposed alternative is Velocity Matching, or VM, an objective that is simulation-free. Rather than rolling the learned dynamics forward and comparing trajectories, it targets the velocity directly. It relies on piecewise-linear interpolation and stochastic collocation, and it approximates the velocity population risk to second order in the sampling step. Being simulation-free is the point. Training no longer has to integrate the learned dynamics forward, and it no longer has to backpropagate through a numerical scheme, so the cost that grew with the rollout horizon, and the degeneration that came with long horizons, are both taken out of the objective. The solver that biased the learned Jacobian is no longer inside the training loop either.

Alongside the objective comes a theoretical framework for velocity-based learning. Its stability estimates control trajectory errors in terms of velocity errors, so getting velocities right is enough to keep trajectories right. Those estimates support universal and continuous-time approximation results when the learned velocity field is parameterized by deep neural networks. In other words, the framework says what a neural velocity field can represent, and it says that a small error in velocity cannot quietly grow into an uncontrolled error along the trajectory.

A bias-variance decomposition of the VM objective then yields three further results. Its higher-order approximation of the velocity population risk gives a consistent minimizer. It relates to the residual of an A-stable continuous-stage Runge-Kutta scheme, which connects VM to implicit integration without ever requiring an implicit solve. And it implicitly matches the action of the learned Jacobian to that of the true one along the observed trajectories, without access to the underlying velocity field or its Jacobian.

The experiments cover nonlinear ordinary differential equations and semi-discretized time-dependent partial differential equations. On both kinds of problem, VM improves accuracy over the trajectory-based objectives it was designed to replace. It also avoids the two costs that motivated it: the expense of simulating during training, and the instabilities that come with long rollouts.

The argument closes its own loop. Trajectory objectives were expensive, hard to optimize at long horizons, and biased by their solver. Matching velocities removes the simulation, and the theory shows that this does not mean giving up on trajectories: velocity errors control trajectory errors, the objective is tied to an A-stable implicit scheme without requiring an implicit solve, and the learned Jacobian is matched to the true one along the observed data. Each of the original complaints receives its own answer.