Momentum

Momentum#

class seli.opt.Momentum(lr: float = 0.001, beta: float = 0.9)[source]#

Bases: Optimizer

Momentum optimizer.

Accelerates optimization by accumulating a velocity vector in the direction of persistent gradient directions. This is analogous to the momentum of a ball rolling down a hill.

The velocity is updated with the gradient and a decay factor. The decay factor is a hyperparameter that controls the influence of previous gradients on the current update.

For well-behaved functions, momentum often leads to faster convergence, when compared to SGD.

Methods Summary

call_param(key, grad, **_)

Process the gradients of a single parameter.

Methods Documentation

call_param(key: str, grad: Float[Array, '*s'], **_) Float[Array, '*s'][source]#

Process the gradients of a single parameter. This function is useful for implementing custom optimizers that essentially run the same function for all parameters. This is the case for most well known optimizers.

Parameters:
  • loss (Float[Array, ""]) – The absolute loss value.

  • key (str) – The key of the parameter.

  • grad (Float[Array]) – The gradients of the parameter.

  • param (Float[Array]) – The parameter values.

Returns:

grad – The processed gradients of the parameter.

Return type:

Float[Array]