add lookahead optimiser

[chengchingwen]

This add the implementation of Lookahead optimiser from this paper and based on the pytorch impl

[DhairyaLGandhi]

Perhaps an AbstractOptimiser type would help with some of the repitition.

[chengchingwen]

But some optimiser don't have momentum term. Maybe we also need a AbstractMomentumOptimiser?

[DhairyaLGandhi]

Thinking if we could use traits here?

[chengchingwen]

https://docs.julialang.org/en/v1/manual/methods/#Trait-based-dispatch-1
if you are talking about this, I guess it's possible.

[chengchingwen]

abstract type OptimiserStyle end
struct Momentumlike <: OptimiserStyle end
struct NonMomentumlike <: OptimiserStyle end

has_momentum(o::Type{Descent}) = NonMomentumlike()
has_momentum(o::Type{Momentum}) = Momentumlike()
has_momentum(o::Type{Adam}) = Momentumlike()

like this?

[DhairyaLGandhi]

Something like that, yeah.

[DhairyaLGandhi]

Also, ref https://github.com/oxinabox/oxinabox.github.io/blob/master/notebook_posts/Dispatch%2C%20Traits%20and%20Metaprogramming%20Over%20Reflection.ipynb

[chengchingwen]

@dhairyagandhi96 I define a AbstractOptimiserand the GradientStyletraits`.

[CarloLucibello]

I would step back to the original problem and just add the Lookahead optimizer

[chengchingwen]

@CarloLucibello Sure! should I revert the commit to 510cfaa or 14956c2 ?

[CarloLucibello]

yes please. If we want to introduce a type hierarchy for the optimizers let's first discuss it in a separate issue/PR

[chengchingwen]

@CarloLucibello Reverted.

[CarloLucibello]

In the paper, I read
"""
We have the choice of maintaining, interpolating, or resetting the internal state (e.g. momentum) of
the inner optimizer. We evaluate this tradeoff on the CIFAR dataset (where every choice improves
convergence) in Appendix D.1 and maintain internal state for the other experiments.
....
D.1 Inner Optimizer State
Throughout our paper, we maintain the state of our inner optimizer for simplicity. For SGD with
heavy-ball momentum, this corresponds to preserving the momentum. Here, we present a sensitivity
study by comparing the convergence of Lookahead when maintaining the momentum, interpolating
the momentum, and resetting the momentum. All three improve convergence versus SGD.
"""
So let's keep this simple and just maintain the state.

This also needs a rebase

[CarloLucibello]

So let's keep this simple and just maintain the state.

I'm saying this because the specific needs of this PR, we probably need a state setter/getter API for the optimizers, so we should think of it more carefully in a separate issue/PR, then we can revisit LookAhead when we have that in place.
So I think it's better to just merge the minimal LookAhead implementation here, with no state reset

[chengchingwen]

@DhairyaLGandhi rebased and updated

[CarloLucibello]

sorry we droppebd the ball here. Closing since optimisers are now in Optimisers.jl