diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 9915b59..a3f8464 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/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 diff --git a/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala b/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala new file mode 100644 index 0000000..982b5b1 --- /dev/null +++ b/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) diff --git a/core/src/main/scala/dimwit/tensortree/FloatTree.scala b/core/src/main/scala/dimwit/tensortree/FloatTree.scala index 4b3cd99..e68663d 100644 --- a/core/src/main/scala/dimwit/tensortree/FloatTree.scala +++ b/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 diff --git a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala index dc8e2fb..3a1e71f 100644 --- a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala +++ b/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 diff --git a/examples/src/main/scala/complex/VariationalAutoencoder.scala b/examples/src/main/scala/complex/VariationalAutoencoder.scala index e5b553f..470ce0f 100644 --- a/examples/src/main/scala/complex/VariationalAutoencoder.scala +++ b/examples/src/main/scala/complex/VariationalAutoencoder.scala @@ -210,33 +210,34 @@ object VariationalAutoencoderExample: val batches = trainImages.chunk(Axis[TrainSample], numSamples / batchSize) val optimizer = GradientDescent(learningRate = learningRate) - def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params): Params = + def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) = val grads = grad(batchLoss(trainKey, batch))(params) - val (newParams, _) = optimizer.update(grads, params, ()) - newParams + val (newParams, newState) = optimizer.update(grads, params, state) + (newParams, newState) val (jitDonate, jitStep, jitReclaim) = jitDonating(trainBatch) - def trainEpoch(key: Random.Key, epoch: Int, params: Params): Params = + def trainEpoch(key: Random.Key, epoch: Int, params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) = val batchKeys = key.split(batches.size) jitReclaim( - batches.zip(batchKeys).foldLeft(jitDonate(params)): - case (batchParams, (batch, key)) => - jitStep(key, batch, batchParams) + batches.zip(batchKeys).foldLeft(jitDonate(params, state)): + case ((batchParams, state), (batch, key)) => + jitStep(key, batch, batchParams, state) ) val keysForEpochs = dataKey.split(numEpochs) val initialParams = Params(encoderParams, decoderParams).map([T <: Tuple] => (n: Labels[T]) ?=> (t: Tensor[T, Float32]) => t *! 0.1f) + val initState = optimizer.init(initialParams) - val trainedParams = (0 until numEpochs).foldLeft(initialParams): - case (params, epoch) => + val (trainedParams, _) = (0 until numEpochs).foldLeft(initialParams, initState): + case ((params, state), epoch) => timed(s"Evaluation $epoch/$numEpochs"): val lossValue = batchLoss(keysForEpochs(epoch), testImages)(params) println(s"Test loss in epoch $epoch: $lossValue") timed(s"Training $epoch/$numEpochs"): dimwit.gc() - trainEpoch(keysForEpochs(epoch), epoch, params) + trainEpoch(keysForEpochs(epoch), epoch, params, state) /* * Evaluation