File size: 4,912 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 | """
Implementation of normalized steepest descent algorithm. We additionally apply multipliers to align the RMS norm update with Adam.
Note that when p=2 is chosen, the algorithm is the same as normalized SGD.
"""
import torch
import math
class NormalizedSteepestDescent(torch.optim.Optimizer):
"""
Implements a normalized version of Stochastic Gradient Descent (with momentum).
The update for each parameter is normalized by its Frobenius norm before
being applied. This is a form of layer-wise adaptivity.
"""
def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12, use_nesterov=False, use_grafting=False):
# scale is used for aligning the update RMS with Adam
if lr < 0.0:
raise ValueError(f"Invalid learning rate: {lr}")
if momentum < 0.0:
raise ValueError(f"Invalid momentum value: {momentum}")
if weight_decay < 0.0:
raise ValueError(f"Invalid weight_decay value: {weight_decay}")
# use_grafting not implemented yet
defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov)
super(NormalizedSteepestDescent, self).__init__(params, defaults)
def __setstate__(self, state):
super(NormalizedSteepestDescent, self).__setstate__(state)
@torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step."""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
norm_p = group['norm_p']
lr = group['lr']
momentum = group['momentum']
weight_decay = group['weight_decay']
eps = group['eps']
scale = group['scale']
use_nesterov = group['use_nesterov']
for p in group['params']:
if p.grad is None:
continue
# Get the gradient
d_p = p.grad
# Apply momentum
if momentum != 0:
param_state = self.state[p]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
else:
buf = param_state['momentum_buffer']
buf.mul_(momentum).add_(d_p, alpha=1-momentum) # m_t = beta * m_{t-1} + (1-beta) * g_t
# This is the update tensor (m_t)
if use_nesterov:
# use nesterov momentum, following Muon and Lion
Gt = torch.lerp(d_p, buf, momentum)
else:
Gt = buf
else:
# If no momentum, the gradient is the update tensor
Gt = d_p
mn = p.numel() # parameter count
if norm_p <= 1 + eps:
# in this case, only the largest coordiantes have update
g_max_abs = torch.max(Gt.abs())
mask = Gt.abs() >= g_max_abs - eps
# Count how many elements have the max value (k in the formula)
k = torch.sum(mask).float()
# The update direction is sign(G_t) where |G_t| is max, and 0 otherwise.
update_direction = torch.zeros_like(Gt)
update_direction[mask] = torch.sign(Gt[mask])
# Calculate the final update magnitude based on the limit
alpha = -lr * scale * torch.sqrt(mn / (k + eps))
elif norm_p == float('inf'):
# in this case, we apply SignSGD
update_direction = torch.sign(Gt)
alpha = -lr * scale
else:
norm_q = 1 / (1 - 1 / norm_p)
assert norm_q > 0
# Calculate the update direction: |G_t|^(q-1) * sign(G_t)
update_direction = torch.pow(Gt.abs(), norm_q - 1) * torch.sign(Gt)
update_norm = torch.linalg.norm(update_direction)
update_direction = update_direction / (update_norm + eps)
# Calculate the pre-factor
alpha = -lr * scale * math.sqrt(mn)
# Apply weight decay
if weight_decay > 0.0:
p.mul_(1.0 - lr * weight_decay)
# Apply the update
p.add_(update_direction, alpha=alpha)
return loss |