| """ |
| utils.grokfast — accelerated grokking by amplifying slow-varying gradient |
| components (Lee et al. 2024, arXiv:2405.20233). |
| |
| Maintain an EMA of gradients across steps; the slow-EMA component |
| corresponds to the generalising circuit. Adding it back into the live |
| gradient (scaled by `lamb`) accelerates the grokking transition 20-100×. |
| """ |
|
|
| from __future__ import annotations |
|
|
|
|
| def gradfilter_ema(model, grads_ema, alpha: float = 0.98, lamb: float = 2.0): |
| """ |
| Call this AFTER `loss.backward()` and BEFORE `optimizer.step()`. |
| |
| Args: |
| model: the network whose gradients we are filtering. |
| grads_ema: dict {param_name: ema_grad}, or None on the first call. |
| alpha: EMA decay (0.98 → very slow, emphasises persistent grads). |
| lamb: amplification factor for the slow component. |
| |
| Returns: |
| Updated `grads_ema` dict — pass it back in on the next step. |
| """ |
| if grads_ema is None: |
| grads_ema = {} |
|
|
| for name, p in model.named_parameters(): |
| if p.requires_grad and p.grad is not None: |
| if name not in grads_ema: |
| grads_ema[name] = p.grad.data.detach().clone() |
| else: |
| grads_ema[name] = ( |
| grads_ema[name] * alpha |
| + p.grad.data.detach() * (1 - alpha) |
| ) |
| p.grad.data = p.grad.data + grads_ema[name] * lamb |
|
|
| return grads_ema |
|
|
