Momentum

Momentum#

class braintools.optim.Momentum(lr=0.001, momentum=0.9, weight_decay=0.0, grad_clip_norm=None, grad_clip_value=None)#

Momentum optimizer.

Implements the momentum variant of stochastic gradient descent, where updates accumulate a velocity that persists across iterations to accelerate convergence in relevant directions.

Parameters:

lr (float | LRScheduler) – Learning rate. Can be a float or LRScheduler instance.
momentum (float) – Momentum factor. The fraction of the gradient to retain from previous steps.
weight_decay (float) – Weight decay (L2 penalty) coefficient.
grad_clip_norm (float | None) – Maximum gradient norm for clipping.
grad_clip_value (float | None) – Maximum gradient value for clipping.

Notes

The momentum update is computed as:

\[ \begin{align}\begin{aligned}v_{t+1} = \mu v_t + g_t\\\theta_{t+1} = \theta_t - \alpha v_{t+1}\end{aligned}\end{align} \]

where \(\mu\) is the momentum factor, \(g_t\) is the gradient at step t, \(\alpha\) is the learning rate, \(v_t\) is the velocity, and \(\theta\) are the parameters.

Examples

Basic Momentum optimizer:

>>> import brainstate
>>> import braintools
>>> import jax.numpy as jnp
>>>
>>> # Create model
>>> model = brainstate.nn.Linear(10, 5)
>>> optimizer = braintools.optim.Momentum(lr=0.01, momentum=0.9)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))

Momentum with weight decay:

>>> optimizer = braintools.optim.Momentum(lr=0.01, momentum=0.9, weight_decay=0.0001)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))

Momentum with learning rate scheduling:

>>> scheduler = braintools.optim.StepLR(base_lr=0.1, step_size=30, gamma=0.1)
>>> optimizer = braintools.optim.Momentum(lr=scheduler, momentum=0.9)
>>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
>>>
>>> for epoch in range(100):
...     # Training code here
...     optimizer.step(grads)
...     if (epoch + 1) % epoch_size == 0:
...         scheduler.step()

Momentum

Contents

Momentum#