Repository: FluxML/Optimisers.jl
Pull request #9 (2021-01-29 00:25:41)

Move all optimizers to Optimisers.jl
This PR finishes the already good work in this repo to move all the optimizers from Flux to Optimisers.jl. The main advantage of the implementation in Optimisers.jl is the use of functors and explicit state. This eliminates the need for any `IdDict`s and simplifies how compositions of basic optimizers work.

~~One notable exception: this implementation is mutable. The simple reason for this is that immutable updates will copy entire models and result in high memory usage. Until we have a compiler optimization to address this, this PR opts to have mutable updates for now. You can read more about the reasoning [in this design note](https://github.com/FluxML/ML-Coordination-Tracker/discussions/22).~~

Remaining tasks:
- [x] docstrings
- [x] clean up tests

cc @DhairyaLGandhi @CarloLucibello @mcabbott @ToucheSir 


Patch:
diff --git a/src/Optimisers.jl b/src/Optimisers.jl
index d72f235..c57aa5a 100644
--- a/src/Optimisers.jl
+++ b/src/Optimisers.jl
@@ -5,7 +5,8 @@ using Functors: functor, fmap, isleaf
 include("interface.jl")
 include("rules.jl")
 
-export Descent, Momentum, Nesterov, RMSProp,
-       ADAM
+export Descent, ADAM, Momentum, Nesterov, RMSProp,
+       ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW, RADAM, OADAM, AdaBelief,
+       WeightDecay, OptimiserChain
 
 end # module
diff --git a/src/interface.jl b/src/interface.jl
index f04727d..b59dad5 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -1,13 +1,11 @@
-function patch(x, x̄)
-  return x .- x̄
-end
+patch(x, x̄) = x .- x̄
 
 function state(o, x)
   if isleaf(x)
     return init(o, x)
   else
     x, _ = functor(x)
-    map(x -> state(o, x), x)
+    return map(x -> state(o, x), x)
   end
 end
 
@@ -20,11 +18,11 @@ function update(o, x::T, x̄, state) where T
   if x̄ === nothing
     return x, state
   elseif isleaf(x)
-    _update(o, x, x̄, state)
+    return _update(o, x, x̄, state)
   else
     x̄, _  = functor(typeof(x), x̄)
-    x, re = functor(typeof(x), x)
+    x, restructure = functor(typeof(x), x)
     xstate = map((x, x̄, state) -> update(o, x, x̄, state), x, x̄, state)
-    re(map(first, xstate)), map(x -> x[2], xstate)
+    return restructure(map(first, xstate)), map(x -> x[2], xstate)
   end
 end
diff --git a/src/rules.jl b/src/rules.jl
index e8e0c64..9c89ea2 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -1,106 +1,516 @@
 """
-    Descent(η)
+    Descent(η = 1f-1)
 
 Classic gradient descent optimiser with learning rate `η`.
-For each parameter `p` and its gradient `p̄`, this runs `p -= η*p̄`.
+For each parameter `p` and its gradient `dp`, this runs `p -= η*dp`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
 """
-mutable struct Descent
-  eta::Float64
+struct Descent{T}
+  eta::T
 end
+Descent() = Descent(1f-1)
 
-init(o::Descent, x) = nothing
+init(o::Descent, x::AbstractArray) = nothing
 
-function apply(o::Descent, x, x̄, state)
-  η = convert(eltype(x̄), o.eta)
-  x̄ .* η, state
+function apply(o::Descent, x, dx, state)
+  η = convert(eltype(dx), o.eta)
+  dx .*= η
+  
+  return dx .* η, state
 end
 
-function (o::Descent)(m, m̄)
-  update(o, m, m̄, state(o, m))[1]
+(o::Descent)(m, dm, st) = update(o, m, dm, st)
+
+"""
+    Momentum(η = 1f-2, ρ = 9f-1)
+
+Gradient descent optimizer with learning rate `η` and momentum `ρ`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Momentum (`ρ`): Controls the acceleration of gradient descent in the
+                  prominent direction, in effect dampening oscillations.
+"""
+struct Momentum{T}
+  eta::T
+  rho::T
 end
+Momentum(η = 1f-2, ρ = 9f-1) = Momentum{typeof(η)}(η, ρ)
+
+init(o::Momentum, x::AbstractArray) = zero(x)
 
-function (o::Descent)(m, m̄, st)
-  update(o, m, m̄, st)
+function apply(o::Momentum, x, dx, state)
+  η, ρ, v = o.eta, o.rho, state
+  @. v = ρ * v - η * dx
+  
+  return -v, v
 end
 
-struct Momentum{T,S}
+(o::Momentum)(m, dm, state) = update(o, m, dm, state)
+
+"""
+    Nesterov(η = 1f-3, ρ = 9f-1)
+
+Gradient descent optimizer with learning rate `η` and Nesterov momentum `ρ`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Nesterov momentum (`ρ`): Controls the acceleration of gradient descent in the
+                           prominent direction, in effect dampening oscillations.
+"""
+struct Nesterov{T}
   eta::T
-  rho::S
+  rho::T
+end
+Nesterov(η = 1f-3, ρ = 9f-1) = Nesterov{typeof(η)}(η, ρ)
+
+init(o::Nesterov, x::AbstractArray) = zero(x)
+
+(o::Nesterov)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::Nesterov, x, dx, state)
+  η, ρ, v = o.eta, o.rho, state
+  d = @. ρ^2 * v - (1+ρ) * η * dx
+  @. v = ρ * v - η * dx
+  
+  return -d, v
 end
 
