Centred RMSProp (#51)

* Centred RMSProp

* try making this a keyword

* description, show

* add a d to keyword

* fixup

* require recent Zygote

Project.toml

@@ -13,6 +13,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 ChainRulesCore = "1"
 Functors = "0.2.8"
+Zygote = "0.6.40"
 julia = "1.6"
 
 [extras]

src/rules.jl

@@ -84,37 +84,55 @@ function apply!(o::Nesterov, state, x, dx)
 end
 
 """
-    RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η)))
+    RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η)); centred = false)
 
 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.
 
+[Centred RMSProp](http://arxiv.org/abs/1308.08500) is a variant which normalises
+gradients by an estimate their variance, instead of their second moment.
+
 # 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)
+- Keyword `centred` (or `centered`): Indicates whether to use centred variant
+                                     of the algorithm.
 """
 struct RMSProp{T}
   eta::T
   rho::T
   epsilon::T
+  centred::Bool
 end
-RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η))) = RMSProp{typeof(η)}(η, ρ, ϵ)
+RMSProp(η = 1f-3, ρ = 9f-1, ϵ = eps(typeof(η)); centred::Bool = false, centered::Bool = false) =
+  RMSProp{typeof(η)}(η, ρ, ϵ, centred | centered)
 
-init(o::RMSProp, x::AbstractArray) = zero(x)
+init(o::RMSProp, x::AbstractArray) = (zero(x), o.centred ? zero(x) : false)
 
 function apply!(o::RMSProp, state, x, dx)
-  η, ρ, ϵ, acc = o.eta, o.rho, o.epsilon, state
+  η, ρ, ϵ = o.eta, o.rho, o.epsilon
+  quad, lin = state
 
-  @.. acc = ρ * acc + (1 - ρ) * abs2(dx)
-  dx′ = @lazy dx * (η / (sqrt(acc) + ϵ))
+  @.. quad = ρ * quad + (1 - ρ) * abs2(dx)
+  if o.centred
+    @.. lin = ρ * lin + (1 - ρ) * dx
+  end
+  dx′ = @lazy dx * η / (sqrt(quad - abs2(lin)) + ϵ)
 
-  return acc, dx′
+  return (quad, lin), dx′
+end
+
+function Base.show(io::IO, o::RMSProp)
+    show(io, typeof(o))
+    print(io, "(")
+    join(io, [o.eta, o.rho, o.epsilon], ", ")
+    print(io, "; centred = ", o.centred, ")")
 end
 
 """

test/rules.jl

@@ -14,10 +14,13 @@ RULES = [
   OptimiserChain(ClipNorm(), Adam(0.001)),
   OptimiserChain(ClipGrad(0.5), Momentum()),
   OptimiserChain(WeightDecay(), OAdam(), ClipGrad(1)),
+  # Not the default:
+  RMSProp(centred = true),
 ]
 
 name(o) = typeof(o).name.name  # just for printing testset headings
 name(o::OptimiserChain) = join(name.(o.opts), " → ")
+name(o::RMSProp) = o.centred ? "RMSProp(centred = true)" : :RMSProp
 
 LOG = Dict()  # for debugging these testsets, this makes it easy to plot each optimiser's loss