Rework optimizers

core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala

@@ -1,12 +1,13 @@
 package dimwit.optimizer
 
 import dimwit.*
-
 import dimwit.Conversions.given
 import dimwit.tensortree.*
 import dimwit.tensortree.FloatTree.*
 import dimwit.tensortree.FloatTree.ops.*
 import dimwit.autodiff.*
+import dimwit.autodiff.Grad
+import dimwit.tensortree.IsFloatTree.given
 
 /** Gradient optimizer interface with functional state management.
   *
@@ -28,108 +29,116 @@ import dimwit.autodiff.*
   *    }}}
   */
 trait GradientOptimizer:
-  type State[_, V]
+  type State[_]
 
   // Core API
-  def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V]
-  def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: State[Params, V])(using IsFloating[V]): (Params, State[Params, V])
+  def init[Params: IsFloatTree](params: Params): State[Params]
+  def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: State[Params]): (Params, State[Params])
 
   // Convenience: iterator with fixed gradient function
-  def iterateWithState[V, Params: TensorTree: FloatTreeFor[V]](init: Params)(df: Params => Grad[Params])(using IsFloating[V]): Iterator[(Params, State[Params, V])] =
+  def iterateWithState[V, Params: IsFloatTree](init: Params)(df: Params => Grad[Params]): Iterator[(Params, State[Params])] =
     Iterator.iterate((init, this.init(init))): (params, state) =>
       val grads = df(params)
       update(grads, params, state)
 
-  def iterate[V, Params: TensorTree: FloatTreeFor[V]](init: Params)(df: Params => Grad[Params])(using IsFloating[V]): Iterator[Params] =
+  def iterate[V, Params: IsFloatTree](init: Params)(df: Params => Grad[Params]): Iterator[Params] =
     iterateWithState(init)(df).map(_._1)
 
-case class GradientDescent(learningRate: Double) extends GradientOptimizer:
+case class GradientDescentState[P](
+    step: Tensor0[Int32]
+)
+
+class GradientDescent(learningRate: Double | LearningRateSchedule) extends GradientOptimizer:
 
-  type State[P, V] = Unit // Stateless optimizer
+  private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate)
 
-  def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Unit = ()
+  type State[P] = GradientDescentState[P]
 
-  def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: Unit)(using IsFloating[V]): (Params, Unit) =
-    val newParams = params -- gradients.value.scale(learningRate)
-    (newParams, ())
+  def init[Params: IsFloatTree](params: Params): GradientDescentState[Params] =
+    GradientDescentState(step = 1)
 
-case class Lion(learningRate: Double, weightDecay: Double = 0.0f, beta1: Double = 0.9f, beta2: Double = 0.99f) extends GradientOptimizer:
+  def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: GradientDescentState[Params]): (Params, GradientDescentState[Params]) =
+    val newParams = params -- gradients.value.scale(learningRateF(state.step.item))
+    (newParams, GradientDescentState(state.step + 1))
+
+case class LionState[P](
+    momentums: P, // momentums
+    step: Tensor0[Int32]
+)
 
-  type State[P, V] = P // momentum state has same structure as params
+class Lion(learningRate: Double | LearningRateSchedule, val weightDecay: Double = 0.0f, val beta1: Double = 0.9f, val beta2: Double = 0.99f) extends GradientOptimizer:
 
-  def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Params =
-    params.map([T <: Tuple] =>
-      (n: Labels[T]) ?=>
-        (t: Tensor[T, V]) =>
-          Tensor(t.shape).fill(0f)
-    )
+  private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate)
 
-  def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, momentums: Params)(using IsFloating[V]): (Params, Params) =
+  type State[P] = LionState[P]
+
+  def init[Params: IsFloatTree](params: Params): LionState[Params] =
+    LionState(params.fillCopy(0f), step = 1)
+
+  def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: LionState[Params]): (Params, LionState[Params]) =
     // the direction (1 or -1)
     // is determined by the sign of the momentum + gradient
