Hi everyone,
I am trying to perform the following experiment and I’d like your advice on what’s the best way to implement it in Pytorch. Given a mini-batch, weight gradients dW^{(t)} are computed based on minimizing a loss function. When we step towards that direction, we get our new weights and can calculate the new loss.

What I’d like to do is based on the weight gradients (in particular, the vector direction) to step at multiple distances from W^(t), each time computing the new loss, i.e. move across the dotted line as if I am experimenting with different learning rates/step sizes and each time calculate the loss as if we were to step that much:

image779×265 51.3 KB

In other words, to call optimizer.step() but then be able to revert bach each time to the initial state, so that I can then call optimizer.step() but with a bigger step size, and to repeat this multiple times. What would the best way be to implement this?
Diagrams taken from the work of Santurkar et al. (https://papers.nips.cc/paper/7515-how-does-batch-normalization-help-optimization.pdf 1).