Repository: FluxML/Optimisers.jl
Pull request #16 (2021-02-24 14:51:11)

RFC: Switch interface for higher-order optimizers
In response to [this discussion](https://github.com/FluxML/ML-Coordination-Tracker/discussions/22#discussioncomment-342967), this PR makes a simple change to the order of the arguments in our interface. Functionally, no changes are made.

Higher-order optimizers (e.g. AdaHessian) consume the first and second derivatives as arguments to the rule. This change allows us to support arbitrary number of derivative terms in the future (though it is limited to only the gradient for now).

Summary of changes:
- `apply(o, x, dx, st)` -> `apply(o, st, x, dx)`
- `update(o, x, dx, st)` -> `update(o, st, x, dx)`
- `o(x, dx, st)` -> `o(st, x, dx)`

The idea is that a future second-order rule would define `apply(o, st, x, grad, hessian) = ...`. And the `update` would be changed to support `Vararg`:
```julia
function update(o, state, x::T, x̄s...) where T
  if all(isnothing, x̄s)
    return x, state
  elseif isleaf(x)
    return _update(o, state, x, x̄s...)
  else
    # don't know if this is correct, but it's just a sketch
    x̄s = foreach(x̄ -> functor(typeof(x), x̄)[1], x̄s)
    x, restructure = functor(typeof(x), x)
    xstate = map((state, x, x̄s...) -> update(o, state, x, x̄s...), state, x, x̄s...)
    return restructure(map(first, xstate)), map(x -> x[2], xstate)
  end
end
```

darsnack: Actually, the following passes the tests:
```julia
function _update(o, st, x, x̄s...)
  x̄, st = apply(o, st, x, x̄s...)
  return patch(x, x̄), st
end

function update(o, state, x::T, x̄s...) where T
  if all(isnothing, x̄s)
    return x, state
  elseif isleaf(x)
    return _update(o, state, x, x̄s...)
  else
    x̄s = map(x̄ -> functor(typeof(x), x̄)[1], x̄s)
    x, restructure = functor(typeof(x), x)
    xstate = map((state, x, x̄s...) -> update(o, state, x, x̄s...), state, x, x̄s...)
    return restructure(map(first, xstate)), map(x -> x[2], xstate)
  end
end
```

So we could add support for higher-order now even though there are no higher-order rules defined.
darsnack: One other concern is composition. Here's a sketch for `OptimiserChain`:
```julia
function apply(o::OptimiserChain, states, x, dx, dxs...)
  new_states = similar(states)
  for (i, (opt, state)) in enumerate(zip(o.opts, states))
    dx, new_states[i] = apply(opt, state, x, dx, dxs...)
  end

  return dx, new_states
end
```

As you can see, this forces the use of the first `dx` which represented the update or "patch" to `x`. The higher-order terms don't get transformed. This would exclude zero-th order optimizers from composition, but I don't have a clear picture of what composition even means in that case.

It also assumes that all the composed optimizers are of the same order. This is something I would like to fix. You could do it like `apply(o::Descent, state, x, xs...) = apply(o, state, x, first(xs))`, but I have yet to figure out how to make that generic.
ToucheSir: One crude approach would be to have first-order optimizers take in a vararg `dxs...`, slice out `dx = dxs[1]` and concat that onto the rest when returning. In other words, each optimizer assumes `length(dxs) >= [optimizer order]` and doesn't touch anything higher. There are probably weird interactions that would come out of gradient transforms like clipping for this, but it's the best I can think of at the moment.
darsnack: Yeah I wanted to enable something that allowed writing rules with a fixed number of `dxs` based on order (syntactically the clearest for the rule author). The other thought was that while a rule can accept an arbitrary number of derivative terms, it produces a single transformed `Δx` that gets `patch`ed onto `x`.

If there was an interface for the optimizer to specify its order, then we could handle calling `apply` with the right number of arguments. For example, if we have `order(o::Descent) = 1` and `order(o::AdaHessian) = 2`:
```julia
function apply(o::OptimiserChain, states, x, dx, dxs...)
  new_states = similar(states)
  for (i, (opt, state)) in enumerate(zip(o.opts, states))
    nterms = order(opt)
    dx, new_states[i] = apply(opt, state, x, dx, dxs[1:nterms])
  end

  return dx, new_states
end
```
If we had an abstract type, you could encode the order into the type.
ToucheSir: > Yeah I wanted to enable something that allowed writing rules with a fixed number of dxs based on order (syntactically the clearest for the rule author)

I was also thinking about how he might design the interface for this. Some design ideas below:
```julia
# 0-order
apply(o, states, x) = states # or x, ()

# >= 2-order
function apply(o, states, x, dx, dxs...)
  states, dx = apply(o, states, x, dx)
  return states, (dx, dxs...)
end
```
This removes the boilerplate for 1st-order optimizers. Given a way to query optimizer order, this could be generalized to any order.
darsnack: Okay, so I am pretty happy with where we are heading. I think we've managed to support arbitrary order optimizers with no extra boilerplate, and a rule write can just add more terms onto `apply` if they need them. Here's a crude (i.e. not good) demonstration of `AdaHessian` (I ignore the hessian estimation part of the algorithm...just assume that `D` is correct).
```julia
using Optimisers
using Statistics: mean

function spatial_smoothing(D, b)
    Dout = similar(D)
    d = length(D)
    for i in 0:(d ÷ b - 1)
        for j in 1:b
            Dout[i * b + j] = sum(D[i * b + k] for k in 1:b) / b
        end
    end

    return Dout
end

struct AdaHessian{T, S<:Integer}
    eta::T
    beta::Tuple{T, T}
    power::T
    block::S
    epsilon::T
end
AdaHessian(η = 1.5f-1, β = (9f-1, 9.99f-1), k = 0.5, b = 1, ϵ = eps(typeof(η))) =
    AdaHessian{typeof(η), typeof(b)}(η, β, k, b, ϵ)

Optimisers.init(o::AdaHessian, x::AbstractArray) = (zero(x), zero(x), o.beta)

function Optimisers.apply(o::AdaHessian, state, x, dx, D)
    η, β, k, b, ϵ = o.eta, o.beta, o.power, o.block, o.epsilon
    mt, vt, βt = state

    Ds = spatial_smoothing(D, b)
    @. mt = β[1] * mt + (1 - β[1]) * dx
    @. vt = β[2] * vt + (1 - β[2]) * Ds^2
    dx = @. mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)^k * η

    return (mt, vt, βt .* β), dx
end

w, b = rand(3, 4), rand(3)

o = AdaHessian()
ps = (w, b)
gs = map(x -> rand(size(x)...), ps)
Ds = map(x -> rand(size(x)...), ps)
st = Optimisers.state(o, ps)
st, ps = Optimisers.update(o, st, ps, gs, Ds)
```

It all runs, and I made no changes to the existing rules other than switching the order of the arguments/returns for `apply` and `update`. Pretty nice!
ToucheSir: Bump. I think this should be good to go after a rebase?
darsnack: Thanks for the reminder. I'll probably only get time to work on this toward the end of this week. I think it needs a rebase, and I want to make sure for all the combinations (1st order + higher order / with + without a rule defined) that we hit the right defaults.

Patch:
diff --git a/src/interface.jl b/src/interface.jl
index b59dad5..e659524 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -9,20 +9,23 @@ function state(o, x)
   end
 end
 
-function _update(o, x, x̄, st)
-  x̄, st = apply(o, x, x̄, st)
-  return patch(x, x̄), st
+function _update(o, st, x, x̄s...)
+  st, x̄ = apply(o, st, x, x̄s...)
+  return st, patch(x, x̄)
 end
 
-function update(o, x::T, x̄, state) where T
-  if x̄ === nothing
-    return x, state
+function update(o, state, x::T, x̄s...) where T
+  if all(isnothing, x̄s)
+    return state, x
   elseif isleaf(x)
-    return _update(o, x, x̄, state)
+    return _update(o, state, x, x̄s...)
   else
-    x̄, _  = functor(typeof(x), x̄)
+    x̄s = map(x̄ -> functor(typeof(x), x̄)[1], x̄s)
     x, restructure = functor(typeof(x), x)
-    xstate = map((x, x̄, state) -> update(o, x, x̄, state), x, x̄, state)
-    return restructure(map(first, xstate)), map(x -> x[2], xstate)
+    xstate = map((state, x, x̄s...) -> update(o, state, x, x̄s...), state, x, x̄s...)
+    return map(first, xstate), restructure(map(last, xstate))
   end
 end
+
+# default all rules to first order calls
+apply(o, state, x, dx, dxs...) = apply(o, state, x, dx)
diff --git a/src/rules.jl b/src/rules.jl
index 54c7979..1f97a11 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -15,13 +15,13 @@ Descent() = Descent(1f-1)
 
 init(o::Descent, x::AbstractArray) = nothing
 
-function apply(o::Descent, x, dx, state)
+function apply(o::Descent, state, x, dx)
   η = convert(eltype(dx), o.eta)
   
-  return dx .* η, state
+  return state, dx .* η
 end
 
-(o::Descent)(m, dm, st) = update(o, m, dm, st)
+(o::Descent)(state, m, dm) = update(o, state, m, dm)
 
 """
     Momentum(η = 1f-2, ρ = 9f-1)
@@ -42,14 +42,14 @@ Momentum(η = 1f-2, ρ = 9f-1) = Momentum{typeof(η)}(η, ρ)
 
 init(o::Momentum, x::AbstractArray) = zero(x)
 
-function apply(o::Momentum, x, dx, state)
+function apply(o::Momentum, state, x, dx)
   η, ρ, v = o.eta, o.rho, state
   @. v = ρ * v - η * dx
   
-  return -v, v
+  return v, -v
 end
 
-(o::Momentum)(m, dm, state) = update(o, m, dm, state)
+(o::Momentum)(state, m, dm) = update(o, state, m, dm)
 
 """
     Nesterov(η = 1f-3, ρ = 9f-1)
@@ -70,14 +70,14 @@ Nesterov(η = 1f-3, ρ = 9f-1) = Nesterov{typeof(η)}(η, ρ)
 
 init(o::Nesterov, x::AbstractArray) = zero(x)
 
-(o::Nesterov)(m, dm, state) = update(o, m, dm, state)
+(o::Nesterov)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::Nesterov, x, dx, state)
+function apply(o::Nesterov, state, x, dx)
   η, ρ, v = o.eta, o.rho, state
   d = @. ρ^2 * v - (1+ρ) * η * dx
   @. v = ρ * v - η * dx
   
