PyTorch Implementation of Michael Jordan’s lab's Perturbed SGD?

[RylanSchaeffer]

There’s a line of work out of Michael Jordan’s lab regarding perturbed stochastic gradient descent that allegedly has advantages over SGD:

• Gradient Descent Can Take Exponential Time to Escape Saddle Points
• How to Escape Saddle Points Efficiently
• Stochastic Gradient Descent Escapes Saddle Points Efficiently

Is there an implementation of Perturbed SGD in PyTorch as an optimizer? I looked through the available optimizers and the answer appears to be no.

[mrshenli]

AFAIK, optim.SGD does not support perturbation. @fmassa can you confirm this or is there a simple work around?

[RylanSchaeffer]

I think it'd be pretty straightforward to modify SGD with a add_perturbation parameter. PSGD is almost identical to SGD; the only modification is adding i.i.d. Gaussian noise:

[vishwakftw]

I think adding noise to gradients is simple. A simple apply_ call should be able to iterate through all model parameters and add the noise to the gradients, after which step is called.

I initially proposed adding Noisy SGD, but the proposal was rejected considering its triviality. Ref: #4332

[soumith]

Our general criteria for acceptance:

• The method has to be well-established (we ask people to submit new or low-cited methods to https://github.com/pytorch/contrib )
  For example, if you see Add QHM and QHAdam optimizers #13213 above, we reject those at that point of time, as they were fresh off the press
• The method has to match the paper
• The method has to have a consistent and thoughtful interface
  • Sometimes, leaving a very well oiled and super high-touch API like optim.SGD as-is and putting advanced options, regularization specializations and other stuff into a separate API make sense as well, to de-risk newly introduced bugs to a very small set of people

cc: @vincentqb

[RylanSchaeffer]

@soumith those seem like reasonable criteria. What determines whether a method is well-established?

I checked and the two earlier papers (excluding the recent 2019 preprint) were published at top-tier venues and have been well cited since:

• Gradient Descent Can Take Exponential Time to Escape Saddle Points was published at NeurIPS in 2017 and has 57 citations
• How to Escape Saddle Points Efficiently was published at ICML 2017 and has 213 citations.

[vishwakftw]

The first paper, IIRC, is an analysis based paper where splines with saddles are constructed and shows that GD takes exponential time to escape them.

The second paper proposes the perturbed gradient method. The method itself has a lot of bells and whistles depending on the type of function it is optimizing (Lipschitz constant, Hessian-Lipschitz constant and so on). The noise is not added at every step, but rather occasionally.

For a pure PSGD paper, this would be a good reference: Escaping From Saddle Points – Online Stochastic Gradient for Tensor Decomposition by Rong Ge and others in 2015.