""" Implementation of a hybrid optimizer. For high-dimensional parameters (tensors with more than 2 dimensions), this optimizer applies steepest descent with respect to a specified matrix norm (induced 1, 2, or infinity norm). The update rules are derived from the provided formulas. For all other parameters (2D or 1D), it defaults to the Adam optimization algorithm. """ import torch import math from optimizer.numerical_methods import zeropower_via_newtonschulz5 class MatrixNormSteepestDescent(torch.optim.Optimizer): """ Implements a hybrid optimization algorithm. It uses matrix norm steepest descent for tensors with dim > 2, and Adam for all other parameters. """ def __init__(self, params, norm_p=2, lr=1e-3, momentum=0.95, betas=(0.9, 0.999), weight_decay=0.1, scale=0.2, eps=1e-8, use_nesterov=False): """ Initializes the optimizer. Args: params (iterable): Iterable of parameters to optimize or dicts defining parameter groups. norm_p (int or float): The matrix norm to use for high-dimensional parameters. Supported values are 1, 2, and float('inf'). lr (float, optional): Learning rate (default: 1e-3). betas (Tuple[float, float], optional): Coefficients used for computing running averages of gradient and its square in Adam (default: (0.9, 0.999)). beta1 is also used as momentum for the matrix norm updates. weight_decay (float, optional): Weight decay (L2 penalty) (default: 0.1). eps (float, optional): Term added to the denominator to improve numerical stability (default: 1e-8). use_nesterov (bool, optional): Enables Nesterov momentum for matrix norm updates (default: False). """ 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}") defaults = dict(norm_p=norm_p, lr=lr, momentum=momentum, betas=betas, weight_decay=weight_decay, scale=scale, eps=eps, use_nesterov=use_nesterov) super(MatrixNormSteepestDescent, self).__init__(params, defaults) def __setstate__(self, state): super(MatrixNormSteepestDescent, 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'] beta1, beta2 = group['betas'] weight_decay = group['weight_decay'] scale = group['scale'] eps = group['eps'] norm_p = group['norm_p'] use_nesterov = group['use_nesterov'] for p in group['params']: if p.grad is None: continue grad = p.grad # Apply decoupled weight decay for all parameters if weight_decay != 0: p.add_(p, alpha=-weight_decay * lr) param_state = self.state[p] # Check parameter dimension to decide update rule if p.dim() >= 2: # --- Matrix Norm Steepest Descent for Tensors with dim >= 2 --- # Momentum calculation if momentum != 0: if 'momentum_buffer' not in param_state: buf = param_state['momentum_buffer'] = torch.clone(grad).detach() else: buf = param_state['momentum_buffer'] buf.mul_(momentum).add_(grad, alpha=1 - momentum) if use_nesterov: Gt = torch.lerp(grad, buf, momentum) else: Gt = buf else: Gt = grad # original_shape = Gt.shape # # Matricize the tensor by reshaping: (d1, d_2*d_3*...) # matrix_Gt = Gt.reshape(original_shape[0], -1) update_magnitude = 0 update_direction = torch.zeros_like(Gt) if norm_p == 1: # Based on the formula for the induced 1-norm update col_max_abs = torch.max(torch.abs(Gt), dim=-2, keepdim=True)[0] update_magnitude = torch.sum(col_max_abs) # Direction is ColumnMax(Gt) mask = (torch.abs(Gt) >= col_max_abs - 1e-12) update_direction[mask] = torch.sign(Gt[mask]) elif norm_p == float('inf'): # Based on the formula for the induced infinity-norm update row_max_abs = torch.max(torch.abs(Gt), dim=-1, keepdim=True)[0] update_magnitude = torch.sum(row_max_abs) # Direction is RowMax(Gt) mask = (torch.abs(Gt) >= row_max_abs - 1e-12) update_direction[mask] = torch.sign(Gt[mask]) elif norm_p == 2: # Implement Muon (similar to the https://arxiv.org/pdf/2502.16982 version since we have the normalization) try: update_direction = zeropower_via_newtonschulz5(G=Gt) except torch.linalg.LinAlgError: # If NewtonSchulz fails, we skip the update for this parameter to avoid crashing. update_magnitude = 0 update_direction.fill_(0) elif norm_p == '2,inf': # +++ NEW: Row-wise Normalized Gradient Descent +++ # Formula: G_i / ||G_i||_2 for each row i # 1. Calculate L2 norm for each row (dim=-1) # Gt shape is likely (Out, In) or (B, Out, In) etc. # We assume the last dimension is the "input" dimension for summation # row_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True) # !!!!!!! here fore embedding layers, we consider using col-wise normalization # This is also sth interesting: column-wise normalization is 1,2 induced steepest descent col_norms = torch.norm(Gt, p=2, dim=-1, keepdim=True) # print(f'The shape of column norm: {col_norms.shape}') # 2. Normalize gradient rows # update_direction = Gt / (row_norms + 1e-12) update_direction = Gt / (col_norms + 1e-12) else: raise ValueError(f'Invalid matrix norm choice for p.dim() >= 2. Only 1, 2, and inf are supported.') # Combine magnitude and direction and reshape back to original tensor shape # final_update = (update_direction * update_magnitude).reshape(original_shape) update_direction = update_direction / (torch.linalg.norm(update_direction) + 1e-12) alpha = -lr * scale * math.sqrt(p.numel()) p.add_(update_direction, alpha=alpha) else: # --- Adam Fallback for Params with dim == 1 --- # Initialize Adam state if 'step' not in param_state: param_state['step'] = 0 param_state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) 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'] # Adam update logic exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2) bias_correction1 = 1 - beta1 ** step bias_correction2 = 1 - beta2 ** step step_size = lr / bias_correction1 denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) p.addcdiv_(exp_avg, denom, value=-step_size) return loss