-function apply(o::Momentum, x, Δ, st)
-  η, ρ = o.eta, o.rho
-  v = st
-  v = @. ρ * v - η * Δ
-  Δ = @. -v
-  Δ, v
+"""
+    RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η)))
+
+Optimizer using the
+[RMSProp](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
+algorithm. Often a good choice for recurrent networks. Parameters other than learning rate
+generally don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Momentum (`ρ`): Controls the acceleration of gradient descent in the
+                  prominent direction, in effect dampening oscillations.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct RMSProp{T}
+  eta::T
+  rho::T
+  epsilon::T
 end
+RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η))) = RMSProp{typeof(η)}(η, ρ, ϵ)
+
+init(o::RMSProp, x::AbstractArray) = zero(x)
 
-function (o::Momentum)(m, m̄, state)
-  update(o, m, m̄, state)
+function apply(o::RMSProp, x, dx, state)
+  η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state
+  @. acc = ρ * acc + (1 - ρ) * dx^2
+  dx = @. dx * (η / (sqrt(acc) + ϵ))
+  
+  return dx, acc
 end
 
-init(o::Momentum, x::AbstractArray) = zero(x)
+(o::RMSProp)(m, dm, state) = update(o, m, dm, state)
 
-struct Nesterov{T,S}
+"""
+    ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
+
+[ADAM](https://arxiv.org/abs/1412.6980) optimiser.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct ADAM{T}
   eta::T
-  rho::S
+  beta::Tuple{T, T}
+  epsilon::T
 end
+ADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = ADAM{typeof(η)}(η, β, ϵ)
 
-init(o::Nesterov, x::AbstractArray) = zero(x)
+init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
 
-function (o::Nesterov)(m, m̄, state)
-  update(o, m, m̄, state)
+(o::ADAM)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::ADAM{T}, x, dx, state) where T
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+  mt, vt, βt = state
+
+  @. mt = β[1] * mt + (one(T) - β[1]) * dx
+  @. 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 .* β)
 end
 
