[optim] Add PolynomialLrSchedule and InvSqrtDecayLrSchedule

Add two new LR schedule types to Levanter's optim config for the LR
schedule sweep experiments described in issue #4082.

PolynomialLrSchedule wraps optax.polynomial_schedule with configurable
power parameter (power=1 linear, power=2 quadratic, power=0.5 sqrt).

InvSqrtDecayLrSchedule implements lr/sqrt(1+c*t/T) with configurable
decay constant, providing a schedule that never reaches zero.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: github-actions[bot]

lib/levanter/src/levanter/optim/config.py

@@ -78,6 +78,49 @@ def build(self, ctx: LrScheduleContext):
         return _inv_decay_schedule(ctx.learning_rate, ctx.min_lr, ctx.decay_steps)
 
 
+@LrSchedule.register_subclass("polynomial")
+@dataclass(frozen=True)
+class PolynomialLrSchedule(LrSchedule):
+    """Polynomial decay: lr * (1 - t/T)^power, reaching min_lr at step T.
+
+    Wraps optax.polynomial_schedule. Power=1 gives linear decay, power=2 gives
+    quadratic decay (drops LR more aggressively early), power=0.5 gives sqrt decay
+    (holds LR higher early, drops sharply at end).
+    """
+
+    power: float = 2.0
+
+    def build(self, ctx: LrScheduleContext):
+        return optax.polynomial_schedule(
+            init_value=ctx.learning_rate,
+            end_value=ctx.min_lr,
+            power=self.power,
+            transition_steps=ctx.decay_steps,
+        )
+
+
+@LrSchedule.register_subclass("inv_sqrt_decay")
+@dataclass(frozen=True)
+class InvSqrtDecayLrSchedule(LrSchedule):
+    """Inverse sqrt decay: lr / sqrt(1 + c * t).
+
+    Unlike InvSqrtLrSchedule (which uses a fixed timescale relative to warmup),
+    this schedule decays from peak LR using a configurable constant `c` that
+    controls how fast the LR drops. The LR never reaches zero.
+    """
+
+    decay_constant: float = 28.6
+
+    def build(self, ctx: LrScheduleContext):
+        c = self.decay_constant
+        lr = ctx.learning_rate
+
+        def schedule(count):
+            return lr / jnp.sqrt(1.0 + c * count / ctx.decay_steps)
+
+        return schedule
+
+
 @LrSchedule.register_subclass("power")
 @dataclass(frozen=True)
 class PowerLrSchedule(LrSchedule):

lib/levanter/tests/test_optimizer_config.py

@@ -248,6 +248,117 @@ def test_wsds_schedule_with_cycle_points():
     assert sched_fn(971) < 1e-3
 
 
+def test_polynomial_schedule_quadratic():
+    """Quadratic decay: (1-t)^2 shape via PolynomialLrSchedule."""
+    from levanter.optim.config import PolynomialLrSchedule
+
+    optimizer = AdamConfig(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        warmup=0.1,
+        min_lr_ratio=0.0,
+        lr_schedule=PolynomialLrSchedule(power=2.0),
+    )
+
+    sched_fn = optimizer.lr_scheduler(1000)
+
+    # Warmup phase
+    assert np.isclose(sched_fn(0), 0.0)
+    assert np.isclose(sched_fn(100), 1e-3)
+
+    # Decay phase: at midpoint (t=450 into 900-step decay), LR = 1e-3 * (1 - 450/900)^2 = 0.25e-3
+    assert np.isclose(sched_fn(550), 0.25e-3, atol=1e-6)
+
+    # End of decay
+    assert np.isclose(sched_fn(999), 0.0, atol=1e-5)
+
+
+def test_polynomial_schedule_linear():
+    """Power=1 should match linear decay."""
+    from levanter.optim.config import PolynomialLrSchedule
+
+    optimizer = AdamConfig(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        warmup=0.0,
+        min_lr_ratio=0.0,
+        lr_schedule=PolynomialLrSchedule(power=1.0),
+    )
+
+    sched_fn = optimizer.lr_scheduler(100)
+
+    assert np.isclose(sched_fn(0), 1e-3)
+    assert np.isclose(sched_fn(50), 0.5e-3, atol=1e-6)
+    assert np.isclose(sched_fn(100), 0.0, atol=1e-6)
+
+
+def test_polynomial_schedule_sqrt():
+    """Power=0.5 (sqrt decay) holds LR higher early, drops faster at end."""
+    from levanter.optim.config import PolynomialLrSchedule
+
+    optimizer = AdamConfig(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        warmup=0.0,
+        min_lr_ratio=0.0,
+        lr_schedule=PolynomialLrSchedule(power=0.5),
+    )
+
+    sched_fn = optimizer.lr_scheduler(100)
+
+    # At midpoint: (1-0.5)^0.5 ≈ 0.707
+    assert np.isclose(sched_fn(50), 1e-3 * 0.5**0.5, atol=1e-5)
+    assert np.isclose(sched_fn(100), 0.0, atol=1e-6)
+
+
+def test_polynomial_schedule_with_min_lr():
+    """Polynomial decay with a floor (min_lr_ratio > 0)."""
+    from levanter.optim.config import PolynomialLrSchedule
+
+    optimizer = AdamConfig(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        warmup=0.0,
+        min_lr_ratio=0.05,
+        lr_schedule=PolynomialLrSchedule(power=2.0),
+    )
+
+    sched_fn = optimizer.lr_scheduler(100)
+
+    # End of decay should reach min_lr = 0.05 * 1e-3
+    assert np.isclose(sched_fn(100), 0.05e-3, atol=1e-6)
+
+
+def test_inv_sqrt_decay_lr_schedule():
+    """InvSqrtDecayLrSchedule: lr / sqrt(1 + c * t / T)."""
+    from levanter.optim.config import InvSqrtDecayLrSchedule
+
+    optimizer = AdamConfig(
+        learning_rate=1e-3,
+        weight_decay=0.0,
+        warmup=0.0,
+        min_lr_ratio=0.0,
+        lr_schedule=InvSqrtDecayLrSchedule(decay_constant=28.6),
+    )
+
+    sched_fn = optimizer.lr_scheduler(1000)
+
+    # At t=0, lr = 1e-3 / sqrt(1) = 1e-3
+    assert np.isclose(sched_fn(0), 1e-3)
+
+    # Monotonically decreasing
+    assert sched_fn(100) < sched_fn(0)
+    assert sched_fn(500) < sched_fn(100)
+    assert sched_fn(999) < sched_fn(500)
+
+    # At t=T, lr = 1e-3 / sqrt(1 + 28.6) ≈ 1e-3 / 5.44 ≈ 0.000184
+    expected_end = 1e-3 / np.sqrt(1 + 28.6)
+    assert np.isclose(sched_fn(1000), expected_end, atol=1e-6)
+
+    # Never reaches zero
+    assert sched_fn(1000) > 0
+
+
 def test_warmup_longer_than_run_does_not_jump():
     optimizer = AdamConfig(
         learning_rate=3e-3,