class

Adam

extendsOptimizer

Adam

source edit

Adam optimiser — adaptive moment estimation (Kingma & Ba, 2014).

Maintains exponential moving averages of the gradient and squared gradient for each parameter, then applies a bias-corrected step scaled element-wise by the inverse second-moment $\sqrt{\hat v}$ .

Weight decay (when non-zero) is added to the gradient before the moment update — i.e. the L2-regularisation form, NOT the decoupled AdamW form. Use AdamW for decoupled decay.

Math

Per parameter, with global step counter $t =$ step_count_:

g_t \leftarrow g_t + \lambda\, \theta_t \quad \text{(if } \lambda > 0\text{)}

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} = \theta_t - \eta \cdot \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}

References

Kingma & Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015). Reddi, Kale & Kumar, "On the Convergence of Adam and Beyond" (ICLR 2018) — AMSGrad variant.

See Also

AdamW : decoupled-weight-decay variant. NAdam : Nesterov-momentum variant. RAdam : variance-rectified variant.

Attributes

lr_float

Step size

\eta

. Default 1e-3.

beta1_, beta2_float

EMA decay rates for the first and second moments. Defaults 0.9 and 0.999 from the original paper.

eps_float

Denominator stabiliser

\epsilon

. Default 1e-8.

weight_decay_float

L2 penalty coefficient

\lambda

added to the gradient. Set to 0 to disable.

amsgrad_bool

Reserved — see Notes. Currently must be false.

step_count_int64

Global step counter; bias-correction factors are functions of this. Persisted in state_dict under state_step.

m_, v_Storage[]

Per-parameter first-moment and second-moment buffers, matching each parameter's shape and dtype.

scalar_cache_AdamScalarCache

Per-step broadcast cache for the GPU kernel.

Notes

step_count_ is a global counter incremented once per step() call regardless of how many parameters are updated, so all parameters share the same bias-correction factor within a step. This matches reference-framework semantics.

AMSGrad (Reddi et al., 2018) is declared via the amsgrad ctor flag but not yet implemented — passing true raises not_implemented at construction time.

Constructors

ctor

Adam

void Adam(int)

source edit

Construct an Adam optimiser bound to params.

Parameters

paramsvector of TensorImpl shared pointers

Parameters to optimise. One m and one v buffer is allocated per parameter on first step() call.

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

L2 penalty coefficient

\lambda

. Default 0.

amsgradbool= None

Reserved. Must be false — see class Notes.

Raises

not_implemented

If amsgrad == true.

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.

Parameters

lrfloat

New step size

\eta

Notes

The 3.4 perf pass folds $\eta$ (and $1 - \eta \lambda$ for AdamW) into AdamScalarCache; mutating $\eta$ here flips scalar_cache_.valid = false so the cached MLX scalars are rebuilt on the next step.

Methods

fnconst

beta1

double beta1()

source edit

First-moment EMA decay rate $\beta_1$ .

fnconst

beta2

double beta2()

source edit

Second-moment EMA decay rate $\beta_2$ .

fnconst

eps

double eps()

source edit

Denominator stabiliser $\epsilon$ .

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.

Parameters

slot_idxsize_t

Index into m_ / v_.

paramTensorImpl shared pointer

Reference parameter — buffers inherit its shape / dtype.

load_state_buffers

void load_state_buffers(const int & bufs)

source edit

Restore per-parameter state from a previous state_buffers.

Parameters

bufsvector of NamedBuffers

Must have the same length as the parameter list and contain exp_avg / exp_avg_sq keys.

Raises

runtime_error

On length / key / dtype / shape mismatch.

set_step_count

void set_step_count(int s)

source edit

Set the global step counter (used by checkpoint loading).

fnconst

state_buffers

→vector of NamedBuffers

int state_buffers()

source edit

Export per-parameter state (exp_avg = $m$ , exp_avg_sq = $v$ ).

Returns

vector of NamedBuffers

One entry per parameter, each containing the named buffers 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

"adam_v1".

fnconst

step_count

int step_count()

source edit

Global step counter $t$ used for bias correction.

update_one

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

source edit

Apply the standard Adam update to a single parameter.

Parameters

slot_idxsize_t

Index into m_ / v_.

paramTensorImpl shared pointer

In-place parameter to update.

