Major updates to automagic3 optimizer. Seems to be functioning more ideally and naturally decays, as it should.

toolkit/optimizers/automagic3.py

@@ -1,3 +1,7 @@
+"""
+NOTE: This is experimental and under active development; expect breaking changes and bugs. Feedback welcome.
+"""
+import math
 from typing import List
 import torch
 
@@ -9,15 +13,29 @@ class Automagic3(torch.optim.Optimizer):
     A learning rate is kept per row of each parameter: one lr per output
     channel for >=2D weights (e.g. one lr per output neuron of a Linear layer)
     and one lr per element for 1D weights (biases, norms). Each step the lr is
-    nudged by how strongly the per-element update direction agrees with the
+    nudged by how consistent the per-element update direction has been.
-    previous step. Agreement is the plain fraction of elements that kept their
+
-    sign (every element counts equally, regardless of update magnitude) and is
+    Agreement is window-consistency: over the last ``polarity_history_count``
-    used as a proportional signal in [-1, 1] that scales a multiplicative
+    update-sign snapshots plus the current sign, the per-element score is the
-    (geometric) bump ``lr *= exp(signal * lr_bump_rate)``, so each step is a fixed
+    fraction sharing the dominant sign (``max(p, 1-p)``, in [0.5, 1]; direction-
-    fractional move that behaves uniformly across the whole [min_lr, max_lr]
+    agnostic, and the current step is just one vote so a single flip against a
-    range and a full up bump is exactly undone by a full down bump. The lr
+    long consistent history barely moves it). It is reduced to one value per row.
-    drifts smoothly, self-stabilises near random agreement, and is clamped to
+    ``polarity_history_count`` is min 1 (default 2); a longer window makes the lr
-    [min_lr, max_lr].
+    react to a sustained trend rather than a single noisy step.
+
+    The agreement maps to a direction in [-1, 1] measured relative to the
+    window's noise floor ``b`` -- the consistency a pure-noise row shows purely
+    by chance, computed automatically from the window size (e.g. 0.75 for a
+    window of <=3 snapshots, ~0.69 for 4). Agreement == ``b`` -> 0 (lr steady),
+    1.0 -> +1 (lr up), below ``b`` -> down. There is no manual target: holding at
+    the noise floor is self-balancing per row -- the lr grows while a row is more
+    consistent than chance and settles at whatever lr makes its consistency meet
+    ``b``, so noisy and clean layers each find their own operating point without
+    tuning, and it neither collapses to min_lr nor runs to max_lr. Measuring
+    relative to ``b`` also makes the signal window-size independent. The direction
+    scales a multiplicative (geometric) bump ``lr *= exp(direction *
+    lr_bump_rate)`` -- a fixed fractional move, uniform across the whole range.
+    lr is clamped to [min_lr, max_lr].
 
     With ``fused=True`` (default) the step is fused into the backward pass via
     ``register_post_accumulate_grad_hook``: each parameter is updated and its
@@ -47,10 +65,10 @@ class Automagic3(torch.optim.Optimizer):
     2. Proportional lr control (was a hard threshold flip). v2 bumped the lr up
        or down by a fixed amount depending on whether agreement crossed a
        threshold, which jitters when agreement hovers near the boundary. v3
-       scales the bump by how strongly the directions agree
+       scales the bump by how far consistency is from its hold point relative to
-       (``2*agreement - 1`` in [-1, 1]). Plain English: the lr nudges gently
+       the noise floor (direction in [-1, 1]). Plain English: the lr nudges gently when the
-       when the signal is weak and firmly when it is strong, and parks itself
+       signal is weak and firmly when it is strong, and parks itself instead of
-       instead of oscillating when gradients are basically noise.
+       oscillating when gradients are basically noise.
 
     3. Multiplicative (geometric) lr bump (was additive). v2 added/subtracted a
        fixed absolute amount, so the same bump was a huge relative jump near
@@ -88,11 +106,15 @@ class Automagic3(torch.optim.Optimizer):
         eps: float = 1e-30,
         clip_threshold: float = 1.0,
         weight_decay: float = 0.0,
+        polarity_history_count: int = 3,  # update-sign snapshots kept; current is compared vs all (min 1)
         fused: bool = True,
     ):
         if lr > 1e-3:
             print(f"Warning! Start lr {lr} is very high; forcing to 1e-6.")
             lr = 1e-6
+        # Agreement compares the current update sign against the stored history,
+        # so at least one snapshot must be kept (1 = compare to previous only).
+        polarity_history_count = max(1, int(polarity_history_count))
         defaults = dict(
             lr=lr,
             min_lr=min_lr,
@@ -102,6 +124,10 @@ class Automagic3(torch.optim.Optimizer):
             eps=eps,
             clip_threshold=clip_threshold,
             weight_decay=weight_decay,
+            polarity_history_count=polarity_history_count,
+            # Noise floor of the consistency measure for this window (history+1),
+            # subtracted off so the lr signal is window-size independent.
+            agreement_floor=self._noise_floor(polarity_history_count + 1),
         )
         super().__init__(params, defaults)
 
@@ -136,6 +162,16 @@ class Automagic3(torch.optim.Optimizer):
     def _rms(t: torch.Tensor) -> torch.Tensor:
         return t.norm(2) / (t.numel() ** 0.5)
 
