github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-08 14:09:02 -06:00
parent 0d018db689
commit 5d6887fd98
1 changed files with 98 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
toolkit/optimizers/automagic3.py
View File

@@ -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
previous step. Agreement is the plain fraction of elements that kept their
sign (every element counts equally, regardless of update magnitude) and is
used as a proportional signal in [-1, 1] that scales a multiplicative
(geometric) bump ``lr *= exp(signal * lr_bump_rate)``, so each step is a fixed
fractional move that behaves uniformly across the whole [min_lr, max_lr]
range and a full up bump is exactly undone by a full down bump. The lr
drifts smoothly, self-stabilises near random agreement, and is clamped to
[min_lr, max_lr].
nudged by how consistent the per-element update direction has been.
Agreement is window-consistency: over the last ``polarity_history_count``
update-sign snapshots plus the current sign, the per-element score is the
fraction sharing the dominant sign (``max(p, 1-p)``, in [0.5, 1]; direction-
agnostic, and the current step is just one vote so a single flip against a
long consistent history barely moves it). It is reduced to one value per row.
``polarity_history_count`` is min 1 (default 2); a longer window makes the 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
(``2*agreement - 1`` in [-1, 1]). Plain English: the lr nudges gently
when the signal is weak and firmly when it is strong, and parks itself
instead of oscillating when gradients are basically noise.
scales the bump by how far consistency is from its hold point relative to
the noise floor (direction in [-1, 1]). Plain English: the lr nudges gently when the
signal is weak and firmly when it is strong, and parks itself instead of
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
# elements that kept their sign, every element counting equally
# regardless of update magnitude. For >=2D params this is reduced to one
# value per output channel; for 1D it is per element. The result drives
# a proportional lr bump in [-1, 1]. The match is summed straight off the
# bool tensor (no full-size float cast) into the per-row reduction.
# Window-consistency agreement. We keep the last polarity_history_count
# update-sign snapshots; over the window of those plus the current sign,
# the per-element agreement is the fraction sharing the *dominant* sign,
# max(p, 1 - p) == 0.5 + |p - 0.5| where p is the fraction positive. This
# is direction-agnostic (a consistently-negative element scores as high
# 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"]
state["last_polarity"] = cur_polarity
pol_hist = state["pol_hist"]
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:
direction = agreement.mul_(2.0).sub_(1.0)
if pol_hist:
# 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
# its reciprocal, so up and down are symmetric in log space (a full
# up bump is exactly undone by a full down bump). lr_bump_rate is thus a
# 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_bump_rate is a 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"] = []