gradStorage

Gradient

g_t

matching param shape / dtype / device.

Notes

Dispatches to adam_step_gpu (GPU stream) or adam_step_cpu (CPU stream) with decoupled_wd = false.

class

Adam

extendsOptimizer

Adam

source edit

Adam optimiser — adaptive moment estimation (Kingma & Ba, 2014).

Maintains exponential moving averages of the gradient and squared gradient for each parameter, then applies a bias-corrected step scaled element-wise by the inverse second-moment $\sqrt{\hat v}$ .

Weight decay (when non-zero) is added to the gradient before the moment update — i.e. the L2-regularisation form, NOT the decoupled AdamW form. Use AdamW for decoupled decay.

Math

Per parameter, with global step counter $t =$ step_count_:

g_t \leftarrow g_t + \lambda\, \theta_t \quad \text{(if } \lambda > 0\text{)}

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} = \theta_t - \eta \cdot \frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon}

References

Kingma & Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015). Reddi, Kale & Kumar, "On the Convergence of Adam and Beyond" (ICLR 2018) — AMSGrad variant.

See Also

AdamW : decoupled-weight-decay variant. NAdam : Nesterov-momentum variant. RAdam : variance-rectified variant.

Attributes

lr_float

Step size

\eta

. Default 1e-3.

beta1_, beta2_float

EMA decay rates for the first and second moments. Defaults 0.9 and 0.999 from the original paper.

eps_float

Denominator stabiliser

\epsilon

. Default 1e-8.

weight_decay_float

L2 penalty coefficient

\lambda

added to the gradient. Set to 0 to disable.

amsgrad_bool

Reserved — see Notes. Currently must be false.

step_count_int64

Global step counter; bias-correction factors are functions of this. Persisted in state_dict under state_step.

m_, v_Storage[]

Per-parameter first-moment and second-moment buffers, matching each parameter's shape and dtype.

scalar_cache_AdamScalarCache

Per-step broadcast cache for the GPU kernel.

Notes

AMSGrad (Reddi et al., 2018) is declared via the amsgrad ctor flag but not yet implemented — passing true raises not_implemented at construction time.

Constructors

ctor

Adam

void Adam(int)

source edit

Construct an Adam optimiser bound to params.

Parameters

paramsvector of TensorImpl shared pointers

Parameters to optimise. One m and one v buffer is allocated per parameter on first step() call.

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

L2 penalty coefficient

\lambda

. Default 0.

amsgradbool= None

Reserved. Must be false — see class Notes.

Raises

not_implemented

If amsgrad == true.

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.

Parameters

lrfloat

New step size

\eta

Notes

Methods

fnconst

beta1

double beta1()

source edit

First-moment EMA decay rate $\beta_1$ .

fnconst

beta2

double beta2()

source edit

Second-moment EMA decay rate $\beta_2$ .

fnconst

eps

double eps()

source edit

Denominator stabiliser $\epsilon$ .

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.

Parameters

slot_idxsize_t

Index into m_ / v_.

paramTensorImpl shared pointer

Reference parameter — buffers inherit its shape / dtype.

load_state_buffers

void load_state_buffers(const int & bufs)

source edit

Restore per-parameter state from a previous state_buffers.

Parameters

bufsvector of NamedBuffers

Must have the same length as the parameter list and contain exp_avg / exp_avg_sq keys.

Raises

runtime_error

On length / key / dtype / shape mismatch.

set_step_count

void set_step_count(int s)

source edit

Set the global step counter (used by checkpoint loading).

fnconst

state_buffers

→vector of NamedBuffers

int state_buffers()

source edit

Export per-parameter state (exp_avg = $m$ , exp_avg_sq = $v$ ).

Returns

vector of NamedBuffers

One entry per parameter, each containing the named buffers 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

"adam_v1".

fnconst

step_count

int step_count()

source edit

Global step counter $t$ used for bias correction.

update_one

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

source edit

Apply the standard Adam update to a single parameter.

Parameters

slot_idxsize_t

Index into m_ / v_.

paramTensorImpl shared pointer

In-place parameter to update.

gradStorage

Gradient

g_t

matching param shape / dtype / device.

Notes

Dispatches to adam_step_gpu (GPU stream) or adam_step_cpu (CPU stream) with decoupled_wd = false.