jax_sgmc.scheduler.polynomial_step_size

jax_sgmc.scheduler.polynomial_step_size(a=1.0, b=1.0, gamma=0.33)

Polynomial decreasing step size schedule.

Implements the original proposal of a polynomial step size schedule \(\epsilon = a(b + n)^{\gamma}\).

Parameters:

• a (Array) – Scale of all step sizes

• b (Array) – Stabilization constant

• gamma (Array) – Decay constant

Return type:

specific_scheduler

Returns:

Returns triplet as described above.