class

AdamW

extendsOptimizer

AdamW

source edit

AdamW optimiser — Adam with decoupled weight decay (Loshchilov & Hutter, 2019).

Decouples the L2 penalty from the adaptive moment estimation: the weight-decay term is applied directly to the parameter, NOT folded into the gradient $g_t$ . This restores the rotation-invariance lost by Adam's adaptive learning rate when L2 is added to $g$ , and is strongly preferred over Adam + L2 for transformer training.

Math

m_t = \beta_1\, m_{t-1} + (1 - \beta_1)\, g_t, \qquad v_t = \beta_2\, v_{t-1} + (1 - \beta_2)\, g_t^2

\hat m_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat v_t = \frac{v_t}{1 - \beta_2^t}

\theta_{t+1} = (1 - \eta \lambda)\, \theta_t - \eta \cdot \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}

References

Loshchilov & Hutter, "Decoupled Weight Decay Regularization" (ICLR 2019). Kingma & Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015).

See Also

Adam : L2-regularisation variant.

Attributes

lr_float

Step size

\eta

. Default 1e-3.

beta1_, beta2_float

EMA decay rates. Defaults 0.9 / 0.999.

eps_float

Denominator stabiliser. Default 1e-8.

weight_decay_float

Decoupled penalty coefficient

\lambda

. Default 1e-2 — larger than Adam's because it is no longer scaled by the adaptive denominator.

step_count_int64

Global step counter.

m_, v_Storage[]

Per-parameter first- and second-moment buffers.

scalar_cache_AdamScalarCache

Per-step broadcast cache, including wd_factor.

Notes

The implementation realises the decay as a multiplicative pre-scale of $\theta_t$ by $(1 - \eta \lambda)$ (cached in AdamScalarCache::wd_factor), then performs the standard Adam step. Equivalent to the additive form above but avoids a separate MLX op invocation per parameter.

Constructors

ctor

AdamW

void AdamW(int)

source edit

Construct an AdamW optimiser bound to params.

Parameters

paramsvector of TensorImpl shared pointers

Parameters to optimise.

lrfloat= None

Step size

\eta

. Default 1e-3.

beta1, beta2float= None

First- and second-moment EMA decay rates. Defaults 0.9 and 0.999.

epsfloat= None

Denominator stabiliser

\epsilon

. Default 1e-8.

weight_decayfloat= None

Decoupled penalty coefficient

\lambda

. Default 1e-2.

Virtual methods

virtconst

lr

double lr()

source edit

Current learning rate $\eta$ .

virt

set_lr

void set_lr(double lr)

source edit

Set the learning rate and invalidate the scalar cache.

Notes

Same rationale as Adam::set_lr — additionally invalidates the cached AdamW wd_factor = $(1 - \eta \lambda)$ scalar.

Methods

init_state_slot

void init_state_slot(int slot_idx, const int & param)

source edit

Allocate zero-initialised $m$ and $v$ buffers for slot_idx.

load_state_buffers

void load_state_buffers(const int & bufs)

source edit

Restore per-parameter state. See Adam::load_state_buffers.

set_step_count

void set_step_count(int s)

source edit

Set the global step counter.

fnconst

state_buffers

int state_buffers()

source edit

Export per-parameter state (exp_avg and exp_avg_sq).

fnconst

state_dict_id

→str

int state_dict_id()

source edit

Schema identifier for state_dict serialisation.

Returns

str

"adamw_v1".

fnconst

step_count

int step_count()

source edit

Global step counter $t$ .

update_one

void update_one(int slot_idx, int & param, const int & grad)

source edit

Apply the decoupled-weight-decay Adam update.

Notes

Dispatches to adam_step_{cpu,gpu} with decoupled_wd = true so the kernel applies $\theta \leftarrow (1 - \eta \lambda) \theta$ before the Adam step.

class

AdamW

extendsOptimizer

AdamW

source edit

AdamW optimiser — Adam with decoupled weight decay (Loshchilov & Hutter, 2019).

Math

m_t = \beta_1\, m_{t-1} + (1 - \beta_1)\, g_t, \qquad v_t = \beta_2\, v_{t-1} + (1 - \beta_2)\, g_t^2

\hat m_t = \frac{m_t}{1 - \beta_1^t}, \qquad \hat v_t = \frac{v_t}{1 - \beta_2^t}

\theta_{t+1} = (1 - \eta \lambda)\, \theta_t - \eta \cdot \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}

References

Loshchilov & Hutter, "Decoupled Weight Decay Regularization" (ICLR 2019). Kingma & Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015).

See Also

Adam : L2-regularisation variant.

Attributes

lr_float

Step size

\eta

. Default 1e-3.

beta1_, beta2_float

EMA decay rates. Defaults 0.9 / 0.999.

eps_float

Denominator stabiliser. Default 1e-8.

weight_decay_float

Decoupled penalty coefficient

\lambda

. Default 1e-2 — larger than Adam's because it is no longer scaled by the adaptive denominator.

step_count_int64

Global step counter.

m_, v_Storage[]

Per-parameter first- and second-moment buffers.

scalar_cache_AdamScalarCache

Per-step broadcast cache, including wd_factor.

Notes

Constructors

ctor

AdamW

void AdamW(int)

source edit

Construct an AdamW optimiser bound to params.

Parameters

paramsvector of TensorImpl shared pointers

Parameters to optimise.

lrfloat= None

Step size

\eta

. Default 1e-3.

beta1, beta2float= None

First- and second-moment EMA decay rates. Defaults 0.9 and 0.999.

epsfloat= None

Denominator stabiliser

\epsilon

. Default 1e-8.

weight_decayfloat= None

Decoupled penalty coefficient

\lambda

. Default 1e-2.

Virtual methods

virtconst

lr

double lr()

source edit

Current learning rate $\eta$ .

virt

set_lr

void set_lr(double lr)

source edit

Set the learning rate and invalidate the scalar cache.

Notes

Same rationale as Adam::set_lr — additionally invalidates the cached AdamW wd_factor = $(1 - \eta \lambda)$ scalar.

Methods

init_state_slot

void init_state_slot(int slot_idx, const int & param)

source edit

Allocate zero-initialised $m$ and $v$ buffers for slot_idx.

load_state_buffers

void load_state_buffers(const int & bufs)

source edit

Restore per-parameter state. See Adam::load_state_buffers.

set_step_count

void set_step_count(int s)

source edit

Set the global step counter.

fnconst

state_buffers

int state_buffers()

source edit

Export per-parameter state (exp_avg and exp_avg_sq).

fnconst

state_dict_id

→str

int state_dict_id()

source edit

Schema identifier for state_dict serialisation.

Returns

str

"adamw_v1".

fnconst

step_count

int step_count()

source edit

Global step counter $t$ .

update_one

void update_one(int slot_idx, int & param, const int & grad)

source edit

Apply the decoupled-weight-decay Adam update.

Notes

Dispatches to adam_step_{cpu,gpu} with decoupled_wd = true so the kernel applies $\theta \leftarrow (1 - \eta \lambda) \theta$ before the Adam step.