Welcome to the repository of our NeurIPS 2022 paper "Chaotic Dynamics are Intrinsic to Neural Network Training with SGD"

This repository is structured in two subfolders: A folder o2grad, containing the O2Grad package and a folder experiments, containing the files used to run the experiments in the respective section of the paper.

To reproduce the experiments, you will need to go through the following 3 steps:

Recreate the Python virtual environment by entering in a terminal from inside the repository root folder:

> python -m venv .chaotic_nips22
> pip install -r requirements.txt

If you prefer conda for package management, we highly recommend you use mamba to install the dependencies. Install by running

> conda install mamba -n base -c conda-forge
> mamba env create -f environment.yml

Setup the o2grad package by running from a terminal inside the repository root folder:

> cd o2grad
> python setup.py develop

Here is a list of the main dependencies:

pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
pip install hydra-core --upgrade
pip install neptune-client
pip install pytorch-lightning
pip install matplotlib

Be sure to choose the right cudatoolkit for Pytorch for the cuda version on the GPU if you plan to run experiments on GPU.

Finally, you will need to modify train.py in experiments/ to suit your needs. We provide two logging options: Neptune.ai (recommended) or Tensorboard. Assuming you have an account on neptune.ai, the simplest solution would be to replace the following lines in train.py:

#L108-109 in train.py:

# api_key='XXXX',
project='username/project_name',

with your own Neptune API key and your own project name to log your experiments there.

Start an experiment by running train.py with the following options (as defined in experiments/config):

By default, the experiments will run with regular SGD. To rerun the pruning experiments, you need to use the CGD optimizer. As an example, if you want to use pruning of the top magnitude components of the negative spectrum for the MLP model and log metrics related to the Lyapunov matrix, you could do this as

train.py model.name=mlp optim=cgd optim.k=100 optim.prune_set=neg logging.log_lyapunov=True

For a full overview of the parameters that can be adapted in the CGD optimizer, see experiments/config/optim/cgd.yaml.

For the distance experiments you need to run a single training with SGD with Lyapunov calculation first:

train.py save_eigenthings='True' seed=42

And then load the respective eigenvectors file that was created after training in './storage/..'. Define perturbation magnitudes and strategies as follows (perturb_strat='all' for all different strategies at once). Use the same seed!

train.py seed=42 model=mlp dataset=usps logging.log_eigenthings=False logging.log_lyapunov=False logging.lyapunov_perturb_eigs='path/to/save' perturb_eps=[1e-1] perturb_strat='all' parallel=5

This will train multiple models in parallel thus calculating the Lyapunov matrix here is not recommended.

For a the subspace similarity experiment with a CNN without momentum on USPS, start the run with:

train.py model.name='cnn' logging.log_lyapunov_subsim=True logging.t_lya_subsim=[10,1e+4]

Figure 2a from the paper.

Below, we provide an experiment table with detailed commands for running the experiments used to generate all figures used in the main paper: