Merge pull request #26 from theomgdev/copilot/improve-chaosgrad-optimizer

Improve ChaosGrad diagnostics correctness and decay stability

Author: theomgdev

odyssnet/training/chaos_optimizer.py

@@ -484,7 +484,8 @@ def step(self, closure=None):
 
                 # ---- Weight decay ----
                 if per_decay > 0.0 and not is_hebbian:
-                    p.data.mul_(1.0 - genesis_lr * per_decay)
+                    decay_factor = max(self._EPS, 1.0 - genesis_lr * per_decay)
+                    p.data.mul_(decay_factor)
 
                 # ---- Parameter update ----
                 p.data.add_(v_hat, alpha=-(genesis_lr * per_lr / denom))
@@ -569,10 +570,11 @@ def get_diagnostics(self, debug: bool = False) -> dict:
         genesis_lr    = self.defaults['lr']
         effective_lrs = []
         init_lrs      = []
-        for state, _ in all_states:
+        for state, group in all_states:
             if 'per_param_lr' in state and 'init_lr' in state:
+                group_genesis_lr = group.get('lr', self.defaults['lr'])
                 effective_lrs.append(state['per_param_lr'] / state['init_lr'])
-                init_lrs.append(state['init_lr'] * genesis_lr)
+                init_lrs.append(state['init_lr'] * group_genesis_lr)
 
         diag: dict = {
             'global_step':      self._global_step,
@@ -601,14 +603,17 @@ def get_diagnostics(self, debug: bool = False) -> dict:
             g_states = [self.state[p] for p in group['params'] if self.state.get(p)]
             if not g_states:
                 continue
+            g_genesis_lr = group.get('lr', self.defaults['lr'])
             g_lrs   = [s['per_param_lr'] / s['init_lr'] for s in g_states if 'per_param_lr' in s]
+            g_init_lrs = [s['init_lr'] * g_genesis_lr for s in g_states if 'init_lr' in s]
             g_betas  = [s['per_param_beta']  for s in g_states if 'per_param_beta'  in s]
             g_alphas = [s['per_param_alpha']  for s in g_states if 'per_param_alpha' in s]
             g_decays = [s['per_param_decay']  for s in g_states if 'per_param_decay' in s]
             group_stats.append({
                 'group_name':       gname,
                 'param_count':      len(g_states),
                 'avg_effective_lr': (sum(g_lrs)    / len(g_lrs))    if g_lrs    else 0.0,
+                'avg_init_lr':      (sum(g_init_lrs) / len(g_init_lrs)) if g_init_lrs else 0.0,
                 'avg_beta':         (sum(g_betas)   / len(g_betas))  if g_betas  else 0.0,
                 'avg_alpha':        (sum(g_alphas)  / len(g_alphas)) if g_alphas else 0.0,
                 'avg_decay':        (sum(g_decays)  / len(g_decays)) if g_decays else 0.0,

tests/training/test_chaos_optimizer_extra.py

@@ -138,6 +138,31 @@ def test_lr_group_override_propagates(self):
         assert torch.allclose(m.W.data.diagonal(), torch.zeros(m.num_neurons))
         assert math.isfinite(m.W.data.norm().item())
 
+    def test_get_diagnostics_respects_group_lr_override(self):
+        m   = _model()
+        opt = _opt(m, lr=1e-4)
+        _one_step_raw(opt, m)
+
+        base_diag = opt.get_diagnostics()
+        base_debug_diag = opt.get_diagnostics(debug=True)
+        for pg in opt.param_groups:
+            pg['lr'] = 5e-5
+        new_diag = opt.get_diagnostics()
+        new_debug_diag = opt.get_diagnostics(debug=True)
+
+        assert new_diag['avg_init_lr'] == pytest.approx(base_diag['avg_init_lr'] * 0.5, rel=1e-6)
+
+        base_group_init = {
+            g['group_name']: g['avg_init_lr']
+            for g in base_debug_diag['param_groups']
+        }
+        new_group_init = {
+            g['group_name']: g['avg_init_lr']
+            for g in new_debug_diag['param_groups']
+        }
+        for gname, base_avg in base_group_init.items():
+            assert new_group_init[gname] == pytest.approx(base_avg * 0.5, rel=1e-6)
+
 
 # ---------------------------------------------------------------------------
 # reset_param_state