-    val updateDirection = (momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign
+    val updateDirection = (state.momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign
 
-    val updatedParams = params -- updateDirection.scale(learningRate) -- params.scale(weightDecay)
-    val newMomentums = momentums **! beta2 ++ gradients.value **! (1f - beta2)
+    val updatedParams = params -- updateDirection.scale(learningRateF(state.step.item)) -- params.scale(weightDecay)
+    val newMomentums = state.momentums **! beta2 ++ gradients.value **! (1f - beta2)
 
-    (updatedParams, newMomentums)
+    (updatedParams, LionState(newMomentums, state.step + 1))
 
-case class AdamState[P, V: IsFloating](
+case class AdamState[P](
     momentums: P, // momentums
     velocities: P, // velocities
-    b1: Tensor0[V], // decay rate for momentums mᵗ
-    b2: Tensor0[V] // decay rate for velocities vᵗ
+    step: Tensor0[Int32]
 )
 
 /** Implements the Adam optimization algorithm.
   *
   * @see [[https://arxiv.org/abs/1412.6980 Adam: A Method for Stochastic Optimization]]
   */
-case class Adam(
-    learningRate: Double, // step size (learning rate)
+class Adam(
+    learningRate: Double | LearningRateSchedule,
     b1: Double = 0.9, // decay rate for momentums mᵗ
     b2: Double = 0.999, // decay rate for velocities vᵗ
     epsilon: Double = 1e-8 // small constant to prevent division by zero
 ) extends GradientOptimizer:
 
+  private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate)
   private val β1 = b1
   private val β2 = b2
+  private val ε = epsilon
 
-  type State[P, V] = AdamState[P, V]
+  type State[P] = AdamState[P]
 
-  def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] =
+  def init[P: IsFloatTree](params: P): State[P] =
     def zeros = params.fillCopy(0f)
-    AdamState[Params, V](zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f))
+    AdamState(zeros, zeros, step = 1)
+
+  def learningRate[Params](state: State[Params]): Double = learningRateF(state.step.item)
 
-  def update[V, Params: TensorTree: FloatTreeFor[V]](
+  def update[Params: IsFloatTree](
       gradients: Grad[Params],
       params: Params,
-      state: State[Params, V]
-  )(using IsFloating[V]): (Params, State[Params, V]) =
+      state: State[Params]
+  ): (Params, State[Params]) =
     // rename state variables to last time step for clarity
     val `mₜ₋₁` = state.momentums
     val `vₜ₋₁` = state.velocities
-    val `β1ₜ₋₁` = state.b1
-    val `β2ₜ₋₁` = state.b2
 
     // rename parameters for internal clarity
-    val α = learningRate
-    val ε = epsilon
-    val `θₜ₋₁` = params
+    val α = learningRate(state)
 
-    // update moments for bias correction
-    val β1ₜ = `β1ₜ₋₁` * β1
-    val β2ₜ = `β2ₜ₋₁` * β2
+    val `θₜ₋₁` = params
 
     // Adam implementation
     val gₜ = gradients.value
     val mᵗ = `β1` **! `mₜ₋₁` ++ (1f - `β1`) **! gₜ
     val vᵗ = `β2` **! `vₜ₋₁` ++ (1f - `β2`) **! gₜ.pow(2)
+    val (β1ₜ, β2ₜ) = (math.pow(β1, state.step.item), math.pow(β2, state.step.item))
     val m̂ = mᵗ `//!` (1f - `β1ₜ`)
     val v̂ = vᵗ `//!` (1f - `β2ₜ`)
     val θₜ = `θₜ₋₁` -- (α **! m̂) `//` (v̂.sqrt ++! ε)
 
-    (θₜ, AdamState(mᵗ, vᵗ, β1ₜ, β2ₜ))
+    (θₜ, AdamState(mᵗ, vᵗ, state.step + 1))
 
 /** Implements the AdamW algorithm (Adam with decoupled weight decay).
   *
@@ -141,21 +150,21 @@ case class Adam(
   * @param learningRate The step size.
   * @param weightDecayFactor The coefficient for weight decay (lambda).
   */
-case class AdamW(
+class AdamW(
     val adam: Adam,
     val weightDecayFactor: Double
 ) extends GradientOptimizer:
 
-  type State[P, V] = adam.State[P, V]
+  type State[P] = adam.State[P]
 
-  def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] = adam.init(params)
+  def init[Params: IsFloatTree](params: Params): State[Params] = adam.init(params)
 
-  def update[V, Params: TensorTree: FloatTreeFor[V]](
+  def update[Params: IsFloatTree](
       gradients: Grad[Params],
       params: Params,
-      state: State[Params, V]
-  )(using IsFloating[V]): (Params, State[Params, V]) =
-    val α = adam.learningRate
+      state: State[Params]
+  ): (Params, State[Params]) =
+    val α = adam.learningRate(state)
     val `θₜ₋₁` = params
     val `λ'` = weightDecayFactor
     val λ = `λ'` * α // Tie weight decay to learning rate

core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala

@@ -0,0 +1,78 @@
+package dimwit.optimizer
+
+type LearningRateSchedule = Int => Double
+
+object LearningRateSchedule:
+  def apply(f: Int => Double): LearningRateSchedule = f
+
+  private[dimwit] def from(learningRate: Double | LearningRateSchedule): LearningRateSchedule =
+    learningRate match
+      case f: LearningRateSchedule => f
+      case d: Double               => _ => d
+
+  extension (s: LearningRateSchedule)
+
+    /** Shifts a schedule forward in time by a specified number of steps.
+      *
+      * For all `t < steps`, the schedule evaluates as if `t = 0`, effectively locking
+      * the learning rate at its initial starting value until the delay has passed.
+      *
+      * @param steps The number of iterations to delay the schedule's progression.
+      * @return A time-shifted schedule.
+      */
+    def delay(steps: Int): LearningRateSchedule = t =>
+      val shiftedT = math.max(t - steps, 0)
+      s(shiftedT)
+
+  /** Combines multiple schedules by taking their pointwise minimum.
+    *
+    * At any given step `t`, this evaluates all provided schedules and returns
+    * the lowest learning rate. This acts as a mathematical lower envelope,
+    * seamlessly handing off from one curve to another when they cross.
+    * It is highly useful for safely composing full schedules, such as capping
+    * a delayed decay curve with a linear warmup phase.
+    *
+    * @param schedules A variable number of schedules to evaluate concurrently.
+    * @return A composite schedule that yields the lowest value across all input schedules at step `t`.
+    * @throws java.lang.UnsupportedOperationException if no schedules are provided.
+    */
+  def pointwiseMin(schedules: LearningRateSchedule*): LearningRateSchedule = t =>
+    schedules.map(s => s(t)).min
+
+  /** Creates a schedule that rises linearly from a fraction of `maxLr` up to `maxLr`.
+    *
+    * At `t = 0`, the learning rate starts slightly above zero (`maxLr / (warmupSteps + 1)`).
+    * It reaches exactly `maxLr` at `t = warmupSteps`, and remains locked at `maxLr` for all subsequent steps.
+    *
+    * @param maxLr       The peak learning rate reached at the end of the warmup.
+    * @param warmupSteps The number of steps over which the learning rate increases.
+    * @return A linear warmup schedule.
+    */
+  def linearWarmup(
+      maxLr: Float,
+      warmupSteps: Int
+  ): LearningRateSchedule = t =>
+    val warmupRatio = math.min((t + 1f) / (warmupSteps + 1f), 1f)
+    maxLr * warmupRatio
+
+  /** Creates a schedule that decays from `maxLr` down to `minLr` following a half-cosine curve.
+    *
+    * This schedule has no concept of warmup; it begins decaying immediately at `t = 0`.
+    * Once `t >= decaySteps`, the learning rate locks permanently at `minLr`.
+    *
+    * @param maxLr      The initial maximum learning rate at `t = 0`.
+    * @param minLr      The final baseline learning rate to reach after decaying.
+    * @param decaySteps The number of steps over which to apply the decay curve.
+    * @return A cosine decay schedule.
+    * @throws java.lang.IllegalArgumentException if `decaySteps` is zero or negative.
+    */
+  def cosineDecay(
+      maxLr: Float,
+      minLr: Float,
+      decaySteps: Int
+  ): LearningRateSchedule =
+    require(decaySteps > 0, "decaySteps must be strictly positive to avoid division by zero")
+    t =>
+      val decayRatio = math.min(t / decaySteps.toFloat, 1f)
+      val coeff = 0.5f * (1.0f + math.cos(math.Pi.toFloat * decayRatio))
+      minLr + coeff * (maxLr - minLr)

core/src/main/scala/dimwit/tensortree/FloatTree.scala

@@ -142,3 +142,24 @@ object FloatTree:
         p.map([T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, V]) => a.asFloat(vtype)).asInstanceOf[F[NewV]]
 
 type FloatTreeFor[V] = [P] =>> FloatTree[P, V]
