Saved in:
Bibliographic Details
Main Authors: Shou, Xiao, Ding, Yanna, Gao, Jianxi
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2505.20221
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913859705503744
author Shou, Xiao
Ding, Yanna
Gao, Jianxi
author_facet Shou, Xiao
Ding, Yanna
Gao, Jianxi
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_20221
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Gradient Flow Matching for Learning Update Dynamics in Neural Network Training
Shou, Xiao
Ding, Yanna
Gao, Jianxi
Machine Learning
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.
title Gradient Flow Matching for Learning Update Dynamics in Neural Network Training
topic Machine Learning
url https://arxiv.org/abs/2505.20221