+    @staticmethod
+    def _noise_floor(window: int) -> float:
+        # Expected window-consistency max(p, 1-p) of a pure-noise element over a
+        # window of `window` independent fair coin flips. This is > 0.5 and grows
+        # with smaller windows (0.75 for window<=3, ~0.688 for 4, ...), so it must
+        # be subtracted off for the lr signal to behave the same at any window.
+        total = sum(math.comb(window, k) * max(k, window - k)
+                    for k in range(window + 1))
+        return total / (window * (2 ** window))
+
     @staticmethod
     def _approx_sq_grad(row: torch.Tensor, col: torch.Tensor) -> torch.Tensor:
         r = (row / row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
@@ -210,7 +246,9 @@ class Automagic3(torch.optim.Optimizer):
         state["lr"] = torch.full(
             lr_shape, float(group["lr"]), dtype=torch.float32, device=p.device
         )
-        state["last_polarity"] = torch.zeros(p.shape, dtype=torch.bool, device=p.device)
+        # Rolling history of the last polarity_history_count update-sign
+        # snapshots (bool); the current step is compared against all of them.
+        state["pol_hist"] = []
         if p.dim() >= 2:
             state["exp_avg_sq_row"] = torch.zeros(
                 p.shape[:-1], dtype=p.dtype, device=p.device
@@ -298,39 +336,61 @@ class Automagic3(torch.optim.Optimizer):
         # max-norm trust region) so no single weight can take an outsized step.
         update.clamp_(-group["clip_threshold"], group["clip_threshold"])
 
-        # Per-row sign agreement vs. the previous step: the plain fraction of
+        # Window-consistency agreement. We keep the last polarity_history_count
-        # elements that kept their sign, every element counting equally
+        # update-sign snapshots; over the window of those plus the current sign,
-        # regardless of update magnitude. For >=2D params this is reduced to one
+        # the per-element agreement is the fraction sharing the *dominant* sign,
-        # value per output channel; for 1D it is per element. The result drives
+        # max(p, 1 - p) == 0.5 + |p - 0.5| where p is the fraction positive. This
-        # a proportional lr bump in [-1, 1]. The match is summed straight off the
+        # is direction-agnostic (a consistently-negative element scores as high
-        # bool tensor (no full-size float cast) into the per-row reduction.
+        # as a consistently-positive one) and the current step is just one vote,
+        # so a single flip against a long consistent history barely moves it
+        # (e.g. 20-of-21 the same -> ~0.95). Values lie in [0.5, 1]: 0.5 is a
+        # 50/50 (chaotic) element, 1.0 is perfectly consistent. Reduced to one
+        # value per output channel for >=2D, per element for 1D.
         cur_polarity = update > 0
-        eqb = cur_polarity == state["last_polarity"]
+        pol_hist = state["pol_hist"]
-        state["last_polarity"] = cur_polarity
 
         lr_t = state["lr"]
         if p.dim() >= 2:
             dims = tuple(range(1, p.dim()))
-            agreement = eqb.sum(dim=dims, dtype=torch.float32).div_(
-                eqb.shape[1:].numel()
-            )
             lr_b = lr_t.view(lr_t.shape[0], *([1] * (p.dim() - 1)))
         else:
-            agreement = eqb.to(torch.float32)
+            dims = None
             lr_b = lr_t
 
-        if state["step"] > 0:
+        if pol_hist:
-            direction = agreement.mul_(2.0).sub_(1.0)
+            # per-element fraction positive over the window (current + history)
+            win = len(pol_hist) + 1
+            frac = cur_polarity.to(torch.float32)
+            for h in pol_hist:
+                frac.add_(h)
+            frac.div_(win)
+            # dominant-sign fraction in [0.5, 1], then reduce to per-row
+            consist = frac.sub_(0.5).abs_().add_(0.5)
+            agreement = consist.mean(dim=dims) if dims is not None else consist
+
+            # Map consistency to a direction in [-1, 1], measured relative to the
+            # window's noise floor b (the consistency a pure-noise row shows by
+            # chance, computed from the window size): agreement == b -> 0 (hold),
+            # 1.0 (perfect) -> +1 (lr up), below b -> down. There is no manual
+            # target -- holding at b is fully automatic and self-balancing per
+            # row: the lr grows while a row is more consistent than chance and
+            # settles at whatever lr makes its consistency meet b, so noisy and
+            # clean layers each find their own operating point. Measuring
+            # relative to b also makes the signal window-size independent.
+            b = group["agreement_floor"]
+            direction = agreement.sub_(b).div_(1.0 - b).clamp_(-1.0, 1.0)
             # Multiplicative (geometric) bump: lr *= exp(direction * lr_bump_rate).
-            # A full up step multiplies lr by exp(lr_bump_rate), a full down step by
+            # lr_bump_rate is a fractional rate (~lr_bump_rate per step for small
-            # its reciprocal, so up and down are symmetric in log space (a full
+            # values), giving a uniform relative move at every scale across
-            # up bump is exactly undone by a full down bump). lr_bump_rate is thus a
+            # [min_lr, max_lr] instead of a fixed absolute amount.
-            # fractional rate (~lr_bump_rate per step for small values), giving a
-            # uniform relative move at every scale across [min_lr, max_lr]
-            # instead of a fixed absolute amount.
             lr_t.mul_(torch.exp(direction.mul_(group["lr_bump_rate"]))).clamp_(
                 min=group["min_lr"], max=group["max_lr"]
             )
+
+        # Record this step's polarity and trim the window to the configured size.
+        pol_hist.append(cur_polarity)
+        if len(pol_hist) > group["polarity_history_count"]:
+            del pol_hist[0]
         state["step"] += 1
 
         wd = group["weight_decay"]
@@ -406,3 +466,7 @@ class Automagic3(torch.optim.Optimizer):
                 st = self.state.get(p)
                 if st is not None and isinstance(st.get("lr"), torch.Tensor):
                     st["lr"] = st["lr"].to(torch.float32)
+                # Polarity history is transient; rebuild it after load rather
+                # than trying to persist/cast a list of bool tensors.
+                if st is not None and "pol_hist" in st:
+                    st["pol_hist"] = []