{"id":19382894,"url":"https://github.com/locuslab/edge-of-stability","last_synced_at":"2025-04-09T11:11:18.362Z","repository":{"id":50331239,"uuid":"342999869","full_name":"locuslab/edge-of-stability","owner":"locuslab","description":null,"archived":false,"fork":false,"pushed_at":"2024-12-07T04:53:14.000Z","size":533,"stargazers_count":65,"open_issues_count":1,"forks_count":21,"subscribers_count":3,"default_branch":"github","last_synced_at":"2025-04-02T09:07:28.777Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/locuslab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-02-28T01:59:43.000Z","updated_at":"2025-02-23T22:40:35.000Z","dependencies_parsed_at":"2025-01-07T02:32:15.702Z","dependency_job_id":null,"html_url":"https://github.com/locuslab/edge-of-stability","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/locuslab%2Fedge-of-stability","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/locuslab%2Fedge-of-stability/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/locuslab%2Fedge-of-stability/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/locuslab%2Fedge-of-stability/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/locuslab","download_url":"https://codeload.github.com/locuslab/edge-of-stability/tar.gz/refs/heads/github","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248027407,"owners_count":21035594,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-10T09:23:48.479Z","updated_at":"2025-04-09T11:11:18.321Z","avatar_url":"https://github.com/locuslab.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability\n\nThis repository contains source code for the ICLR 2021 paper [Gradient Descent on Neural Networks Typically Occurs\nat the Edge of Stability](https://openreview.net/forum?id=jh-rTtvkGeM) by Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar.\n\n![](https://www.cs.cmu.edu/~jeremiac/gd.gif)\n\nThe structure of this README is:\n1. [Preliminaries](#anchors-in-markdown)\n2. [Quick start](#quick-start)\n3. [Complete documentation](#complete-documentation)\n\n### Preliminaries\n\nTo run the code, you need to set two environment variables:\n1. Set the `DATASETS` environment variable to a directory where datasets will be stored.\n For example: `export DATASET=\"/my/directory/datasets\"`.\n2. Set the `RESULTS` environment variable to a directory where results will be stored.\n For example: `export RESULTS=\"/my/directory/results\"`.\n\n### Quick start\n\nLet's walk through how to use `gd.py` and `flow.py`.\n\n#### Gradient descent\n\nThe script `gd.py` trains a neural network using gradient descent.\nThe required arguments are:\n\n```\ngd.py [dataset] [arch_id] [loss] [lr] [max_steps]\n```\n\nFor example:\n```\npython src/gd.py cifar10-5k fc-tanh  mse  0.01 100000 --acc_goal 0.99 --neigs 2  --eig_freq 100\n```\nThe above command will train a fully-connected tanh network (`fc-tanh`) on a 5k subset of CIFAR-10 (`cifar10-5k`) using the square loss (`mse`).\nWe will run vanila gradient descent with step size 0.01 (`lr`).\nTraining will terminate when either the train accuracy reaches 99% (`train_acc`) or when 100,000 (`max_steps`)\niterations have passed.\nEvery 50 (`eig_freq`) iterations, the top 2 (`neigs`) eigenvalues of the training loss Hessian will be computed and recorded.\nThe training results will be saved in the following output directory:\n```\n${RESULTS}/cifar10-5k/fc-tanh/seed_0/mse/gd/lr_0.01\n```\n(See `get_gd_directory()` in `utilities.py` for the formula by which this output directory is named.)\n\nWithin this output directory, the following files will be created, each containing a PyTorch tensor:\n - `train_loss_final`, `test_loss_final`, `train_acc_final`, `test_acc_final`: the train and test\n losses and accuracies, recorded at each iteration\n - `eigs_final`: the top 2 eigenvalues, measured every 50 (`eig_freq`) iterations.\n\nThe following matplotlib code will plot the train loss, train accuracy, and sharpness.\n```python\nimport torch\nimport matplotlib.pyplot as plt\nfrom os import environ\n\ndataset = \"cifar10-5k\"\narch = \"fc-tanh\"\nloss = \"mse\"\ngd_lr = 0.01\ngd_eig_freq = 50\n\ngd_directory = f\"{environ['RESULTS']}/{dataset}/{arch}/seed_0/{loss}/gd/lr_{gd_lr}\"\n\ngd_train_loss = torch.load(f\"{gd_directory}/train_loss_final\")\ngd_train_acc = torch.load(f\"{gd_directory}/train_acc_final\")\ngd_sharpness = torch.load(f\"{gd_directory}/eigs_final\")[:,0]\n\nplt.figure(figsize=(5, 5), dpi=100)\n\nplt.subplot(3, 1, 1)\nplt.plot(gd_train_loss)\nplt.title(\"train loss\")\n\nplt.subplot(3, 1, 2)\nplt.plot(gd_train_acc)\nplt.title(\"train accuracy\")\n\nplt.subplot(3, 1, 3)\nplt.scatter(torch.arange(len(gd_sharpness)) * gd_eig_freq, gd_sharpness, s=5)\nplt.axhline(2. / gd_lr, linestyle='dotted')\nplt.title(\"sharpness\")\nplt.xlabel(\"iteration\")\n```\n\n![demo](figures/demo1.png)\n\n#### Gradient flow\n\nThe script `flow.py` trains a neural network using gradient flow --- that is, by using the Runge-Kutta\n algorithm to numerically integrate the gradient flow ODE.\n \n The required arguments are:\n ```\nflow.py [dataset] [arch_id] [loss] [tick] [max_time]\n```\nFor example:\n```\npython src/flow.py cifar10-5k fc-tanh mse 1.0  1000 --acc_goal 0.99 --neigs 2  --eig_freq 1\n```\nHere, the argument ``tick = 1.0`` means that the train/test losses and accuracies will be computed and saved after each 1.0 units of time, \nand the argument `max_time  = 1000` means that training will stop after a maximum of 1000 units of time (or until the train accuracy reaches the `acc_goal` of 0.99).\nThe other flags mean the same thing as in the `gd.py` example above.\n\nSee the detailed `flow.py` documentation below for details on how the Runge Kutta step size is set.\n\nThe following matplotlib code will plot the train loss, train accuracy, and sharpness:\n\n```python\nflow_tick = 1.0\nflow_eig_freq = 1\n\nflow_directory = f\"{environ['RESULTS'] }/{dataset}/{arch}/seed_0/{loss}/flow/tick_{flow_tick}\"\n\nflow_train_loss = torch.load(f\"{flow_directory}/train_loss_final\")\nflow_train_acc = torch.load(f\"{flow_directory}/train_acc_final\")\nflow_sharpness = torch.load(f\"{flow_directory}/eigs_final\")[:, 0]\n\nplt.figure(figsize=(5, 5), dpi=100)\n\nplt.subplot(3, 1, 1)\nplt.plot(torch.arange(len(flow_train_loss)) * flow_tick, flow_train_loss)\nplt.title(\"train loss\")\n\nplt.subplot(3, 1, 2)\nplt.plot(torch.arange(len(flow_train_acc)) * flow_tick, flow_train_acc)\nplt.title(\"train accuracy\")\n\nplt.subplot(3, 1, 3)\nplt.scatter(torch.arange(len(flow_sharpness)) * flow_tick * flow_eig_freq, flow_sharpness, s=5)\nplt.title(\"sharpness\")\nplt.xlabel(\"time\")\n```\n\n![demo](figures/demo2.png)\n\n \n \n#### Comparing gradient descent to gradient flow\n\nNow suppose that we'd like to track the discrepancy between the gradient descent trajectory and\nthe gradient flow trajectory.\nThat is, to the extent that gradient descent is a discretization of the gradient flow ODE,\nwe might expect N steps of gradient descent at step size 0.01 to reach the same place as \ngradient flow after \"0.01 N\" units of time. \n\nBy adding the flags `--nproj 500 --iterate_freq 100` to the `src/gd.py` command above, we can \ninstruct `gd.py` to save a 500-dimensional random projection of the iterates (the network weights) every 100 iterations.\nSince the learning rate was 0.01, this amounts to saving a random projection of the iterates every 1.0 units of time.\nA PyTorch tensor containing the random projections of the iterates will be stored in the output directory, under the filename `iterates_final`.\n\nSimilarly, by adding the flags `--nproj 500 --iterate_freq 1` to the `src/flow.py` command, we can\ninstruct `flow.py` to save a 500-dimensional random projection of the iterates every 1 tick (which, in turn, means every 1 unit of time, since `tick`= 1 ).\n \nThus, for both gradient descent and gradient flow, we're saving (random projections of) the iterates every 1.0 units of time.\nTherefore, we can directly compare these saved iterates in order to assess whether gradient descent follows the same\ntrajectory as gradient flow.\n \n The following matplotlib code plots the distance between the gradient descent trajectory and the gradient flow trajectory: \n\n```python\ngd_iterate_freq = 50\nflow_iterate_freq = 1\n\n# the GD iterates are saved every \"gd_lr * gd_iterate_freq\" units of time.\n# the GF iterates are saved every \"flow_tick * flow_iterate_freq\" units of time.\n# to directly compare the trajectories, these two quantities should be equal.\nassert gd_lr * gd_iterate_freq == flow_tick * flow_iterate_freq\n\ngd_iterates = torch.load(f\"{gd_directory}/iterates_final\")\nflow_iterates = torch.load(f\"{flow_directory}/iterates_final\")\nlength = min(len(gd_iterates), len(flow_iterates))\n\ntimes = torch.arange(length) * flow_tick * flow_iterate_freq\ndistance = (gd_iterates[:length, :] - flow_iterates[:length, :]).norm(dim=1)\n\n# the time at which the gradient descent sharpness first crosses the threshold (2  / gd_lr)\ncross_threshold_time = (gd_sharpness \u003e (2 / gd_lr)).nonzero()[0][0] * gd_lr * gd_iterate_freq\n\nplt.figure(figsize=(5, 2), dpi=100)\nplt.scatter(times, distance)\nplt.axvline(cross_threshold_time, linestyle='dotted')\nplt.ylim((0, plt.ylim()[1]))\nplt.title(\"distance betwen gradient descent and gradient flow\")\nplt.xlabel(\"time\")\nplt.ylabel(\"distance\")\n\n```\n![demo](figures/demo3.png)\n\nThe dotted vertical line marks the iteration when the gradient descent sharpness crosses the gradient descent stability threshold 2 / 0.01.\nWe can see that the gradient descent trajectory approximately tracks the gradient flow\n trajectory before this time, but departs immediately afterwards.\n\n\n## Complete documentation\n\n#### Gradient descent\n\nThe script `src/gd.py`trains a network using full-batch gradient descent.\n\nThe required parameters of `src/gd.py` are:\n\n- `dataset` [string]: the dataset to train on. The possible values are:\n    - `cifar10`: the full CIFAR-10 dataset\n    - `cifar10-5k`: the first 5,000 examples from the full CIFAR-10 dataset.  Along similar lines, you can also enter `cifar10-1k`, `cifar10-2k`, `cifar10-10k`, and `cifar10-20k`.\n    - `chebyshev-3-20`: a toy regression dataset comprising 20 points evenly spaced between -1 and 1 labeled noiselessly\n    by the Chebyshev polynomial of degree 3.  `chebyshev-4-20` and `chebyshev-5-20` are similar, but with\n    polynomials of degree 4 and 5. \n    - we will soon add the other datasets from the paper (wikitext, the synthetic datasets) to this repo  \n- `arch_id` [string]: which network architectures to train.  See `load_architecture()` in `archs.py` for a full list of the permissible values.\n- `loss` [string]: which loss function to use.  The possible values are:\n    - `ce`: cross-entropy loss\n    - `mse`: mean squared error (square loss).  When the network has a scalar-valued output (e.g. in regression tasks), this is\n    defined as 0.5 times the squared difference between the network output and the target scalar.\n    More generally, when the network has a vector-valued output (e.g. in multiclass classification), this is defined\n     as 0.5 times the squared L2 norm between the network outputs and the target vector.\n     For classification tasks, we construct the target vector by assigning a 1 at the location of\n     the true class, and assigning 0's everywhere else.\n - `lr` [float]: the learning rate\n - `max_steps` [int]: the maximum number of gradient descent steps to train for.\n\nThe optional parameters of `src/gd.py` are:\n - `seed` [int]: the random seed used when initializing the network weights.\n - `opt` [string, defaults to 'gd']: which (full-batch) gradient descent variant to use.  The options are:\n    - `gd`: vanilla gradient descent.\n    - `polyak`: Polyak-style momentum.\n     If you use this option, make sure to pass in a value for `beta` (the momentum parameter). \n    - `nesterov`: Nesterov-style momentum.\n    If you use this option, make sure to pass in a value for `beta` (the momentum parameter).\n - `beta` [float, defaults to 0.0]: if you use Polyak or Nesterov momentum (i.e. if `gd` = `polyak` or `nesterov`), this is the value of the momentum\n parameter.  This parameter is ignored if you run vanilla gradient descent.\n - `physical_batch_size` [int, defaults to 1000]: the maximum number of examples that we try to fit on the GPU at once.\n For large networks, you may have to reduce this from the default of 1000.\n Note that for non-BN networks, changing this value will have no impact on the result of the training process.\n On the other hand, for BN networks, this value determines the \"ghost batch size.\"\n - `acc_goal` [float between 0 and 1]: if this parameter is set, we terminate training if the train accuracy ever reaches above `acc_goal`.\n - `loss_goal` [float]: if this parameter is set, we terminate training if the train loss ever reaches below `loss_goal`.\n - `neigs` [int, defaults to 0]: the number of top eigenvalues to compute.\n - `eig_freq` [int, defaults to -1]: the frequency at which we compute the top `neigs` Hessian eigenvalues.\nFor example, if `eig_freq` is 10, we compute the leading Hessian eigenvalues every 10 GD iterations.\nIf `eig_freq` is set to `-1` (the default value), we never compute the top Hessian eigenvalues.\n- `nproj`: [int, defaults to 0]: the number of random projections to use when we save random projections of the iterates.\n- `iterate_freq` [int, defaults to -1]: the frequency at which we save random projections of the iterates.\nFor example, if `iterate_freq` is 10, we save random projections of the iterates every 10 GD iterations.\nIf `iterate_freq` is set to -1 (the default value), we never save random projections of the iterates.\n- `save_freq` [int, defaults to -1]: the frequency at which we save the results of training (the train/test losses and \naccuracies, and the computed eigenvalues).  If `save_freq` is set to -1 (the default value), then we only save these\nat the very end of training.\n- `abridged_size` [int, defaults to 5000]: when computing the top Hessian eigenvalues, use an \"abridged\" dataset of this\nsize (obtained by choosing the first `abridged_size` examples in the training dataset.).\n- `save_model` [bool, defaults to False]: if `true`, save the model weights at the end of training.\n\n#### Gradient flow\n\nThe script `src/flow.py` trains a network using gradient flow, i.e. by using the Runge-Kutta \nalgorithm to numerically integrate the gradient flow ODE.\n\n \nThe Runge-Kutta algorithm requires a step size parameter.\nSimilar to gradient descent, for Runge-Kutta to be stable \nthis step size must be small in regions where the sharpness is high.\nWe therefore set the step size using the heuristic\n```\nrunge kutta step size = min(alpha / [estimated sharpness], max_step_size) \n```\nwhere ``--alpha`` and `--max_step_size` are flags that can be passed to `src/flow.py`, \nand `[estimated sharpness]` is the most recent cached value of the sharpness.\nRunge-Kutta is stable when its step size is less than c / sharpness, where c is a \nparticular constant between 2 and 3. \nTherefore, `alpha` is set by default to be 1.0, though a value as high as 2.0 should\nalso be fine in most cases.\nThe `max_step_size` parameter is intended to be an additional stopgap\nto account for the fact that the sharpness may increase rapidly, causing the true sharpness to \nbe a lot higher than the most recently cached sharpness.\n\nThe required parameters of `src/flow.py` are:\n\n- `dataset` [string]: see above\n- `arch_id` [string]: see above\n- `loss` [string]: see above\n- `tick` [float]: the train / test losses and accuracies will be computed and saved every `tick` units of time.\n - `max_time` [float]: the maximum time (ODE time, not wall clock time) to run gradient flow for.\n\nThe optional parameters of `src/flow.py` are:\n - `alpha` [float, defaults to 1.0]: the Runge-Kutta step size is `min(alpha / [estimated sharpness], max_step_size)`\n - `max_step_size` [float, default to 999]: the Runge-Kutta step size is `min(alpha / [estimated sharpness], max_step_size)`\n - `seed` [int]: see above\n - `physical_batch_size` [int, defaults to 1000]: see above\n - `acc_goal` [float between 0 and 1]: see above\n - `loss_goal` [float]: see above\n - `neigs` [int, defaults to 0]: see above\n - `eig_freq` [int, defaults to -1]: see above\n - `nproj`: [int, defaults to 0]: see above\n- `iterate_freq` [int, defaults to -1]: see above\n- `save_freq` [int, defaults to -1]: see above\n- `abridged_size` [int, defaults to 5000]: see above\n- `save_model` [bool, defaults to False]: see above\n\n## Update (December 2024) - Adam\n\nI'm adding code for running \"adaptive edge of stability\" (https://arxiv.org/abs/2207.14484) experiments with Adam.  The code is in the new file `src/adam.py`.\n\nTo train the fully-connected tanh network using Adam with $\\eta$ = 5e-5, $\\beta_1$ = 0.9, $\\beta_2$ = 0.99, $\\epsilon$ = 1e-7 (default), while recording the preconditioned sharpness every five iterations, run the command:\n```\npython src/adam.py cifar10-5k fc-tanh mse 5e-5 20000 --loss_goal 0.05 --neigs 4  --eig_freq 5 --beta1 0.9 --beta2 0.99\n```\nI ran this code for step sizes 5e-5, 1e-4, 2e-4, 4e-4, and then generated the following plot of train loss and preconditioned sharpness:\n![demo](figures/adam.png)\nThe dashed line is the prediction of $(2 + 2 \\beta_1)/((1 - \\beta_1)\\eta )$\n\nTo train the same network using RMSProp, run the above line but pass `--beta1 0.0`, i.e.\n```\npython src/adam.py cifar10-5k fc-tanh mse 5e-5 20000 --loss_goal 0.05 --neigs 4  --eig_freq 5 --beta1 0.0 --beta2 0.99\n```\n\nI ran this code for step sizes 5e-6, 1e-5, 2e-5, 5e-5 and then generated the following plot of train loss and preconditioned sharpness:\n![demo](figures/rmsprop.png)\nThe dashed line is the prediction of $2 / \\eta$.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flocuslab%2Fedge-of-stability","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flocuslab%2Fedge-of-stability","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flocuslab%2Fedge-of-stability/lists"}