MomSPS 1.0.0

Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance

MomSPS

This is the official repository for the code used to run the experiments in the paper that proposed the MomSPS optimizer. The optimizer is implemented in PyTorch.

Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance
D. Oikonomou, N. Loizou

Installation

You can install the package with

pip install momsps

Usage

Import the optimizers in Python with

from momsps import MomSPS
opt = MomSPS(model.parameters(), lr=1)

or

from momsps import MomSPS_smooth
opt = MomSPS_smooth(model.parameters(), lr=1)

Important: you only need to adapt one line of your training script, described below. MomSPS needs access to the value of the (mini-batch) loss in the .step() method. The easiest way to do this is

loss = YOUR_LOSS_FUNCTION()
loss.backward()
opt.step(loss=loss)           # the line you need to change

Citation

If you use MomSPS, please cite

@article{Oikonomou2024,
  author        = {Dimitris Oikonomou and Nicolas Loizou},
  title         = {Stochastic Polyak Step-sizes and Momentum: Convergence Guarantees and Practical Performance},
  year          = {2024},
}