+
+/** A typeclass that proves P is a FloatTree, hiding the specific float type V
+  * from method signatures while keeping the evidence available.
+  */
+trait IsFloatTree[P]:
+  type V
+  given isFloating: IsFloating[V]
+  given floatTree: FloatTree[P, V]
+  given tensorTree: TensorTree[P]
+
+object IsFloatTree:
+  // The compiler automatically packages a FloatTree[P, V] into an IsFloatTree[P]
+  given pack[P, V0](using ft: FloatTree[P, V0], tt: TensorTree[P], isF: IsFloating[V0]): IsFloatTree[P] with
+    type V = V0
+    val isFloating = isF
+    val floatTree = ft
+    val tensorTree = tt
+
+  given unpackFloatTree[P](using isFT: IsFloatTree[P]): FloatTree[P, isFT.V] = isFT.floatTree
+  given unpackIsFloating[P](using isFT: IsFloatTree[P]): IsFloating[isFT.V] = isFT.isFloating
+  given unpackTensorTree[P](using isFT: IsFloatTree[P]): TensorTree[P] = isFT.tensorTree

core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala

@@ -14,6 +14,28 @@ class GradientOptimizerSuite extends DimwitTest:
       val minX = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
       minX.item shouldBe -1.0f +- 0.1f
 
