import torch import math from torch.optim.optimizer import Optimizer from optimizer.numerical_methods import zeropower_via_newtonschulz5 class EigenBlockOptimizer(Optimizer): """ Implements Eigen-Block Basis Optimizer. Features: - Hessian Approximation via Shampoo/SOAP style preconditioners (L, R). - Eigen-Block Rotation: Groups close eigenvalues and rotates within blocks to maximize L1 sparsity (using Newton-Schulz). - Momentum: Applies momentum to gradients in the original space. - Dual Mode: Supports both SignSGD (default) and Adam updates in the optimal basis. """ def __init__(self, params, lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0, eigen_tol=0.1, eps=1e-8, use_adam=False): """ Args: lr (float): Learning rate. betas (tuple): (beta1, beta2). beta1: Momentum coefficient. beta2: Preconditioner EMA coefficient (and Adam 2nd moment if use_adam=True). eigen_tol (float): Relative tolerance for grouping eigenvalues. Groups lambdas where (lambda_i - lambda_start) < tol * max_lambda. use_adam (bool): If True, uses Adam updates in the projected space. If False, uses SignSGD updates in the projected space. """ defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, eigen_tol=eigen_tol, eps=eps, use_adam=use_adam) super(EigenBlockOptimizer, 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'] beta1, beta2 = group['betas'] weight_decay = group['weight_decay'] eigen_tol = group['eigen_tol'] eps = group['eps'] use_adam = group['use_adam'] for p in group['params']: if p.grad is None: continue grad = p.grad if grad.dim() != 2: # Fallback for 1D vectors (bias/layernorm): Standard Adam or SignSGD state = self.state[p] if len(state) == 0: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p) if use_adam: state['exp_avg_sq'] = torch.zeros_like(p) state['step'] += 1 exp_avg = state['exp_avg'] # Update momentum exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) if weight_decay > 0: p.mul_(1 - lr * weight_decay) if use_adam: exp_avg_sq = state['exp_avg_sq'] exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) denom = exp_avg_sq.sqrt().add_(eps) # Bias correction bias_correction1 = 1 - beta1 ** state['step'] bias_correction2 = 1 - beta2 ** state['step'] step_size = lr * math.sqrt(bias_correction2) / bias_correction1 p.addcdiv_(exp_avg, denom, value=-step_size) else: # SignSGD with momentum p.add_(exp_avg.sign(), alpha=-lr) continue # --- Handle 2D Matrices (Weights) --- state = self.state[p] # Init state if len(state) == 0: state['step'] = 0 state['momentum_buffer'] = torch.zeros_like(p) # 1st moment (M) # Shampoo/SOAP style preconditioners state['L'] = torch.zeros((grad.size(0), grad.size(0)), device=grad.device, dtype=grad.dtype) state['R'] = torch.zeros((grad.size(1), grad.size(1)), device=grad.device, dtype=grad.dtype) if use_adam: # Adam 2nd moment (V') in the PROJECTED space # Note: We store this same size as p, assuming basis consistency over time state['adam_v_prime'] = torch.zeros_like(p) state['step'] += 1 L = state['L'] R = state['R'] momentum = state['momentum_buffer'] # 1. Update Momentum (M_t) in Original Space momentum.mul_(beta1).add_(grad, alpha=1 - beta1) # 2. Update Preconditioners (Hessian Approximation) # Using EMA of GG^T and G^TG L.mul_(beta2).add_(grad @ grad.T, alpha=1 - beta2) R.mul_(beta2).add_(grad.T @ grad, alpha=1 - beta2) # 3. Eigen Decomposition # Add 1e-12 for numerical stability L_evals, U_L = torch.linalg.eigh(L + 1e-12 * torch.eye(L.size(0), device=L.device)) R_evals, U_R = torch.linalg.eigh(R + 1e-12 * torch.eye(R.size(0), device=R.device)) # 4. Rotate Momentum to Eigen-basis: M_tilde = U_L^T M U_R # We rotate M instead of G because we are doing Momentum-based updates M_tilde = U_L.T @ momentum @ U_R # 5. Compute Optimal Basis Q_L (Left) Q_L = U_L.clone() # Use max eigenvalue to scale tolerance tol_L = eigen_tol * (L_evals.abs().max().item() + 1e-6) boundaries_L = self._find_block_boundaries(L_evals, tol_L) # --- NEW: Record L Block Stats --- self._store_block_stats(state, 'L', boundaries_L, L_evals.shape[0]) for start, end in boundaries_L: if end - start <= 1: continue # Extract rows corresponding to this block from M_tilde # We use M_tilde to determine the rotation N within the block M_sub = M_tilde[start:end, :] # Formula: N = msign(A * sign(A)^T) # Using M_sub aligns the basis with the momentum direction's sparsity A = M_sub @ M_sub.sign().T N_block = zeropower_via_newtonschulz5(A) # Update Q_L columns: Q_L_new = Q_L_old @ N_block Q_L[:, start:end] = U_L[:, start:end] @ N_block # 6. Compute Optimal Basis Q_R (Right) Q_R = U_R.clone() tol_R = eigen_tol * (R_evals.abs().max().item() + 1e-6) boundaries_R = self._find_block_boundaries(R_evals, tol_R) # --- NEW: Record R Block Stats --- self._store_block_stats(state, 'R', boundaries_R, R_evals.shape[0]) for start, end in boundaries_R: if end - start <= 1: continue M_sub = M_tilde[:, start:end] # Formula: N = msign(sign(A)^T * A) A = M_sub.sign().T @ M_sub N_block = zeropower_via_newtonschulz5(A) Q_R[:, start:end] = U_R[:, start:end] @ N_block # 7. Project Momentum and Gradient into the FULL optimal basis Q # M_opt = Q_L^T M Q_R M_opt = Q_L.T @ momentum @ Q_R if use_adam: # If using Adam, we also need the projected Gradient to update V' # G_opt = Q_L^T G Q_R G_opt = Q_L.T @ grad @ Q_R adam_v_prime = state['adam_v_prime'] # Update V'_t using projected gradient square adam_v_prime.mul_(beta2).addcmul_(G_opt, G_opt, value=1 - beta2) # Bias Correction bias_correction1 = 1 - beta1 ** state['step'] bias_correction2 = 1 - beta2 ** state['step'] M_hat = M_opt / bias_correction1 V_hat = adam_v_prime / bias_correction2 # Adam Update in Basis Update_inner = M_hat / (V_hat.sqrt() + eps) else: # SignSGD with Momentum in Basis Update_inner = M_opt.sign() # 8. Rotate update back to original space # D_t = Q_L * Update_inner * Q_R^T Update_step = Q_L @ Update_inner @ Q_R.T if weight_decay > 0: p.mul_(1 - lr * weight_decay) p.add_(Update_step, alpha=-lr) return loss def _find_block_boundaries(self, evals, tol): """ Groups eigenvalues based on difference from the FIRST element in the block. Iterates through sorted eigenvalues. If evals[i] - evals[start] > tol, a new block is started at i. """ # evals is expected to be sorted (torch.linalg.eigh returns ascending) boundaries = [0] n = len(evals) if n == 0: return [] start_idx = 0 for i in range(1, n): # Check difference against the anchor (start_idx) diff = evals[i] - evals[start_idx] if diff > tol: # Close current block, start new one at i boundaries.append(i) start_idx = i boundaries.append(n) # Convert to (start, end) tuples blocks = [] for k in range(len(boundaries) - 1): blocks.append((boundaries[k], boundaries[k+1])) return blocks def _store_block_stats(self, state, name, blocks, total_dim): """ 计算 10%, 25%, 50%, 75%, 90%, 100% 处的特征值索引属于第几个 Block，并存储到 state 中。 """ if 'eigen_stats' not in state: state['eigen_stats'] = {} # 1-based block index (第几个格子) stats = {} percentiles = [0.10, 0.25, 0.50, 0.75, 0.90, 1.00] for p in percentiles: # 计算对应的 index (0-based) # 例如 100% 对应最后一个元素 (total_dim - 1) target_idx = int(p * total_dim) - 1 target_idx = max(0, min(target_idx, total_dim - 1)) # 查找该 index 属于哪个 block found_block_idx = 0 # default for i, (start, end) in enumerate(blocks): if start <= target_idx < end: found_block_idx = i + 1 # 使用 1-based 计数，方便观察 break # 记录 metric，例如 'p10': 1 stats[f'p{int(p*100)}'] = found_block_idx stats['total_blocks'] = len(blocks) stats['total_dim'] = total_dim state['eigen_stats'][name] = stats