This repository contains code for the paper Studying K-FAC Heuristics by Viewing Adam through a Second-Order Lens, published at ICML 2024.
Our complete development environment under Python 3.10 is specified in
local_requirements.txt
, with a list of top-level requirements given in
Pipfile
. In theory, pipenv install
in a fresh virtual environment will set
everything up; in practice, JAX in particular may need manual intervention
depending on your local CUDA and cuDNN versions.
At the time of writing, we depend on a bugfix to the KFAC-JAX library, which is
specified in kfac_jax.patch
. This can be applied from the project root with
$ patch -p0 -i kfac_jax.patch
Datasets are not bundled with the repository, so before first use they will need
to be downloaded by calling the constructors with download=True
.
Each dataset and algorithm is specified by a YAML configuration file in
configs/
, where AdamQLR_Damped_NoLRClipping.yaml
is the AdamQLR (Tuned)
algorithm described in our paper, and AdamQLR_NoHPO_NoLRClipping.yaml
is the AdamQLR
(Untuned) setting. To perform a single training run, simply pass
the corresponding files to train.py
with the -c
flag, e.g.:
$ python train.py -c ./configs/fashion_mnist.yaml ./configs/AdamQLR_Damped_NoLRClipping.yaml
A complete hyperparameter optimisation routine, including 50 repetitions of the
best hyperparameters found, can be performed by calling
hyperparameter_optimisation.py
with the corresponding configuration files:
$ python hyperparameter_optimisation.py -c ./configs/fashion_mnist.yaml ./configs/AdamQLR_Damped_NoLRClipping.yaml ./configs/ASHA.yaml
This same file also contains helper functions for running sensitivity studies.
Hyperparameter optimisation runs based on overall runtime rather than number of
epochs may be performed by substituting ./configs/ASHA_time_training.yaml
or
./configs/ASHA_time_validation.yaml
in place of ./configs/ASHA.yaml
.
To replicate all our experimental results, the various run_*.sh
scripts may
be useful.
Logs are produced by Tensorboard in a runs/
directory by default; the paths
can be changed with the config/command-line flag --log-root
.
All our experimental plots are produced using paper_plots.py
, though you may
need to update the paths to match your local configuration.