optax.LookaheadParams

class optax.LookaheadParams(fast: optax.Params, slow: optax.Params)

Holds a pair of slow and fast parameters for the lookahead optimizer.

Gradients should always be calculated with the fast parameters (i.e., params.fast). The slow parameters should be used for testing and inference as they generalize better. See the reference for a detailed discussion.

fast

Fast parameters (use these for gradient computation).

Type:

base.Params

slow

Slow parameters (use these for inference).

Type:

base.Params

References

Zhang et al, Lookahead Optimizer: k steps forward, 1 step back, 2019