File size: 3,081 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 | """
Implementation of normalized SGD algorithm. We additionally apply 0.2 * \sqrt{mn} to align the RMS norm update with Adam.
"""
import torch
import math
class NormalizedSGD(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, lr=1e-3, momentum=0.9, weight_decay=0.1, scale=0.2, eps=1e-12):
# 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}")
defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay, scale=scale, eps=eps)
super(NormalizedSGD, self).__init__(params, defaults)
def __setstate__(self, state):
super(NormalizedSGD, 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:
lr = group['lr']
momentum = group['momentum']
weight_decay = group['weight_decay']
eps = group['eps']
scale = group['scale']
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)
update_vec = buf
else:
# If no momentum, the gradient is the update tensor
update_vec = d_p
# +++ CORE MODIFICATION: Normalize the update vector +++
# Calculate the Frobenius norm of the update vector
norm = torch.linalg.norm(update_vec)
# Apply weight decay
if weight_decay > 0.0:
p.mul_(1.0 - lr * weight_decay)
# Normalize and apply the update
# p_new = p_old - lr * (m_t / ||m_t||_F)
scale_align = scale * math.sqrt(p.numel())
p.add_(update_vec, alpha=-lr * scale_align / (norm + eps))
return loss |