RMSprop optimiser — gradient normalisation by EMA of squared magnitude, with optional centring and Polyak-style momentum.
The plain rule replaces Adagrad's monotonic accumulator with an exponential moving average, so the effective learning rate stops shrinking and tracks the current gradient regime instead.
Math
Uncentred:
Centred — subtract the squared mean to recover the true variance:
With momentum the normalised gradient is folded into a velocity buffer first:
References
Tieleman & Hinton, "Lecture 6.5 — RMSProp" (Coursera 2012).
See Also
Adadelta, Adagrad, Rprop
Parameters
paramsvector of TensorImpllrfloat, default 1e-2alphafloat, default 0.99epsfloat, default 1e-8weight_decayfloat, default 0.0momentumfloat, default 0.0centeredbool, default falseAttributes
square_avg_vector of Storagegrad_avg_vector of Storagecentered == true).moment_buf_vector of Storagemomentum != 0).Notes
State buffers are allocated lazily — code that only ever uses the plain variant does not pay the memory cost of the centred or momentum-augmented variants.
Constructors
1Construct the optimiser and resize per-slot state vectors.
Parameters
paramsvector of TensorImpllrfloat, default 1e-2alphafloat, default 0.99epsfloat, default 1e-8weight_decayfloat, default 0.0momentumfloat, default 0.0centeredbool, default falseVirtual methods
2Methods
3Allocate per-slot state buffers for parameter i.
Parameters
iintNotes
square_avg_[i] is always allocated. grad_avg_[i] is
allocated only when centered_ is true and moment_buf_[i]
only when momentum_ is non-zero — see Attributes on
RMSprop.
Return the versioned state-dict identifier.
Returns
str"rmsprop_v1".
Apply one RMSprop update to parameter i.
Parameters
iintgStoragep on this step.Notes
Dispatches to a fused MLX kernel on GPU and a scalar Accelerate loop on CPU; the four variant combinations (plain / centred / momentum / both) share a single branchy path that elides the unused buffer reads.