File size: 5,359 Bytes
ad93bc0 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | # activation_gradient_preconditioned_adam.py
import torch
import math
from torch.optim.optimizer import Optimizer
class ASGO(Optimizer):
"""
Implements Activation-Gradient-Preconditioned Adam.
This optimizer uses second-order statistics from both input activations (X_t)
and output gradients (G_Yt) to form two preconditioners. This is closely
related to the Shampoo optimizer.
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.8), weight_decay=0.1, eps=1e-10, scale=0.2, use_grafting=False, use_nesterov=False):
# betas[0]: momentum (β1)
# betas[1]: adam second moment (β2 in pseudocode, for V')
# betas[2]: preconditioners (β3 in pseudocode, for L and R)
# if not (isinstance(betas, (list, tuple)) and len(betas) == 3):
# raise ValueError("Betas must be a tuple/list of 3 elements.")
defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eps=eps, scale=scale, use_grafting=use_grafting, use_nesterov=use_nesterov)
super(ASGO, self).__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
lr = group['lr']
if len(group['betas']) == 2:
beta1, beta2_adam = group['betas']
beta2_precond = beta2_adam
else:
beta1, beta2_adam, beta2_precond = group['betas']
weight_decay = group['weight_decay']
eps = group['eps']
scale = group['scale']
use_grafting = group['use_grafting']
use_nesterov = group['use_nesterov']
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
state = self.state[p]
if len(state) == 0:
state['step'] = 0
state['momentum_buffer'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# R_t: Right preconditioner
state['right_preconditioner'] = torch.zeros((p.shape[-1], p.shape[-1]), device=p.device)
# V'_t: Adam's 2nd moment buffer for grafting
if use_grafting:
# beta2_adam = beta2_precond
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
state['step'] += 1
step = state['step']
momentum_buffer = state['momentum_buffer']
right_preconditioner = state['right_preconditioner']
if use_grafting:
exp_avg_sq = state['exp_avg_sq']
exp_avg_sq.mul_(beta2_adam).addcmul_(grad, grad, value=1 - beta2_adam)
GTG = torch.mm(grad.t(), grad)
# Step 5: Update momentum
momentum_buffer.mul_(beta1).add_(grad, alpha=1 - beta1)
if use_nesterov:
Gt = torch.lerp(grad, momentum_buffer, beta1)
else:
Gt = momentum_buffer
# Step 6 & 7: Update preconditioners
# left_preconditioner.mul_(beta2_precond).add_(GGT, alpha=1 - beta2_precond)
right_preconditioner.mul_(beta2_precond).add_(GTG, alpha=1 - beta2_precond)
# Step 8 & 9: Eigendecomposition of preconditioners
# we use float32 for eigendecomposition to ensure numerical stability
# eps_L = torch.eye(left_preconditioner.size(0), device=p.device) * eps
# S_L, U = torch.linalg.eigh(left_preconditioner.float() + eps_L)
eps_R = torch.eye(right_preconditioner.size(0), device=p.device) * eps
# S_R, V = torch.linalg.eigh(right_preconditioner.float() + eps_R)
# update_direction = Gt @ V @ torch.diag(S_R.rsqrt()) @ V.t()
U, S_R, Vh = torch.linalg.svd(right_preconditioner.float() + eps_R)
update_direction = Gt @ U @ torch.diag(S_R.rsqrt()) @ Vh
# normalize the update
update_direction = update_direction / (torch.linalg.norm(update_direction) + eps)
if use_grafting:
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2_adam ** step
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
# The magnitude is calculated based on the effective gradient
adam_update_magnitude = torch.linalg.norm((Gt / bias_correction1) / denom)
alpha = -lr * adam_update_magnitude
else:
# Default scaling if not grafting
alpha = -lr * scale * math.sqrt(p.numel())
if weight_decay > 0.0:
p.add_(p, alpha=-lr * weight_decay)
p.add_(update_direction, alpha=alpha)
return loss |