Lookahead Optimizer Wrapper

Lookahead keeps two copies of the parameters. An inner optimizer updates the fast copy on every step. At regular synchronization steps, the slow copy moves partway toward the newly updated fast copy, and both copies are set to that interpolated value.

A synchronization step

After the inner optimizer produces fast_candidate, use:

new_slow = slow + slow_step_size * (fast_candidate - slow)
new_fast = new_slow

Between synchronization steps, retain the old slow copy and keep the inner optimizer's fast result. Lookahead changes the path of the updates; it does not guarantee better convergence for every problem.

This task starts from fresh state

Both copies initially equal params, and the inner optimizer is plain SGD. Perform only the first step:

• If sync_period > 1, no synchronization happens yet. Return params - lr * grads; the slow copy is unchanged.
• If sync_period = 1, synchronize immediately. Return params - slow_step_size * lr * grads.

For params = [1, 2], grads = [2, -2], lr = 0.1 and slow_step_size = 0.5, the fast SGD candidate is [0.8, 2.2]. Return that candidate when the period is 5. With a period of 1, interpolate halfway from the initial values and return [0.9, 2.1].

Using the wrapper

Build optax.lookahead(optax.sgd(lr), sync_period=int(sync_period), slow_step_size=slow_step_size). Initialize both copies with optax.LookaheadParams(fast=params, slow=params) and use this paired structure for initialization, update and applying updates. The supplied gradient array is for the fast parameters; do not duplicate it into a pair.

A longer run must keep the returned paired parameters and optimizer state between calls. Reinitializing every call would restart the synchronization count. This exercise returns only the first step's fast values.