Minimalistic_EoS

Code for the paper Understanding Edge-of-Stability Training Dynamics with a Minimalist Example by Xingyu Zhu*, Zixuan Wang*, Xiang Wang, Mo Zhou, and Rong Ge.

Paper: arXiv:2210.03294

Overview

This repository contains a self-contained experiments for studying edge-of-stability (EoS) behavior in minimalist gradient descent examples. The paper studies settings where the loss can oscillate while still decreasing over longer time scales, and where the final sharpness approaches the stability threshold 2 / eta.

The notebooks reproduce trajectory, loss, sharpness, and bifurcation visualizations for the minimalist scalar-network examples.

Contents

• 2num_traj_vis.ipynb: two-parameter trajectory visualization for a fixed learning rate.
• 2num_traj_vis_difflr.ipynb: two-parameter trajectory comparison across learning rates.
• 4num_traj_vis.ipynb: four-parameter trajectory visualization for a fixed learning rate.
• 4num_traj_vis_difflr.ipynb: four-parameter trajectory comparison across learning rates.
• bifurcation_examples.ipynb: bifurcation and polynomial-map visualizations.
• utilsrc/: scalar-network definitions, training utilities, sharpness computations, and plotting helpers.
• utils.py: helper utilities used by the bifurcation notebook.
• algos/: local Hessian/eigenvalue utilities used by the experiments.

Run Jupyter from the repository root so local imports such as utilsrc, utils, and algos resolve correctly:

cd Minimalistic_EoS
mkdir -p figs

The notebooks save generated plots into ./figs/.

Citation

@inproceedings{
  zhu2023understanding,
  title={Understanding Edge-of-Stability Training Dynamics with a Minimalist Example},
  author={Xingyu Zhu and Zixuan Wang and Xiang Wang and Mo Zhou and Rong Ge},
  booktitle={The Eleventh International Conference on Learning Representations },
  year={2023},
  url={https://openreview.net/forum?id=p7EagBsMAEO}
}