Gradient Flow Matching for Learning Update Dynamics in Neural Network Training

TL;DR

Gradient Flow Matching (GFM) models neural network training as a dynamical system, accurately forecasting weight trajectories and convergence across architectures by learning optimizer-aware vector fields.

Contribution

GFM introduces a continuous-time framework that captures optimizer update rules, enabling accurate extrapolation of training dynamics and convergence prediction.

Findings

GFM achieves competitive forecasting accuracy with Transformer-based models.

GFM outperforms LSTM and classical baselines in training trajectory prediction.

GFM generalizes across different neural architectures and initializations.

Abstract

Training deep neural networks remains computationally intensive due to the itera2 tive nature of gradient-based optimization. We propose Gradient Flow Matching (GFM), a continuous-time modeling framework that treats neural network training as a dynamical system governed by learned optimizer-aware vector fields. By leveraging conditional flow matching, GFM captures the underlying update rules of optimizers such as SGD, Adam, and RMSprop, enabling smooth extrapolation of weight trajectories toward convergence. Unlike black-box sequence models, GFM incorporates structural knowledge of gradient-based updates into the learning objective, facilitating accurate forecasting of final weights from partial training sequences. Empirically, GFM achieves forecasting accuracy that is competitive with Transformer-based models and significantly outperforms LSTM and other classical baselines. Furthermore, GFM generalizes across neural architectures and initializations, providing a unified framework for studying optimization dynamics and accelerating convergence prediction.