# src/optimizer/burst_optimizer.py import sys from pathlib import Path import torch from torch.optim import Optimizer import geoopt def _ensure_fisher_rao_path() -> None: candidates = [ Path(__file__).resolve().parents[2] / "fisher_rao", # repo-local package parent Path.home() / "Projects" / "Fisher_Rao", ] for root in candidates: if root.exists() and str(root.parent) not in sys.path: sys.path.insert(0, str(root.parent)) return _ensure_fisher_rao_path() try: from fisher_rao.burst_modulation import FisherInfoBurstModulator except ImportError: FisherInfoBurstModulator = None # type: ignore[misc, assignment] class GeooptBurstOptimizer(Optimizer): """ Highly adaptive Riemannian optimizer for the Sphere manifold, balancing speed and stability: - Adaptive quantile-based gradient clipping - Burst triggering driven by clipped Riemannian gradient norms (real-time adaptation) - Dynamic momentum damping (higher retention → more aggressive when loss high) - Dynamic burst factor (higher multiplier → more aggressive bursts when loss high) - Loss-scaled noise injection and adaptive burst threshold - Clear phased step() for readability """ def __init__( self, params, lr=0.02, momentum=0.92, burst_threshold=7.0, burst_factor_max=6.0, burst_factor_min=2.0, max_theta=torch.pi / 12, damping=0.99, damping_min=0.92, damping_loss_scale=100000.0, damping_ema_alpha=0.05, twist_rate=1.0, stagnation_window=40, stagnation_thresh=1e-3, stagnation_noise_amp=0.01, warm_up_steps=3000, burst_schedule_interval=250, grad_normalization=True, max_rgrad_norm=float("inf"), min_burst_threshold=4.0, max_burst_threshold=10.0, adapt_burst_threshold=True, good_improve_multiplier=8.0, verbose=True, adaptive_grad_clip=True, clip_quantile=0.95, clip_multiplier=2.0, use_fisher_modulation=False, fisher_modulator=None, fisher_info_scale=100.0, ): defaults = dict( lr=lr, momentum=momentum, burst_threshold=burst_threshold, burst_factor_max=burst_factor_max, burst_factor_min=burst_factor_min, max_theta=max_theta, damping=damping, damping_min=damping_min, damping_loss_scale=damping_loss_scale, damping_ema_alpha=damping_ema_alpha, twist_rate=twist_rate, warm_up_steps=warm_up_steps, burst_schedule_interval=burst_schedule_interval, grad_normalization=grad_normalization, max_rgrad_norm=max_rgrad_norm, adaptive_grad_clip=adaptive_grad_clip, clip_quantile=clip_quantile, clip_multiplier=clip_multiplier, ) super().__init__(params, defaults) self.stagnation_window = stagnation_window self.stagnation_thresh = stagnation_thresh self.base_stagnation_noise_amp = stagnation_noise_amp self.min_burst_threshold = min_burst_threshold self.max_burst_threshold = max_burst_threshold self.adapt_burst_threshold = adapt_burst_threshold self.good_improve_multiplier = good_improve_multiplier self.verbose = verbose self.use_fisher_modulation = use_fisher_modulation if fisher_modulator is not None: self.fisher_modulator = fisher_modulator elif use_fisher_modulation and FisherInfoBurstModulator is not None: self.fisher_modulator = FisherInfoBurstModulator(info_scale=fisher_info_scale) else: self.fisher_modulator = None if use_fisher_modulation and self.fisher_modulator is None: raise ImportError( "use_fisher_modulation=True requires fisher_rao package at ~/Projects/Fisher_Rao" ) self.current_step = 0 self.burst_count = 0 self.loss_history = [] self.smoothed_loss = None # Validate and store manifold self.manifold = None for group in self.param_groups: for p in group["params"]: if hasattr(p, "manifold"): if self.manifold is None: self.manifold = p.manifold elif self.manifold is not p.manifold: raise ValueError("All parameters must share the same manifold.") if not isinstance(self.manifold, geoopt.manifolds.Sphere): raise ValueError("Optimizer designed for geoopt Sphere manifold.") # Initial setup with torch.no_grad(): for group in self.param_groups: group["twist"] = 0.0 for p in group["params"]: p.data = self.manifold.projx(p.data) @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() self.current_step += 1 loss_val = None if loss is not None: loss_val = loss.item() if torch.is_tensor(loss) else float(loss) if self.smoothed_loss is None: self.smoothed_loss = loss_val else: alpha = self.param_groups[0]["damping_ema_alpha"] self.smoothed_loss = (1 - alpha) * self.smoothed_loss + alpha * loss_val global_twist_proxy = 0.0 global_num_instances = 0 stagnation_noise_scale = 1.0 for group in self.param_groups: # Warm-up effective_lr = group["lr"] if group["warm_up_steps"] > 0: progress = min(1.0, self.current_step / group["warm_up_steps"]) effective_lr *= progress # First pass: collect rgrad norms for adaptive clipping rgrad_norms = [] for p in group["params"]: if p.grad is None: continue x = p.data egrad = p.grad.data rgrad = self.manifold.egrad2rgrad(x, egrad) instance_norms = rgrad.norm(dim=-1) rgrad_norms.append(instance_norms) # Determine clip norm if rgrad_norms: all_norms = torch.cat(rgrad_norms) if group["adaptive_grad_clip"] and all_norms.numel() > 1: clip_norm = torch.quantile(all_norms, group["clip_quantile"]) * group["clip_multiplier"] else: clip_norm = group["max_rgrad_norm"] else: clip_norm = group["max_rgrad_norm"] # Dynamic damping & burst factor (aggressive when loss high) if self.smoothed_loss is not None: loss_factor = 1 - torch.exp(torch.tensor(-self.smoothed_loss / group["damping_loss_scale"])) effective_damping = group["damping_min"] + (group["damping"] - group["damping_min"]) * loss_factor effective_burst_factor = group["burst_factor_min"] + (group["burst_factor_max"] - group["burst_factor_min"]) * loss_factor else: effective_damping = group["damping"] effective_burst_factor = group["burst_factor_max"] # Fisher-information modulation (conservative in high-curvature regions) fisher_factors = None effective_burst_threshold = group["burst_threshold"] if self.use_fisher_modulation and self.fisher_modulator is not None and loss is not None: params_with_grad = [p for p in group["params"] if p.grad is not None] if params_with_grad: log_lik = -loss # surrogate log-likelihood fisher_factors = self.fisher_modulator.update(params_with_grad, log_lik) effective_damping = min( 0.999, effective_damping + fisher_factors["damping_boost"], ) effective_burst_factor *= fisher_factors["burst_factor_scale"] effective_burst_threshold = group["burst_threshold"] * fisher_factors["threshold_scale"] stagnation_noise_scale = fisher_factors["stagnation_boost"] if self.verbose and self.current_step % 500 == 0: print( f"[Fisher] info={fisher_factors['smoothed_fisher_info']:.2e} " f"sens={fisher_factors['sensitivity']:.3f} " f"burst_scale={fisher_factors['burst_factor_scale']:.3f}" ) # Second pass: clip, accumulate twist, update momentum for p in group["params"]: if p.grad is None: continue state = self.state[p] if "v" not in state: state["v"] = torch.zeros_like(p.data) state["old_x"] = p.data.clone() x = p.data egrad = p.grad.data rgrad = self.manifold.egrad2rgrad(x, egrad) # Clip rgrad_norm = rgrad.norm(dim=-1, keepdim=True) scale = torch.min(torch.ones_like(rgrad_norm), clip_norm / rgrad_norm.clamp(min=1e-8)) clipped_rgrad = rgrad * scale clipped_norm = clipped_rgrad.norm(dim=-1, keepdim=True) # Optional normalization if group["grad_normalization"]: mean_clipped_norm = clipped_norm.mean().clamp(min=1e-8) clipped_rgrad = clipped_rgrad * (1.0 / mean_clipped_norm) # Accumulate twist twist_inc = clipped_norm.mean() * group["twist_rate"] group["twist"] += twist_inc.item() global_twist_proxy += twist_inc.item() global_num_instances += x.shape[0] # Update momentum velocity state["v"] = group["momentum"] * state["v"] + effective_lr * clipped_rgrad # Burst logic twist_triggered = group["twist"] > effective_burst_threshold scheduled = (group["burst_schedule_interval"] > 0 and self.current_step % group["burst_schedule_interval"] == 0) apply_burst = twist_triggered or scheduled factor = min(20.0, effective_burst_factor if apply_burst else 1.0) if self.verbose: if scheduled: print(f"[Scheduled Burst] Step {self.current_step}") if factor > 1.0 + 1e-3: print( f"[Burst] Factor={factor:.2f} " f"(twist_trigger={twist_triggered}, scheduled={scheduled}, twist={group['twist']:.2f})" ) if self.current_step % 500 == 0: print( f"[Dynamic] Step {self.current_step} | Damping={effective_damping:.4f} | " f"Burst Factor={effective_burst_factor:.2f} | Smoothed Loss≈{self.smoothed_loss:.2f}" ) if apply_burst: group["twist"] = 0.0 self.burst_count += 1 # Apply updates for p in group["params"]: state = self.state[p] if "v" not in state: continue update_vec = factor * state["v"] # Cap step size theta = update_vec.norm(dim=-1, keepdim=True).clamp(min=1e-8) cap_mask = theta > group["max_theta"] if cap_mask.any(): update_vec = update_vec * torch.where( cap_mask, group["max_theta"] / theta, torch.ones_like(theta) ) # Retraction new_x = self.manifold.retr(p.data, update_vec) state["old_x"].copy_(p.data) p.data = new_x # Transport velocity state["v"] = self.manifold.transp(state["old_x"], new_x, state["v"] * effective_damping) # Stagnation detection & adaptation if loss_val is not None: self.loss_history.append(loss_val) if len(self.loss_history) > self.stagnation_window: self.loss_history.pop(0) if len(self.loss_history) == self.stagnation_window: improvement = self.loss_history[0] - loss_val if improvement < self.stagnation_thresh: noise_amp = self.base_stagnation_noise_amp * stagnation_noise_scale if self.verbose: print(f"[Stagnation] Plateau detected → injecting noise (amp={noise_amp:.4f})") for group in self.param_groups: for p in group["params"]: state = self.state[p] if "v" in state: noise = noise_amp * torch.randn_like(state["v"]) noise = noise - self.manifold.inner(p.data, noise, keepdim=True) * p.data state["v"] += noise if self.adapt_burst_threshold: ref_threshold = self.param_groups[0]["burst_threshold"] if improvement > self.good_improve_multiplier * self.stagnation_thresh: new_threshold = max(self.min_burst_threshold, ref_threshold * 0.9) elif improvement < self.stagnation_thresh: new_threshold = min(self.max_burst_threshold, ref_threshold * 1.111) else: new_threshold = ref_threshold if abs(new_threshold - ref_threshold) > 1e-6: for g in self.param_groups: g["burst_threshold"] = new_threshold if self.verbose: direction = "↓" if new_threshold < ref_threshold else "↑" print(f"[Adapt] Burst threshold {direction} {new_threshold:.2f}") return loss # Updated src/optimizer/burst_optimizer.py excerpt — HierarchicalGeooptBurstOptimizer only # Changes: # - Stronger response to stagnation (higher multiplier cap + sensitivity) # - Much higher headroom on boosts (50× burst, 10× theta) # - Conservative resets: only restore originals on *significant* improvement (>2× thresh) # This prevents premature reset cascades and lets strong boosts persist longer when needed # src/optimizer/burst_optimizer.py (updated Hierarchical caps) class HierarchicalGeooptBurstOptimizer: def __init__( self, optimizer, stagnation_window=200, stagnation_thresh=1e-3, burst_boost=2.0, theta_boost=1.3, name="Macro", verbose=True, ): self.optimizer = optimizer self.stagnation_window = stagnation_window self.stagnation_thresh = stagnation_thresh self.burst_boost = burst_boost self.theta_boost = theta_boost self.name = name self.verbose = verbose self.loss_history = [] self.originals = {id(g): (g["burst_factor_max"], g["max_theta"]) for g in self.param_groups} @property def param_groups(self): return self.optimizer.param_groups def step(self, closure=None): loss = self.optimizer.step(closure=closure) if loss is not None: loss_val = loss.item() if torch.is_tensor(loss) else loss self.loss_history.append(loss_val) if len(self.loss_history) > self.stagnation_window: self.loss_history.pop(0) if len(self.loss_history) == self.stagnation_window: improvement = self.loss_history[0] - loss_val if improvement < self.stagnation_thresh: lack_of_improvement = max(0.0, self.stagnation_thresh - improvement) severity = lack_of_improvement multiplier = 1.0 + min(8.0, severity / self.stagnation_thresh * 5.0) if self.verbose: print( f"[{self.name} Burst] Long-term stagnation (severity={severity:.1e}) → " f"boosting burst_max x{self.burst_boost * multiplier:.2f}, theta x{self.theta_boost * multiplier:.2f}" ) for group in self.param_groups: orig_bf, orig_mt = self.originals[id(group)] boosted_bf = orig_bf * self.burst_boost * multiplier boosted_theta = orig_mt * self.theta_boost * multiplier group["burst_factor_max"] = min(orig_bf * 80.0, boosted_bf) group["max_theta"] = min(orig_mt * 15.0, boosted_theta) else: if improvement > 2 * self.stagnation_thresh: modified = any( abs(g["burst_factor_max"] - self.originals[id(g)][0]) > 1e-6 or abs(g["max_theta"] - self.originals[id(g)][1]) > 1e-6 for g in self.param_groups ) if modified and self.verbose: print(f"[{self.name} Reset] Significant progress → restoring original burst_factor_max/max_theta") for group in self.param_groups: orig_bf, orig_mt = self.originals[id(group)] group["burst_factor_max"] = orig_bf group["max_theta"] = orig_mt return loss def zero_grad(self): self.optimizer.zero_grad()