AdamW step with decoupled weight decay — morie_adamw

Adam applies the decay outside the adaptive 1/sqrt(v) scaling, so parameters with a large gradient history are not decayed less. Supply theta or the step reduces to plain Adam.

Usage

morie_adamw_step(
  g,
  beta1 = 0.9,
  beta2 = 0.999,
  lr = 0.001,
  wd = 0.01,
  eps = 1e-08,
  theta = NULL,
  state = NULL
)

Arguments

g: Gradient of the unregularised loss.
beta1, beta2: Moment decay rates in [0, 1).
lr: Step size.
wd: Decoupled weight-decay coefficient.
eps: Denominator floor.
theta: Current parameters, needed for the decay term.
state: List from the previous call, or NULL.

Value

List with update, state, m, v, decay, t, step_norm.

References

Loshchilov, I., & Hutter, F. (2019). Decoupled weight decay regularization. ICLR 2019.

Examples

round(morie_adamw_step(0, lr = 0.1, wd = 0.5, theta = 2)$update, 6)
#> [1] -0.1