Lookahead (Zhang et al., 2019) wraps any optimiser. The inner optimiser updates "fast" weights for k steps; then the "slow" weights move part of the way towards them and the fast weights are reset to the slow ones:
Implement a class Lookahead(params, inner_step, k=5, alpha=0.5), where params is a NumPy array (the initial weights) and inner_step(params, grad) returns updated weights:
step(grad): apply the inner step to the fast weights. After every k-th call, perform the slow update and sync. Return the current fast weights.fast and slow