File size: 3,993 Bytes
eae86f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """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):
# Initialization only. The DCU float32 eigensolver can fail for repeated
# eigenvalues, especially the zero gradients of alpha=0 residual blocks.
# Scaling and CPU float64 preserve the symmetric eigenproblem exactly.
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):
# Descending Rayleigh quotient prevents arbitrary eigenvector swaps.
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'])
# SOAP initializes the preconditioner on the first gradient.
# The first scheduled LR is zero, so no training is skipped.
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
|