+    it("should work with lr schedule"):
+      def learningRateSchedule(step: Int): Double =
+        if step <= 5 then 0.0
+        else 0.1
+      val optimizer = GradientDescent(learningRate = learningRateSchedule)
+      val x1 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(5).next()
+      x1.item shouldBe 2.0f +- 0.1
+      val x2 = optimizer.iterate(Tensor0(2.0f))(x => Grad(2 * (x + 1))).drop(1000).next()
+      x2.item shouldBe -1.0f +- 0.1f
+
+    it("should compute the step (single step)"):
+      val optimizer = GradientDescent(learningRate = 0.1)
+      val initParams = Tensor0(2.0f)
+      val initState = optimizer.init(initParams)
+
+      val grad = Grad(Tensor0(6.0f))
+      val (nextParams, nextState) = optimizer.update(grad, initParams, initState)
+
+      initState.step.item shouldBe 1
+      nextParams.item shouldBe 1.4f +- 1e-5f
+      nextState.step.item shouldBe 2
+
   describe("Adam"):
     it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"):
       val optimizer = Adam(learningRate = 0.1)
@@ -31,8 +53,7 @@ class GradientOptimizerSuite extends DimwitTest:
       nextParams.item shouldBe 1.9f +- 1e-5f
       nextState.momentums.item shouldBe 0.6f +- 1e-5f
       nextState.velocities.item shouldBe 0.036f +- 1e-5f
-      nextState.b1.item shouldBe 0.9f +- 1e-5f
-      nextState.b2.item shouldBe 0.999f +- 1e-5f
+      nextState.step.asInt32.item shouldBe 2
 
   describe("AdamW"):
     it("should converge towards the minimum of f(x) = (x+1)^2 at x = -1"):
@@ -42,14 +63,14 @@ class GradientOptimizerSuite extends DimwitTest:
 
     it("should apply decoupled weight decay (single step)"):
       val adam = Adam(learningRate = 0.1)
-      val adamW = AdamW(adam, weightDecayFactor = 0.1)
+      val adamW = AdamW(Adam(learningRate = 0.1), weightDecayFactor = 0.1)
 
-      val initParams = Tensor0(2.0)
-      val grad = Grad(Tensor0(6.0))
+      val initParams = Tensor0(2.0f)
+      val grad = Grad(Tensor0(6.0f))
       val (adamParams, _) = adam.update(grad, initParams, adam.init(initParams))
       val (adamWParams, _) = adamW.update(grad, initParams, adamW.init(initParams))
 
-      adamWParams.item shouldBe (adamParams.item - 0.02) +- 1e-5
+      adamWParams.item shouldBe (adamParams.item - 0.02f) +- 1e-5f
 
   describe("Lion"):
     it("should converge towards the minimum of f(x) = (x+1)^2"):
@@ -59,11 +80,12 @@ class GradientOptimizerSuite extends DimwitTest:
 
     it("should compute the exact sign-based update and momentum (single step)"):
       val optimizer = Lion(learningRate = 0.1, beta1 = 0.9, beta2 = 0.99)
-      val initParams = Tensor0(2.0)
+      val initParams = Tensor0(2.0f)
       val initMomentum = optimizer.init(initParams)
 
-      val grad = Grad(Tensor0(6.0))
-      val (nextParams, nextMomentum) = optimizer.update(grad, initParams, initMomentum)
+      val grad = Grad(Tensor0(6.0f))
+      val (nextParams, nextState) = optimizer.update(grad, initParams, initMomentum)
 
-      nextParams.item shouldBe 1.9d +- 1e-5d
-      nextMomentum.item shouldBe 0.06d +- 1e-5d
+      nextParams.item shouldBe 1.9f +- 1e-5f
+      nextState.momentums.item shouldBe 0.06f +- 1e-5f
+      nextState.step.item shouldBe 2