Unifying Central Flow, Rod Flow, and the Edge Coupling

Abstract:At the edge of stability, the iterates of gradient descent oscillate along the sharp directions of the Hessian. This poses a significant challenge for continuous-time models as the usual continuous-time approximation -- the gradient flow -- admits no oscillations. Several continuous-time models have been proposed to capture the behavior at the edge of stability, among them central flow and rod flow. Although both models can be interpreted as performing a coarse-graining in time, the connection between them remains poorly understood. Here we show that both flows can be derived from a third object, the edge coupling: a functional on consecutive iterates whose critical points are the fixed points and period-two orbits of gradient descent. It is revealed that the two flows correspond to coarse-grainings at different time scales. Rod flow coarse-grains over only the fastest time scale, while central flow coarse-grains over both the fast and intermediate time scales. Furthermore, using the edge coupling, we derive rod flows and central flows for momentum methods such as heavy ball and Adam. We verify the derived flows experimentally on representative neural network architectures, where they track the discrete iterates throughout all phases of training -- including the transition from the stable regime to the edge of stability.

Machine Learning (stat.ML); Machine Learning (cs.LG)

arXiv:2611.08358 [stat.ML]

(or arXiv:2611.08358v1 [stat.ML] for this version)

https://doi.org/10.48550/arXiv.2611.08358 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

unifying_flows/LICENSES/Apache-2.0.txt

unifying_flows/LICENSES/MIT-minGPT.txt

unifying_flows/README.md

unifying_flows/notebooks/experiment_adam_cnn_cifar10.ipynb

unifying_flows/notebooks/experiment_adam_mlp_cifar10.ipynb

unifying_flows/notebooks/experiment_adam_mlp_cifar10_rodavg.ipynb

unifying_flows/notebooks/experiment_adam_vit_cifar10.ipynb

unifying_flows/notebooks/experiment_gd_cnn_cifar10.ipynb

unifying_flows/notebooks/experiment_gd_mlp_cifar10.ipynb

unifying_flows/notebooks/experiment_gd_mlp_cifar10_covsm.ipynb

unifying_flows/notebooks/experiment_gd_mlp_cifar10_floor.ipynb

unifying_flows/notebooks/experiment_gd_mlp_cifar10_rodavg.ipynb

unifying_flows/notebooks/experiment_gd_mlp_cifar10_selfstab.ipynb

unifying_flows/notebooks/experiment_gd_vit_cifar10.ipynb

unifying_flows/notebooks/experiment_hb_cnn_cifar10.ipynb

unifying_flows/notebooks/experiment_hb_mlp_cifar10.ipynb

unifying_flows/notebooks/experiment_hb_mlp_cifar10_rodavg.ipynb

unifying_flows/notebooks/experiment_hb_vit_cifar10.ipynb

unifying_flows/notebooks/experiment_nesterov_cnn_cifar10.ipynb

unifying_flows/notebooks/experiment_nesterov_mlp_cifar10.ipynb

unifying_flows/notebooks/experiment_nesterov_mlp_cifar10_rodavg.ipynb

unifying_flows/notebooks/experiment_nesterov_vit_cifar10.ipynb

unifying_flows/notebooks/experiment_rmsprop_cnn_cifar10.ipynb

unifying_flows/notebooks/experiment_rmsprop_mlp_cifar10.ipynb

unifying_flows/notebooks/experiment_rmsprop_mlp_cifar10_rodavg.ipynb

unifying_flows/notebooks/experiment_rmsprop_vit_cifar10.ipynb

unifying_flows/params/adam_cnn.yaml

unifying_flows/params/adam_mlp.yaml

unifying_flows/params/adam_rodavg_k2.yaml

unifying_flows/params/adam_rodavg_k3.yaml

unifying_flows/params/adam_two_cycle.yaml

unifying_flows/params/adam_vit.yaml

unifying_flows/params/adamw_cnn.yaml

unifying_flows/params/adamw_mlp.yaml

unifying_flows/params/adamw_rodavg_k2.yaml

unifying_flows/params/adamw_rodavg_k3.yaml

unifying_flows/params/adamw_two_cycle.yaml

unifying_flows/params/adamw_vit.yaml

unifying_flows/params/gd_cnn.yaml

unifying_flows/params/gd_covsm.yaml

unifying_flows/params/gd_floor.yaml

unifying_flows/params/gd_flow_dtype.yaml

unifying_flows/params/gd_mlp.yaml

unifying_flows/params/gd_mlp_long.yaml

unifying_flows/params/gd_mlp_xent.yaml

unifying_flows/params/gd_rodavg_k2.yaml

unifying_flows/params/gd_rodavg_k3.yaml

unifying_flows/params/gd_selfstab_eta0.05.yaml

unifying_flows/params/gd_selfstab_eta0.075.yaml

unifying_flows/params/gd_vit.yaml

unifying_flows/params/hb_cnn.yaml

unifying_flows/params/hb_mlp.yaml

unifying_flows/params/hb_rodavg_k2.yaml

unifying_flows/params/hb_rodavg_k3.yaml

unifying_flows/params/hb_two_cycle.yaml

unifying_flows/params/hb_vit.yaml

unifying_flows/params/nesterov_cnn.yaml

unifying_flows/params/nesterov_mlp.yaml

unifying_flows/params/nesterov_rodavg_k2.yaml

unifying_flows/params/nesterov_rodavg_k3.yaml

unifying_flows/params/nesterov_two_cycle.yaml

unifying_flows/params/nesterov_vit.yaml

unifying_flows/params/rmsprop_cnn.yaml

unifying_flows/params/rmsprop_mlp.yaml

unifying_flows/params/rmsprop_mlp_xent.yaml

unifying_flows/params/rmsprop_rodavg_k2.yaml

unifying_flows/params/rmsprop_rodavg_k3.yaml

unifying_flows/params/rmsprop_vit.yaml

unifying_flows/requirements.txt

unifying_flows/src/architectures/__init__.py

unifying_flows/src/architectures/utils.py

unifying_flows/src/datasets.py

unifying_flows/src/eig_solvers.py

unifying_flows/src/functional.py

unifying_flows/src/lobpcg.py

unifying_flows/src/loss_criterion.py

unifying_flows/src/loss_function.py

unifying_flows/src/sdcp.py

unifying_flows/src/update_rules.py

unifying_flows/src/utils.py

(77 additional files not shown) You must enabled JavaScript to view entire file list.

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.