File size: 5,478 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 | import torch
import math
class ColumnWiseAdamW(torch.optim.Optimizer):
"""
Implements a Column-wise AdamW optimization algorithm.
The first-order momentum (exp_avg) is tracked element-wise, but the
second-order momentum (exp_avg_sq) is tracked column-wise for tensors
with dim >= 2.
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2):
"""
Initializes the optimizer.
Args:
params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
lr (float, optional): Learning rate (default: 1e-3).
betas (Tuple[float, float], optional): Coefficients used for computing running averages
of gradient and its square (default: (0.9, 0.999)).
eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8).
weight_decay (float, optional): Weight decay coefficient (default: 1e-2).
"""
if lr < 0.0:
raise ValueError(f"Invalid learning rate: {lr}")
if not 0.0 <= betas[0] < 1.0:
raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
if not 0.0 <= betas[1] < 1.0:
raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
if weight_decay < 0.0:
raise ValueError(f"Invalid weight_decay value: {weight_decay}")
if eps < 0.0:
raise ValueError(f"Invalid epsilon value: {eps}")
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
super(ColumnWiseAdamW, self).__init__(params, defaults)
@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']
beta1, beta2 = group['betas']
weight_decay = group['weight_decay']
eps = group['eps']
for p in group['params']:
if p.grad is None:
continue
grad = p.grad
# --- 1. Decoupled Weight Decay ---
if weight_decay != 0:
p.add_(p, alpha=-lr * weight_decay)
param_state = self.state[p]
# --- 2. State Initialization ---
if 'step' not in param_state:
param_state['step'] = 0
# First moment is always element-wise
param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Second moment is column-wise for matrices (dim >= 2)
if p.dim() >= 2:
# By setting the 0-th dimension to 1, we keep exactly one value per column
# For nn.Linear weight (out_features, in_features), shape becomes (1, in_features)
shape = list(p.shape)
shape[0] = 1
param_state['exp_avg_sq'] = torch.zeros(shape, dtype=p.dtype, device=p.device)
else:
param_state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
exp_avg = param_state['exp_avg']
exp_avg_sq = param_state['exp_avg_sq']
param_state['step'] += 1
step = param_state['step']
# --- 3. First Moment Update ---
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
# --- 4. Second Moment Update ---
if p.dim() >= 2:
# Column-wise formulation: [v_t]_j = beta_2 * [v_{t-1}]_j + (1 - beta_2) * ||[g_t]_j||_2^2
# Summing the squares over the 0-th dimension achieves the squared L2 norm for each column.
# Using grad * grad.conj().real supports both real and complex tensors (identical to PyTorch backend).
col_sq_norm = torch.sum(grad * grad.conj(), dim=0, keepdim=True).real
exp_avg_sq.mul_(beta2).add_(col_sq_norm, alpha=1 - beta2)
else:
# Fallback to standard element-wise Adam for 1D tensors (e.g., bias)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
# --- 5. Bias Correction & Step ---
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2 ** step
# 核心改动：对于高维张量（dim >= 2），补充 \sqrt{d_out} 以对齐 Adam 的 scale
if p.dim() >= 2:
d_out = p.shape[0]
step_size = (lr / bias_correction1) * math.sqrt(d_out)
else:
step_size = lr / bias_correction1
# Thanks to keepdim=True, exp_avg_sq shape (1, in_features) naturally broadcasts
# against exp_avg shape (out_features, in_features) during division.
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
p.addcdiv_(exp_avg, denom, value=-step_size)
return loss |