GRADSTOP: Early Stopping of Gradient Descent via Posterior Sampling

This repository contains the necessary code to repeat the experiments found in our paper, GRADSTOP: Early Stopping of Gradient Descent via Posterior Sampling.

Dependencies

We advise using a virtual environment when running code in this repository. To install the packages, run the following commands:

virtualenv venv # create a new virtual environment at ./venv
pip3 install -r requirements.txt

However, if you opt not to create a new virtual environment, simply omit the first line.

Alternatively, if you prefer Conda, use

conda env create -f env.yaml

Demonstration of the method example

To recreate the Demonstration figure (Figure 2), run the notebooks/demo_log_regression_for_heart_disease.ipynb notebook.

Reproducing the runs with medical data

To reproduce the runs with medical data (Table 1 in the paper), you can run the experiments/med_data_run.py . To run the experiment for one random seed, use

python3 med_data_run.py $run_id

where $run_id is any integer (random seed).

Dataset acquisition and caching for transfer learning

To download and cache local copies for datasets, start by signing up to OpenML and place your OpenML api key to openml_apikey.txt. Then you can move either directly to performing experiments or pre-download a local cache of the training data by calling

python3 data.py

Note that a local cache will nonetheless be created whenever the datasets are loaded for the first time. This cache will be stored in ./data/.

Transfer learning on multiple datasets

The experiment to carry out model evaluation on multiple datasets and recreate Figure 3 and Figure 4 in the paper (Model evaluation) is contained in experiments/transfer_learning.py. The experiment is designed to be run for one seed at a time using, e.g., a slurm array job, but a single run can be instanced using

python3 transfer_learning.py $run_id

where $run_id is any integer.

Alternatively, to run the batch of ten random initialisations to recreate the dataset in Figure 3, simply call

sbatch scripts/exp_transfer_learning.sh

in a slurm environment. Note that you will need to amend the above shell script to accommodate for differences in your slurm setup, such as partition and queue names and the module loading for Python.

Transfer learning with noisy labels

Similarly, the experiment to carry out model evaluation on multiple datasets with noisy labels in the transfer set and to recreate Figure 5 in the paper (Transfer learning with noisy labels) is contained in experiments/run_with_noisy_labels.py. For one random seed, run

python3 transfer_learning.py $run_id

where $run_id is any integer.

MCMC demonstration with GradSTOP

Demonstration of parameter uncertainty estimates by Theorem 4.7 of the paper and MCMC can be found in notebooks/demo_mcmc_and_gradstop_for_medical_data.ipynb.

Paper citation

@inproceedings{jamshidi2025,
  title = {{GRADSTOP}: Early Stopping of Gradient Descent via Posterior Sampling},
  author = {Jamshidi, Arash and Sepp\"al\"ainen, Lauri and Haitsiukevich, Katsiaryna and Luu, Hoang Phuc Hau and Bj\"orklund, Anton and Puolam\"aki, Kai },
  year = {2025},
  booktitle = {28th European Conference on Artificial Intelligence (ECAI)},
}