Skip to content

← Optax step 5 of 25

Medium Primitives

Optax RMSprop Step

Jump to the task

RMSprop rescales each gradient using a running average of its squared values. Parameters with larger recent squared gradients get a larger denominator. This exercise uses Optax's default, uncentered RMSprop without momentum or bias correction.

For each parameter, the update is

new_v = 0.9 * v + 0.1 * grads ** 2
new_params = params - lr * grads / sqrt(new_v + 1e-8)

Epsilon is inside the square root here. Different optimizer definitions can place it elsewhere, so follow this convention.

One step from fresh state

Initialize the squared-gradient buffer to zero for this exercise. Thus new_v is 0.1 * grads ** 2. With a parameter of 1, gradient 0.1 and lr=0.01, the denominator is sqrt(0.001 + 1e-8) and the new parameter is approximately 0.968377.

When 0.1 * grads ** 2 is much larger than 1e-8, the first change has magnitude approximately lr / sqrt(0.1), or 3.162 times lr. This approximation does not apply to tiny gradients. A zero gradient makes no change. There is no bias correction in the configuration used here.

Inputs

Implement rmsprop_step(params, grads, lr):

  • params and grads are nonempty 1D floating-point JAX arrays of the same shape.
  • lr is a nonnegative scalar learning rate.
  • Return the updated parameter array with the same shape.

Create optax.rmsprop(lr) and initialize its state from params. Call its update function once with the supplied gradients and state, then apply the returned updates with optax.apply_updates. The updates already include the learning-rate scaling and descent sign; do not multiply by lr or negate them again.

Each function call starts from a fresh state. In a real training loop, initialize once and carry the returned optimizer state between steps; reinitializing on every step would discard the running average.

The widget below illustrates carrying a value through successive steps. Its displayed rule is 0.5 * carry + x, a simple moving-sum example, not the RMSprop recurrence above. It does not calculate optimizer updates.

Loading visualization…