-  return -d, v
+  return v, -d
 end
 
 """
@@ -105,15 +105,15 @@ RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η))) = RMSProp{typeof(η)}(η, ρ
 
 init(o::RMSProp, x::AbstractArray) = zero(x)
 
-function apply(o::RMSProp, x, dx, state)
+function apply(o::RMSProp, state, x, dx)
   η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state
   @. acc = ρ * acc + (1 - ρ) * dx^2
   dx = @. dx * (η / (sqrt(acc) + ϵ))
   
-  return dx, acc
+  return acc, dx
 end
 
-(o::RMSProp)(m, dm, state) = update(o, m, dm, state)
+(o::RMSProp)(state, m, dm) = update(o, state, m, dm)
 
 """
     ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
@@ -137,9 +137,9 @@ ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = ADAM{typeof(η)}(
 
 init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::ADAM)(m, dm, state) = update(o, m, dm, state)
+(o::ADAM)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::ADAM{T}, x, dx, state) where T
+function apply(o::ADAM{T}, state, x, dx) where T
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, vt, βt = state
 
@@ -147,7 +147,7 @@ function apply(o::ADAM{T}, x, dx, state) where T
   @. vt = β[2] * vt + (one(T) - β[2]) * dx ^ 2
   dx = @. mt / (one(T) - βt[1]) / (sqrt(vt / (one(T) - βt[2])) + ϵ) * η
 
