| """Independent SOAP: Adam in Shampoo eigen-coordinates (Appendix C.2). |
| |
| Uses orthogonal iteration (QR) to track eigenvectors every two steps. First |
| moments live in physical parameter coordinates; second moments are permuted |
| when eigenvectors are re-ordered. Vectors/scalars use unrotated Adam. |
| """ |
| import torch |
|
|
|
|
| class SOAP(torch.optim.Optimizer): |
| def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, |
| shampoo_beta=0.999, precondition_frequency=2): |
| super().__init__(params, dict(lr=lr, betas=betas, eps=eps, |
| shampoo_beta=shampoo_beta, |
| precondition_frequency=precondition_frequency)) |
|
|
| @staticmethod |
| def _basis(cov): |
| |
| |
| |
| work = cov.detach().to(device='cpu', dtype=torch.float64) |
| scale = float(work.abs().max()) |
| if scale == 0.0: |
| return torch.eye(cov.shape[0], device=cov.device, dtype=cov.dtype) |
| work = (work + work.T) * (0.5 / scale) |
| _, q = torch.linalg.eigh(work) |
| return q.flip(-1).to(device=cov.device, dtype=cov.dtype) |
|
|
| @staticmethod |
| def _iterate(cov, q): |
| |
| scores = (q * (cov @ q)).sum(0) |
| order = torch.argsort(scores, descending=True) |
| q = q[:, order] |
| scale = cov.abs().amax().clamp_min(1e-30) |
| q, _ = torch.linalg.qr((cov / scale) @ q) |
| return q, order |
|
|
| @torch.no_grad() |
| def step(self, closure=None): |
| result = None |
| if closure is not None: |
| with torch.enable_grad(): |
| result = closure() |
| for group in self.param_groups: |
| b1, b2 = group['betas'] |
| sb, freq = group['shampoo_beta'], group['precondition_frequency'] |
| for p in group['params']: |
| if p.grad is None: |
| continue |
| g = p.grad |
| s = self.state[p] |
| if not s: |
| s['step'] = 0 |
| s['m'] = torch.zeros_like(p) |
| s['v'] = torch.zeros_like(p) |
| if p.ndim == 2: |
| s['left'] = (1.0 - sb) * (g @ g.T) |
| s['right'] = (1.0 - sb) * (g.T @ g) |
| s['ql'] = self._basis(s['left']) |
| s['qr'] = self._basis(s['right']) |
| |
| |
| s['initialized_only'] = True |
| continue |
| s['step'] += 1 |
| step = s['step'] |
| s['m'].lerp_(g, 1.0 - b1) |
| if p.ndim == 2: |
| gl = s['ql'].T @ g @ s['qr'] |
| ml = s['ql'].T @ s['m'] @ s['qr'] |
| else: |
| gl, ml = g, s['m'] |
| s['v'].mul_(b2).addcmul_(gl, gl, value=1.0 - b2) |
| denom = (s['v'] / (1.0 - b2 ** step)).sqrt().add_(group['eps']) |
| update = (ml / (1.0 - b1 ** step)) / denom |
| if p.ndim == 2: |
| update = s['ql'] @ update @ s['qr'].T |
| p.add_(update, alpha=-group['lr']) |
| if p.ndim == 2: |
| s['left'].mul_(sb).addmm_(g, g.T, beta=1.0, alpha=1.0 - sb) |
| s['right'].mul_(sb).addmm_(g.T, g, beta=1.0, alpha=1.0 - sb) |
| if step % freq == 0: |
| s['ql'], il = self._iterate(s['left'], s['ql']) |
| s['qr'], ir = self._iterate(s['right'], s['qr']) |
| s['v'] = s['v'][il][:, ir].contiguous() |
| return result |
|
|