-function apply(o::Nesterov, x, Δ, st)
-  η, ρ = o.eta, o.rho
-  v = st
-  d = @. ρ^2 * v - (1+ρ) * η * Δ
-  v = @. ρ*v - η*Δ
-  Δ = -d
-  Δ, v
+"""
+    RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
+
+[Rectified ADAM](https://arxiv.org/abs/1908.03265) optimizer.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct RADAM{T}
+  eta::T
+  beta::Tuple{T, T}
+  epsilon::T
 end
+RADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = RADAM{typeof(η)}(η, β, ϵ)
+
+init(o::RADAM, x::AbstractArray) = (zero(x), zero(x), o.beta, 1)
 
-struct RMSProp{T,S}
+(o::RADAM)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::RADAM, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+  ρ∞ = 2/(1-β[2])-1
+
+  mt, vt, βt, t = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. vt = β[2] * vt + (1 - β[2]) * dx^2
+  ρ = ρ∞ - 2*t * βt[2] / (1 - βt[2])
+  if ρ > 4
+    r = sqrt((ρ - 4) * (ρ - 2) * ρ∞/((ρ∞ - 4) * (ρ∞ - 2) * ρ))
+    dx = @. mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ) * η * r
+  else
+    dx = @. mt / (1 - βt[1]) * η
+  end
+
+  return dx, (mt, vt, βt .* β, t + 1)
+end
+
+"""
+    AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
+
+[AdaMax](https://arxiv.org/abs/1412.6980) is a variant of ADAM based on the ∞-norm.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct AdaMax{T}
   eta::T
-  rho::S
+  beta::Tuple{T, T}
+  epsilon::T
 end
+AdaMax(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaMax{typeof(η)}(η, β, ϵ)
 
-init(o::RMSProp, x::AbstractArray) = zero(x)
+init(o::AdaMax, x::AbstractArray) = (zero(x), zero(x), o.beta)
+
+(o::AdaMax)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::AdaMax, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+
+  mt, ut, βt = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. ut = max(β[2] * ut, abs(dx))
+  dx = @. (η/(1 - βt[1])) * mt/(ut + ϵ)
+
+  return dx, (mt, ut, βt .* β)
+end
+
+"""
+    OADAM(η = 1f-3, β = (5f-1, 9f-1), ϵ = eps(typeof(η)))
+
+[OADAM](https://arxiv.org/abs/1711.00141) (Optimistic ADAM)
+is a variant of ADAM adding an "optimistic" term suitable for adversarial training.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct OADAM{T}
+  eta::T
+  beta::Tuple{T, T}
+  epsilon::T
+end
+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)
+
+function apply(o::OADAM, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+
+  mt, vt, βt, dx_ = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. vt = β[2] * vt + (1 - β[2]) * dx^2
+  @. dx = -dx_
+  @. dx_ = η * mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
+  dx = @. dx + 2*dx_
+
+  return dx, (mt, vt, βt .* β, dx_)
+end
+
+"""
+    ADAGrad(η = 1f-1, ϵ = eps(typeof(η)))
+
+[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimizer. It has
+parameter specific learning rates based on how frequently it is updated.
+Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct ADAGrad{T}
+  eta::T
+  epsilon::T
+end
+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)
+
+function apply(o::ADAGrad, x, dx, state)
+  η, ϵ = o.eta, o.epsilon
+  acc = state
+
+  @. acc += dx^2
+  dx = @. dx * η / (sqrt(acc) + ϵ)
 
-function apply(o::RMSProp, x, Δ, st)
-  η, ρ = o.eta, o.rho
-  acc = st
-  acc = ρ .* acc .+ (1 .- ρ) .* Δ.^2
-  Δ = Δ .* (η ./ (.√acc .+ ϵ))
-  Δ, acc
+  return dx,  acc
 end
 
-function (o::RMSProp)(m, m̄, state)
-  update(o, m, m̄, state)
+"""
+    ADADelta(ρ = 9f-1, ϵ = eps(typeof(ρ)))
+
+[ADADelta](https://arxiv.org/abs/1212.5701) is a version of ADAGrad adapting its learning
+rate based on a window of past gradient updates.
+Parameters don't need tuning.
+
+# Parameters
+- Rho (`ρ`): Factor by which the gradient is decayed at each time step.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct ADADelta{T}
+  rho::T
+  epsilon::T
+end
+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)
+
+function apply(o::ADADelta, x, dx, state)
+  ρ, ϵ = o.rho, o.epsilon
+  acc, Δacc = state
+
+  @. acc = ρ * acc + (1 - ρ) * dx^2
+  # DON'T remove epsilon from numerator
+  # or even out of the square roots
+  dx = @. dx * sqrt(Δacc + ϵ) / sqrt(acc + ϵ)
+  @. Δacc = ρ * Δacc + (1 - ρ) * dx^2
+  
+  return dx, (acc, Δacc)
 end
 
-struct ADAM{T,K}
+"""
+    AMSGrad(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
+
+The [AMSGrad](https://openreview.net/forum?id=ryQu7f-RZ) version of the ADAM
+optimiser. Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct AMSGrad{T}
   eta::T
-  beta::Tuple{K,K}
+  beta::Tuple{T, T}
+  epsilon::T
 end
+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))
 
-const ϵ = 1f-8
+(o::AMSGrad)(m, dm, state) = update(o, m, dm, state)
 
-function (o::ADAM)(m, m̄, state)
-  update(o, m, m̄, state)
+function apply(o::AMSGrad, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+
+  mt, vt, v̂t = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. vt = β[2] * vt + (1 - β[2]) * dx ^ 2
+  @. v̂t = max(v̂t, vt)
+  dx = @. η * mt / (sqrt(v̂t) + ϵ)
+
+  return dx, (mt, vt, v̂t)
 end
 
-init(o::ADAM, x::AbstractArray) = (zero(x), zero(x), o.beta)
-init(o::ADAM, x) = nothing
+"""
+    NADAM(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
 
-function apply(o::ADAM, x, Δ, st)
-  η, β = o.eta, o.beta
-  mt, vt, βp = st
-  mt = β[1] .* mt .+ (1f0 .- β[1]) .* Δ
-  vt = β[2] .* vt .+ (1f0 .- β[2]) .* Δ .^ 2
-  Δ =  mt ./ (1 .- βp[1]) ./ (.√(vt ./ (1f0 .- βp[2])) .+ ϵ) .* η
-  return Δ, (mt, vt, βp .* β)
+[NADAM](https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ) is a Nesterov variant of ADAM.
+Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+struct NADAM{T}
+  eta::T
+  beta::Tuple{T, T}
+  epsilon::T
 end
+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)
+
+function apply(o::NADAM, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+
+  mt, vt, βt = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. vt = β[2] * vt + (1 - β[2]) * dx^2
+  dx = @. (β[1] * mt / (1 - β[1] * βt[1]) + (1 - β[1]) * dx / (1 - βt[1])) / 
+          (sqrt(vt * β[2] / (1 - βt[2])) + ϵ) * η
+
+  return dx, (mt, vt, βt .* β)
+end
+
+"""
+    ADAMW(η = 1f-3, β = (9f-1, 9.99f-1), γ = 0, ϵ = eps(typeof(η)))
+
+[ADAMW](https://arxiv.org/abs/1711.05101) is a variant of ADAM fixing (as in repairing) its
+weight decay regularization.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Weight decay (`γ`): Decay applied to weights during optimisation.
+- Machine epsilon (`ϵ`): Constant to prevent division by zero
+                         (no need to change default)
+"""
+ADAMW(η = 1f-3, β = (9f-1, 9.99f-1), γ = 0, ϵ = eps(typeof(η))) =
+  OptimiserChain(ADAM{typeof(η)}(η, β, ϵ), WeightDecay(γ))
+
+"""
+    AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η)))
+
+The [AdaBelief](https://arxiv.org/abs/2010.07468) optimiser is a variant of the well-known
+ADAM optimiser.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- Machine epsilon (`ϵ::Float32`): Constant to prevent division by zero
+                                  (no need to change default)
+"""
+struct AdaBelief{T}
+  eta::T
+  beta::Tuple{T, T}
+  epsilon::T
+end
+AdaBelief(η = 1f-3, β = (9f-1, 9.99f-1), ϵ = eps(typeof(η))) = AdaBelief{typeof(η)}(η, β, ϵ)
+
+init(o::AdaBelief, x::AbstractArray) = (zero(x), zero(x))
+
+(o::AdaBelief)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::AdaBelief, x, dx, state)
+  η, β, ϵ = o.eta, o.beta, o.epsilon
+  mt, st = state
+
+  @. mt = β[1] * mt + (1 - β[1]) * dx
+  @. st = β[2] * st + (1 - β[2]) * (dx - mt)^2
+  dx = @. η * mt / (sqrt(st) + ϵ)
+  
+  return dx, (mt, st)
+end
+
+"""
+    WeightDecay(γ = 5f-4)
+
+Decay weights by `γ`.
+
+# Parameters
+- Weight decay (`γ`): Decay applied to weights during optimisation.
+"""
+struct WeightDecay{T}
+  wd::T
+end
+WeightDecay() = WeightDecay(5f-4)
+
+init(o::WeightDecay, x::AbstractArray) = nothing
+
+(o::WeightDecay)(m, dm, state) = update(o, m, dm, state)
+
+function apply(o::WeightDecay, x, dx, state)
+  dx = @. dx + o.wd * x
+
+  return dx, state
+end
+
+"""
+    OptimiserChain(opts...)
+
+Compose a chain (sequence) of optimisers so that each `opt` in `opts`
+updates the gradient in the order specified.
+"""
+struct OptimiserChain{O}
+  opts::O
+end
+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)
+
+function apply(o::OptimiserChain, x, dx, states)
+  new_states = similar(states)
+  for (i, (opt, state)) in enumerate(zip(o.opts, states))
+    dx, new_states[i] = apply(opt, x, dx, state)
+  end
 
+  return dx, new_states
+end
\ No newline at end of file


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 8dc040a..482a656 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,18 +4,33 @@ using Statistics
 
 @testset "Optimisers" begin
   Random.seed!(84)
-  w′ = rand(3,3)
-  @testset for o in (Descent(0.1), Momentum(0.01, 0.9), Nesterov(0.001, 0.9), RMSProp(0.001, 0.9),
-                     ADAM(0.001, (0.9, 0.99)))
-    w = rand(3,3)
-    st = Optimisers.init(o,w)
-    loss(x, y) = mean((x .- y) .^ 2)
+  w′ = (α = rand(3, 3), β = rand(3, 3))
+  @testset for o in (Descent(), ADAM(), Momentum(), Nesterov(), RMSProp(),
+                     ADAGrad(), AdaMax(), ADADelta(), AMSGrad(), NADAM(),
+                     ADAMW(), RADAM(), OADAM(), AdaBelief())
+    w = (α = rand(3, 3), β = rand(3, 3))
+    st = Optimisers.state(o, w)
+    loss(x, y) = mean((x.α .* x.β .- y.α .* y.β) .^ 2)
     l = loss(w, w′)
     for i = 1:10^4
-      gs = gradient(x -> loss(x,w′), w)
+      gs = gradient(x -> loss(x, w′), w)
       w, st = o(w, gs..., st)
     end
     @test loss(w, w′) < 0.01
   end
-
 end
+
+@testset "OptimiserChain" begin
+  Random.seed!(84)
+  w = randn(10, 10)
+  w′ = randn(10, 10)
+  loss(x, w, w′) = mean((w*x .- w′*x) .^ 2)
+  opt = OptimiserChain(WeightDecay(), ADAM(0.001))
+  st = Optimisers.state(opt, w)
+  for t = 1:10^5
+    x = rand(10)
+    gs = gradient(w -> loss(x, w, w′), w)
+    w, st = Optimisers.update(opt, w, gs..., st)
+  end
+  @test loss(rand(10, 10), w, w′) < 0.01
+end
\ No newline at end of file