-  return dx, (mt, vt, βt .* β)
+  return (mt, vt, βt .* β), dx
 end
 
 """
@@ -172,9 +172,9 @@ RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = RADAM{typeof(η)}
 
 init(o::RADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, 1)
 
-(o::RADAM)(m, dm, state) = update(o, m, dm, state)
+(o::RADAM)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::RADAM, x, dx, state)
+function apply(o::RADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   ρ∞ = 2/(1-β[2])-1
 
@@ -190,7 +190,7 @@ function apply(o::RADAM, x, dx, state)
     dx = @. mt / (1 - βt[1]) * η
   end
 
-  return dx, (mt, vt, βt .* β, t + 1)
+  return (mt, vt, βt .* β, t + 1), dx
 end
 
 """
@@ -215,9 +215,9 @@ AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaMax{typeof(η
 
 init(o::AdaMax, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::AdaMax)(m, dm, state) = update(o, m, dm, state)
+(o::AdaMax)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::AdaMax, x, dx, state)
+function apply(o::AdaMax, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, ut, βt = state
@@ -226,7 +226,7 @@ function apply(o::AdaMax, x, dx, state)
   @. ut = max(β[2] * ut, abs(dx))
   dx = @. (η/(1 - βt[1])) * mt/(ut + ϵ)
 
-  return dx, (mt, ut, βt .* β)
+  return (mt, ut, βt .* β), dx
 end
 
 """
@@ -252,9 +252,9 @@ OADAM(η = 1f-3, β = (5f-1, 9f-1), ϵ = eps(typeof(η))) = OADAM{typeof(η)}(η
 
 init(o::OADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, zero(x))
 
-(o::OADAM)(m, dm, state) = update(o, m, dm, state)
+(o::OADAM)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::OADAM, x, dx, state)
+function apply(o::OADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, βt, dx_ = state
@@ -265,7 +265,7 @@ function apply(o::OADAM, x, dx, state)
   @. dx_ = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
   dx = @. dx + 2*dx_
 
-  return dx, (mt, vt, βt .* β, dx_)
+  return (mt, vt, βt .* β, dx_), dx
 end
 
 """
@@ -289,16 +289,16 @@ ADAGrad(η = 1f-1, ϵ = eps(typeof(η))) = ADAGrad{typeof(η)}(η, ϵ)
 
 init(o::ADAGrad, x::AbstractArray) = fill!(similar(x), o.epsilon)
 
-(o::ADAGrad)(m, dm, state) = update(o, m, dm, state)
+(o::ADAGrad)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::ADAGrad, x, dx, state)
+function apply(o::ADAGrad, state, x, dx)
   η, ϵ = o.eta, o.epsilon
   acc = state
 
   @. acc += dx^2
   dx = @. dx * η / (sqrt(acc) + ϵ)
 
-  return dx,  acc
+  return acc, dx
 end
 
 """
@@ -321,9 +321,9 @@ ADADelta(ρ = 9f-1, ϵ = eps(typeof(ρ))) = ADADelta{typeof(ρ)}(ρ, ϵ)
 
 init(o::ADADelta, x::AbstractArray) = (zero(x), zero(x))
 
-(o::ADADelta)(m, dm, state) = update(o, m, dm, state)
+(o::ADADelta)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::ADADelta, x, dx, state)
+function apply(o::ADADelta, state, x, dx)
   ρ, ϵ = o.rho, o.epsilon
   acc, Δacc = state
 
@@ -333,7 +333,7 @@ function apply(o::ADADelta, x, dx, state)
   dx = @. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)
   @. Δacc = ρ * Δacc + (1 - ρ) * dx^2
   
-  return dx, (acc, Δacc)
+  return (acc, Δacc), dx
 end
 
 """
@@ -360,9 +360,9 @@ AMSGrad(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AMSGrad{typeof(
 init(o::AMSGrad, x::AbstractArray) =
   (fill!(similar(x), o.epsilon), fill!(similar(x), o.epsilon), fill!(similar(x), o.epsilon))
 
-(o::AMSGrad)(m, dm, state) = update(o, m, dm, state)
+(o::AMSGrad)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::AMSGrad, x, dx, state)
+function apply(o::AMSGrad, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, v̂t = state
@@ -372,7 +372,7 @@ function apply(o::AMSGrad, x, dx, state)
   @. v̂t = max(v̂t, vt)
   dx = @. η * mt / (sqrt(v̂t) + ϵ)
 
-  return dx, (mt, vt, v̂t)
+  return (mt, vt, v̂t), dx
 end
 
 """
@@ -398,9 +398,9 @@ NADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = NADAM{typeof(η)}
 
 init(o::NADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-(o::NADAM)(m, dm, state) = update(o, m, dm, state)
+(o::NADAM)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::NADAM, x, dx, state)
+function apply(o::NADAM, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
 
   mt, vt, βt = state
@@ -410,7 +410,7 @@ function apply(o::NADAM, x, dx, state)
   dx = @. (β[1] * mt / (1 - β[1] * βt[1]) + (1 - β[1]) * dx / (1 - βt[1])) / 
           (sqrt(vt * β[2] / (1 - βt[2])) + ϵ) * η
 
-  return dx, (mt, vt, βt .* β)
+  return (mt, vt, βt .* β), dx
 end
 
 """
@@ -454,9 +454,9 @@ AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaBelief{typ
 
 init(o::AdaBelief, x::AbstractArray) = (zero(x), zero(x))
 
-(o::AdaBelief)(m, dm, state) = update(o, m, dm, state)
+(o::AdaBelief)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::AdaBelief, x, dx, state)
+function apply(o::AdaBelief, state, x, dx)
   η, β, ϵ = o.eta, o.beta, o.epsilon
   mt, st = state
 
@@ -464,7 +464,7 @@ function apply(o::AdaBelief, x, dx, state)
   @. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
   dx = @. η * mt / (sqrt(st) + ϵ)
   
-  return dx, (mt, st)
+  return (mt, st), dx
 end
 
 """
@@ -482,12 +482,12 @@ WeightDecay() = WeightDecay(5f-4)
 
 init(o::WeightDecay, x::AbstractArray) = nothing
 
-(o::WeightDecay)(m, dm, state) = update(o, m, dm, state)
+(o::WeightDecay)(state, m, dm) = update(o, state, m, dm)
 
-function apply(o::WeightDecay, x, dx, state)
+function apply(o::WeightDecay, state, x, dx)
   dx = @. dx + o.wd * x
 
-  return dx, state
+  return state, dx
 end
 
 """
@@ -503,15 +503,15 @@ OptimiserChain(opts...) = OptimiserChain(opts)
 
 init(o::OptimiserChain, x::AbstractArray) = [init(opt, x) for opt in o.opts]
 
-(o::OptimiserChain)(m, dm, state) = update(o, m, dm, state)
+(o::OptimiserChain)(state, m, dms...) = update(o, state, m, dms...)
 
-function apply(o::OptimiserChain, x, dx, states)
+function apply(o::OptimiserChain, states, x, dx, dxs...)
   new_states = similar(states)
   for (i, (opt, state)) in enumerate(zip(o.opts, states))
-    dx, new_states[i] = apply(opt, x, dx, state)
+    new_states[i], dx = apply(opt, state, x, dx, dxs...)
   end
 
-  return dx, new_states
+  return new_states, dx
 end
 
 for Opt in (:Descent, :ADAM, :Momentum, :Nesterov, :RMSProp,


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 3aa848a..db89f4b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,7 +14,7 @@ using Statistics
     l = loss(w, w′)
     for i = 1:10^4
       gs = gradient(x -> loss(x, w′), w)
-      w, st = o(w, gs..., st)
+      st, w = o(st, w, gs...)
     end
     @test loss(w, w′) < 0.01
   end
@@ -30,7 +30,7 @@ end
   for t = 1:10^5
     x = rand(10)
     gs = gradient(w -> loss(x, w, w′), w)
-    w, st = Optimisers.update(opt, w, gs..., st)
+    st, w = Optimisers.update(opt, st, w, gs...)
   end
   @test loss(rand(10, 10), w, w′) < 0.01
 end