Repository: FluxML/Optimisers.jl (julia)

Issue:
Implementation of `AdamW` differs from PyTorch
Hi, thank you for developing and maintaining this awesome library and ecosystem!

I'm not entirely sure but could it be that the documentation for the `AdamW` optimizer is a bit misleading? If I understand correctly, then [its definition](https://github.com/FluxML/Flux.jl/blob/master/src/optimise/optimisers.jl#L502) of

```
AdamW(η = 0.001, β = (0.9, 0.999), decay = 0) = Optimiser(Adam(η, β), WeightDecay(decay))
```

means that it performs this update (where $-\eta A$ is Adam's update):

$$
\begin{align*}
\theta_t \leftarrow \theta_{t-1} - \eta A + \texttt{decay} \  \theta_{t-1}
\end{align*}
$$

However, the paper on AdamW (which is linked to by the docs) parametrizes this differently as:

$$
\begin{align*}
\theta_t \leftarrow \theta_{t-1} - \eta (\alpha A + \lambda \theta_{t-1})
\end{align*}
$$

I.e. Flux's `eta` corresponds to the paper's $\eta\alpha$ and Flux's `decay` corresponds to the paper's $\eta \lambda$.

This is probably super unimportant (in that case, sorry for the noise) but since I just noticed this during bug hunting in an implementation of mine (which uses AdamW), I thought I'd report it.

Pull request description:
See https://github.com/FluxML/Flux.jl/issues/2433 for the details

Close https://github.com/FluxML/Flux.jl/issues/2433

Patch:
diff --git a/Project.toml b/Project.toml
index 41c9709c..0a19f49c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Optimisers"
 uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
 authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
-version = "0.3.4"
+version = "0.4.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
diff --git a/README.md b/README.md
index 6aecb4b3..73770a08 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,12 @@ This was written as the new training system for [Flux.jl](https://github.com/Flu
 and also used by [Lux.jl](https://github.com/avik-pal/Lux.jl).
 But it can be used separately on any array, or anything else understood by [Functors.jl](https://github.com/FluxML/Functors.jl).
 
+
+> [!WARNING]
+> With version 0.4 the default update rule for AdamW has changed to match the pytorch implementation.
+> The previous rule, which is closer to the original paper, can be obtained by setting `AdamW(..., couple=false)`.
+> See [this issue](https://github.com/FluxML/Flux.jl/issues/2433) for more details.
+
 ## Installation
 
 ```julia
diff --git a/src/rules.jl b/src/rules.jl
index f3df9d63..00bc21a7 100644
--- a/src/rules.jl
+++ b/src/rules.jl
@@ -491,8 +491,8 @@ function apply!(o::NAdam, state, x::AbstractArray{T}, dx) where T
 end
 
 """
-    AdamW(η = 0.001, β = (0.9, 0.999), λ = 0, ϵ = 1e-8)
-    AdamW(; [eta, beta, lambda, epsilon])
+    AdamW(η = 0.001, β = (0.9, 0.999), λ = 0, ϵ = 1e-8; couple = true)
+    AdamW(; [eta, beta, lambda, epsilon, couple])
 
 [AdamW](https://arxiv.org/abs/1711.05101) is a variant of Adam fixing (as in repairing) its
 weight decay regularization.
@@ -506,12 +506,54 @@ Implemented as an [`OptimiserChain`](@ref) of [`Adam`](@ref) and [`WeightDecay`]
 - Weight decay (`λ == lambda`): Controls the strength of ``L_2`` regularisation.
 - Machine epsilon (`ϵ == epsilon`): Constant to prevent division by zero
                          (no need to change default)
-"""
-AdamW(η, β = (0.9, 0.999), λ = 0.0, ϵ = 1e-8) =
-  OptimiserChain(Adam(η, β, ϵ), WeightDecay(λ))
+- Keyword `couple`: If `true`, the weight decay is coupled with the learning rate, as in pytorch's AdamW.
+                    This corresponds to an update of the form `x = x - η * (dx + λ * x)`, where `dx` is the
+                    update from Adam with learning rate 1.
+                    If `false`, the weight decay is decoupled from the learning rate, in the spirit of the original paper.
+                    This corresponds to an update of the form `x = x - η * dx - λ * x`.
+                    Default is `true`.
+
+!!! warning "Breaking change in v0.4"
+    With version 0.4 the default update rule for AdamW has changed to match the pytorch implementation.
+    The previous rule, which is closer to the original paper, can be obtained by setting `AdamW(..., couple=false)`.
+    See [this issue](https://github.com/FluxML/Flux.jl/issues/2433) for more details.
+"""
+struct AdamW{T1,T2,T3,T4} <: AbstractRule
+  eta::T1
+  beta::T2
+  epsilon::T3
+  lambda::T4
+  couple::Bool
+end
+
+function AdamW(η, β = (0.9, 0.999), λ = 0.0, ϵ = 1e-8; couple::Bool = true)
+  η < 0 && throw(DomainError(η, "the learning rate cannot be negative"))
+  AdamW(η, β, λ, ϵ, couple)
+end
+
+AdamW(; eta = 0.001, beta = (0.9, 0.999), lambda= 0.0,  epsilon = 1e-8, kw...) =
+  AdamW(eta, beta, lambda, epsilon; kw...)
+
+init(o::AdamW, x::AbstractArray{T}) where T = (zero(x), zero(x), T.(o.beta))
+
+function apply!(o::AdamW, state, x::AbstractArray{T}, dx) where T
+  η, β, ϵ, λ = T(o.eta), T.(o.beta), T(o.epsilon), T(o.lambda)
+  mt, vt, βt = state
 
-AdamW(; eta = 0.001, beta = (0.9, 0.999), lambda = 0, epsilon = 1e-8) =
-  OptimiserChain(Adam(eta, beta, epsilon), WeightDecay(lambda))
+  # standard Adam update with learning rate eta=1
+  @.. mt = β[1] * mt + (1 - β[1]) * dx
+  @.. vt = β[2] * vt + (1 - β[2]) * abs2(dx)
+  dx′ = @lazy mt / (1 - βt[1]) / (sqrt(vt / (1 - βt[2])) + ϵ)
+
+  # apply learning rate and weight decay
+  if o.couple
+    dx′′ = @lazy η * (dx′ + λ * x)
+  else
+    dx′′ = @lazy η * dx′ + λ * x
+  end
+
+  return (mt, vt, βt .* β), dx′′
+end
 
 """
     AdaBelief(η = 0.001, β = (0.9, 0.999), ϵ = 1e-16)


Test patch:
diff --git a/test/rules.jl b/test/rules.jl
index 52a3580e..b539a7ec 100644
--- a/test/rules.jl
+++ b/test/rules.jl
@@ -15,7 +15,7 @@ RULES = [
   OptimiserChain(ClipGrad(0.5), Momentum()),
   OptimiserChain(WeightDecay(), OAdam(), ClipGrad(1)),
   # Not the default:
-  RMSProp(centred = true),
+  RMSProp(centred = true), AdamW(couple=false),
 ]
 
 name(o) = typeof(o).name.name  # just for printing testset headings
diff --git a/test/runtests.jl b/test/runtests.jl
index fc0fe57f..ae2d9d06 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -332,8 +332,7 @@ end
 
     @testset "keyword arguments" begin
       @test Nesterov(rho=0.8, eta=0.1) === Nesterov(0.1, 0.8)
-      @test AdamW(lambda=0.3).opts[1] == Adam()
-      @test AdamW(lambda=0.3).opts[2] == WeightDecay(0.3)
+      @test AdamW(lambda=0.3, eta=0.1) == AdamW(0.1, (0.9, 0.999), 0.3, 1.0e-8)
     end
 
     @testset "forgotten gradient" begin