From 7e1c38205b44cb67ea7e3a1ad0fa9a2ff24d2b5b Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Sat, 1 Aug 2026 16:51:02 +0200 Subject: [PATCH 1/4] move tensor tree related methods to separate package --- .../dimwit/optimizer/GradientOptimizer.scala | 1 - .../scala/dimwit/tensortree/PyTreeSuite.scala | 83 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 9915b59..23ff702 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -1,7 +1,6 @@ package dimwit.optimizer import dimwit.* - import dimwit.Conversions.given import dimwit.tensortree.* import dimwit.tensortree.FloatTree.* diff --git a/core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala b/core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala new file mode 100644 index 0000000..d1407a3 --- /dev/null +++ b/core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala @@ -0,0 +1,83 @@ +package dimwit.tensortree + +import dimwit.* +import dimwit.jax.Jax +import me.shadaj.scalapy.py +class ToPyTreeSuite extends DimwitTest: + + describe("TensorTree Identity (fromPyTree(toPyTree(x)) == x)"): + + it("1-level case class"): + case class Params( + val w: Tensor1[A, Float32], + val b: Tensor0[Float32] + ) + val params = Params( + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), + Tensor0(0.5f) + ) + + val tc = TensorTree[Params] + val reconstructed = tc.fromPyTree(tc.toPyTree(params)) + + reconstructed.w should approxEqual(params.w) + reconstructed.b should approxEqual(params.b) + + it("2-level case class"): + case class LayerParams( + val w: Tensor2[A, B, Float32], + val b: Tensor0[Float32] + ) + case class ModelParams( + val layer1: LayerParams, + val layer2: LayerParams + ) + + val params = ModelParams( + LayerParams( + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f))), + Tensor0(0.25f) + ), + LayerParams( + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.5f, 0.6f), Array(0.7f, 0.8f))), + Tensor0(0.75f) + ) + ) + + val tc = TensorTree[ModelParams] + val reconstructed = tc.fromPyTree(tc.toPyTree(params)) + + reconstructed.layer1.w should approxEqual(params.layer1.w) + reconstructed.layer1.b should approxEqual(params.layer1.b) + reconstructed.layer2.w should approxEqual(params.layer2.w) + reconstructed.layer2.b should approxEqual(params.layer2.b) + + it("tuple"): + val myTuple = ( + Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), + Tensor0(0.5f) + ) + + val tc = TensorTree[(Tensor1[A, Float32], Tensor0[Float32])] + val reconstructed = tc.fromPyTree(tc.toPyTree(myTuple)) + + reconstructed._1 should approxEqual(myTuple._1) + reconstructed._2 should approxEqual(myTuple._2) + + it("case class with list"): + case class Params( + val layerWeights: List[Tensor2[A, B, Float32]] + ) + val params = Params( + List( + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f))), + Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.1f, 1.2f), Array(1.3f, 1.4f))) + ) + ) + + val tc = TensorTree[Params] + val reconstructed = tc.fromPyTree(tc.toPyTree(params)) + + reconstructed.layerWeights.size shouldBe params.layerWeights.size + reconstructed.layerWeights(0) should approxEqual(params.layerWeights(0)) + reconstructed.layerWeights(1) should approxEqual(params.layerWeights(1)) From d35ad7f8614fb18ebbaa3c5dba310883d5cc641d Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 4 Aug 2026 08:42:11 +0200 Subject: [PATCH 2/4] Rework optimizers * Remove beta1 and beta2 from AdamState * Remove generic V from all optimizers, replace with IsFloatingTree abstraction. * Add learning rate schedule --- .../dimwit/optimizer/GradientOptimizer.scala | 181 +++++++++++++----- .../scala/dimwit/tensortree/FloatTree.scala | 21 ++ .../optimizer/GradientOptimizerSuite.scala | 19 +- 3 files changed, 166 insertions(+), 55 deletions(-) diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 23ff702..9caccdc 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -6,6 +6,50 @@ import dimwit.tensortree.* import dimwit.tensortree.FloatTree.* import dimwit.tensortree.FloatTree.ops.* import dimwit.autodiff.* +import dimwit.autodiff.Grad +import dimwit.tensortree.IsFloatTree.given + +/** A high-performance, stateful generator for sequential mathematical series. + * + * ==Idea== + * This trait bridges a purely functional interface (`Long => R`) with a highly + * optimized, mutable internal implementation. It allows implementing classes + * to compute iterative sequences using raw primitive mutation in `doStep()`. + * + * ==Example== + * For example exponential decay can be represented as a sequence on R without + * sacrificing performance or relying on expensive operations like `math.pow`: + * {{{ + * class ExponentialDecay(rate: Double) extends SequenceFunction[Double]: + * private var value = 1.0 + * protected def doStep(): Unit = value *= rate + * protected def state(): Double = value + * }}} + * while allowing mathematically valid abstractions by exposing this stateful logic + * as a simple, pure `Long => Double` function to the consuming optimizer. + * + * ==Restriction== + * A SequenceFunction can only be used if we can guarantee calling with increasing step. + * + * @tparam R The type of the state. + */ +trait SequenceFunction[R] extends (Int => R): + + private var currentStep: Int = 0 + + protected def doStep(): Unit + protected def state(): R + + /** Override this to provide context-specific debugging information. */ + protected def outOfOrderErrorMessage(step: Long, current: Long): String = + s"Misuse of SequenceFunction, step should never become smaller: $step > $current" + + def apply(step: Int): R = + assert(step > currentStep, outOfOrderErrorMessage(step, currentStep)) + while currentStep < step do + doStep() + currentStep += 1 + state() /** Gradient optimizer interface with functional state management. * @@ -27,43 +71,39 @@ 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: +class GradientDescent(learningRate: Double) extends GradientOptimizer: - type State[P, V] = Unit // Stateless optimizer + type State[P] = Unit // Stateless optimizer - def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Unit = () + def init[Params: IsFloatTree](params: Params): Unit = () - def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: Unit)(using IsFloating[V]): (Params, Unit) = + def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: Unit): (Params, Unit) = val newParams = params -- gradients.value.scale(learningRate) (newParams, ()) -case class Lion(learningRate: Double, weightDecay: Double = 0.0f, beta1: Double = 0.9f, beta2: Double = 0.99f) extends GradientOptimizer: +class Lion(val learningRate: Double, val weightDecay: Double = 0.0f, val beta1: Double = 0.9f, val beta2: Double = 0.99f) extends GradientOptimizer: - type State[P, V] = P // momentum state has same structure as params + type State[P] = P // momentum state has same structure as params - 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) - ) + def init[Params: IsFloatTree](params: Params): Params = + params.fillCopy(0f) - def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, momentums: Params)(using IsFloating[V]): (Params, Params) = + def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, momentums: Params): (Params, Params) = // the direction (1 or -1) // is determined by the sign of the momentum + gradient val updateDirection = (momentums **! beta1 ++ gradients.value **! (1f - beta1)).sign @@ -73,62 +113,74 @@ case class Lion(learningRate: Double, weightDecay: Double = 0.0f, beta1: Double (updatedParams, newMomentums) -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 | (Int => Double), 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 = learningRate match + case f: (Int => Double) => f + case d: Double => _ => d private val β1 = b1 private val β2 = b2 + private val ε = epsilon + + private val betaBiasCorrection: Int => (Double, Double) = new SequenceFunction[(Double, Double)]: + private var β1ₜ = 1d; private var β2ₜ = 1d + def doStep(): Unit = β1ₜ *= β1; β2ₜ *= β2 + def state(): (Double, Double) = (β1ₜ, β2ₜ) + + override protected def outOfOrderErrorMessage(step: Long, current: Long): String = + s"Adam optimizer step moved backwards or stalled (requested step $step, but currently at $current). " + + "Because Adam maintains internal chronological state for bias correction, a single instance must progress strictly forward in time. " + + "This error typically occurs if you:\n" + + " 1. Share the same Adam instance across parallel/concurrent training threads.\n" + + "Solution: Instantiate a fresh Adam optimizer for each independent training loop or worker." - 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 update[V, Params: TensorTree: FloatTreeFor[V]]( + def learningRate[Params](state: State[Params]): Double = learningRateF(state.step.item) + + 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ₜ) = betaBiasCorrection(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). * @@ -140,24 +192,63 @@ 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 val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁` val (θₜ, adamState) = adam.update(gradients, decayedParams, state) (θₜ, adamState) + +object LearningRateSchedules: + + /** A schedule maps the current iteration tensor to a learning rate tensor */ + type Schedule = Int => Double + + extension (s: Schedule) + /** Mathematical intersection: seamlessly hands off when one curve crosses the other */ + infix def min(other: Schedule): Schedule = t => math.min(s(t), other(t)) + + /** Shifts a schedule forward in time. + * For all t < steps, the schedule sees t = 0 (locking it at its initial value). + */ + def delay(steps: Int): Schedule = t => + val shiftedT = math.max(t - steps, 0) + s(shiftedT) + + /** Rises linearly from near 0 up to maxLr over `warmupSteps`, then locks at maxLr. + */ + def linearWarmup( + maxLr: Float, + warmupSteps: Int + ): Schedule = t => + val warmupRatio = math.min((t + 1f) / (warmupSteps + 1f), 1f) + maxLr * warmupRatio + + /** Starts at maxLr and cosine decays down to minLr over `decaySteps`, then locks at minLr. + * (Has no concept of warmup; assumes it starts decaying at t=0). + */ + def cosineDecay( + maxLr: Float, + minLr: Float, + decaySteps: Int + ): Schedule = + 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..e62411f 100644 --- a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala +++ b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala @@ -31,8 +31,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 +41,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 +58,11 @@ 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 grad = Grad(Tensor0(6.0f)) val (nextParams, nextMomentum) = 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 + nextMomentum.item shouldBe 0.06f +- 1e-5f From aaf493f6a553a4421d5530b827bf483ab4a9004f Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 4 Aug 2026 08:11:52 +0200 Subject: [PATCH 3/4] Cleanup Optimizers: Remove SequenceFunction (until proven necessary), move LRSchedule in custom file --- .../dimwit/optimizer/GradientOptimizer.scala | 101 +----------------- .../optimizer/LearningRateSchedule.scala | 73 +++++++++++++ .../scala/dimwit/tensortree/PyTreeSuite.scala | 83 -------------- 3 files changed, 77 insertions(+), 180 deletions(-) create mode 100644 core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala delete mode 100644 core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 9caccdc..baca4f1 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -9,48 +9,6 @@ import dimwit.autodiff.* import dimwit.autodiff.Grad import dimwit.tensortree.IsFloatTree.given -/** A high-performance, stateful generator for sequential mathematical series. - * - * ==Idea== - * This trait bridges a purely functional interface (`Long => R`) with a highly - * optimized, mutable internal implementation. It allows implementing classes - * to compute iterative sequences using raw primitive mutation in `doStep()`. - * - * ==Example== - * For example exponential decay can be represented as a sequence on R without - * sacrificing performance or relying on expensive operations like `math.pow`: - * {{{ - * class ExponentialDecay(rate: Double) extends SequenceFunction[Double]: - * private var value = 1.0 - * protected def doStep(): Unit = value *= rate - * protected def state(): Double = value - * }}} - * while allowing mathematically valid abstractions by exposing this stateful logic - * as a simple, pure `Long => Double` function to the consuming optimizer. - * - * ==Restriction== - * A SequenceFunction can only be used if we can guarantee calling with increasing step. - * - * @tparam R The type of the state. - */ -trait SequenceFunction[R] extends (Int => R): - - private var currentStep: Int = 0 - - protected def doStep(): Unit - protected def state(): R - - /** Override this to provide context-specific debugging information. */ - protected def outOfOrderErrorMessage(step: Long, current: Long): String = - s"Misuse of SequenceFunction, step should never become smaller: $step > $current" - - def apply(step: Int): R = - assert(step > currentStep, outOfOrderErrorMessage(step, currentStep)) - while currentStep < step do - doStep() - currentStep += 1 - state() - /** Gradient optimizer interface with functional state management. * * This API provides the following two styles of usage: @@ -124,31 +82,19 @@ case class AdamState[P]( * @see [[https://arxiv.org/abs/1412.6980 Adam: A Method for Stochastic Optimization]] */ class Adam( - learningRate: Double | (Int => Double), + 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 = learningRate match - case f: (Int => Double) => f - case d: Double => _ => d + case f: LearningRateSchedule => f + case d: Double => _ => d private val β1 = b1 private val β2 = b2 private val ε = epsilon - private val betaBiasCorrection: Int => (Double, Double) = new SequenceFunction[(Double, Double)]: - private var β1ₜ = 1d; private var β2ₜ = 1d - def doStep(): Unit = β1ₜ *= β1; β2ₜ *= β2 - def state(): (Double, Double) = (β1ₜ, β2ₜ) - - override protected def outOfOrderErrorMessage(step: Long, current: Long): String = - s"Adam optimizer step moved backwards or stalled (requested step $step, but currently at $current). " + - "Because Adam maintains internal chronological state for bias correction, a single instance must progress strictly forward in time. " + - "This error typically occurs if you:\n" + - " 1. Share the same Adam instance across parallel/concurrent training threads.\n" + - "Solution: Instantiate a fresh Adam optimizer for each independent training loop or worker." - type State[P] = AdamState[P] def init[P: IsFloatTree](params: P): State[P] = @@ -175,7 +121,7 @@ class Adam( val gₜ = gradients.value val mᵗ = `β1` **! `mₜ₋₁` ++ (1f - `β1`) **! gₜ val vᵗ = `β2` **! `vₜ₋₁` ++ (1f - `β2`) **! gₜ.pow(2) - val (β1ₜ, β2ₜ) = betaBiasCorrection(state.step.item) + 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 ++! ε) @@ -213,42 +159,3 @@ class AdamW( val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁` val (θₜ, adamState) = adam.update(gradients, decayedParams, state) (θₜ, adamState) - -object LearningRateSchedules: - - /** A schedule maps the current iteration tensor to a learning rate tensor */ - type Schedule = Int => Double - - extension (s: Schedule) - /** Mathematical intersection: seamlessly hands off when one curve crosses the other */ - infix def min(other: Schedule): Schedule = t => math.min(s(t), other(t)) - - /** Shifts a schedule forward in time. - * For all t < steps, the schedule sees t = 0 (locking it at its initial value). - */ - def delay(steps: Int): Schedule = t => - val shiftedT = math.max(t - steps, 0) - s(shiftedT) - - /** Rises linearly from near 0 up to maxLr over `warmupSteps`, then locks at maxLr. - */ - def linearWarmup( - maxLr: Float, - warmupSteps: Int - ): Schedule = t => - val warmupRatio = math.min((t + 1f) / (warmupSteps + 1f), 1f) - maxLr * warmupRatio - - /** Starts at maxLr and cosine decays down to minLr over `decaySteps`, then locks at minLr. - * (Has no concept of warmup; assumes it starts decaying at t=0). - */ - def cosineDecay( - maxLr: Float, - minLr: Float, - decaySteps: Int - ): Schedule = - 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/optimizer/LearningRateSchedule.scala b/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala new file mode 100644 index 0000000..bed3e03 --- /dev/null +++ b/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala @@ -0,0 +1,73 @@ +package dimwit.optimizer + +type LearningRateSchedule = Int => Double + +object LearningRateSchedule: + def apply(f: Int => Double): LearningRateSchedule = f + + 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/test/scala/dimwit/tensortree/PyTreeSuite.scala b/core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala deleted file mode 100644 index d1407a3..0000000 --- a/core/src/test/scala/dimwit/tensortree/PyTreeSuite.scala +++ /dev/null @@ -1,83 +0,0 @@ -package dimwit.tensortree - -import dimwit.* -import dimwit.jax.Jax -import me.shadaj.scalapy.py -class ToPyTreeSuite extends DimwitTest: - - describe("TensorTree Identity (fromPyTree(toPyTree(x)) == x)"): - - it("1-level case class"): - case class Params( - val w: Tensor1[A, Float32], - val b: Tensor0[Float32] - ) - val params = Params( - Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), - Tensor0(0.5f) - ) - - val tc = TensorTree[Params] - val reconstructed = tc.fromPyTree(tc.toPyTree(params)) - - reconstructed.w should approxEqual(params.w) - reconstructed.b should approxEqual(params.b) - - it("2-level case class"): - case class LayerParams( - val w: Tensor2[A, B, Float32], - val b: Tensor0[Float32] - ) - case class ModelParams( - val layer1: LayerParams, - val layer2: LayerParams - ) - - val params = ModelParams( - LayerParams( - Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f))), - Tensor0(0.25f) - ), - LayerParams( - Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.5f, 0.6f), Array(0.7f, 0.8f))), - Tensor0(0.75f) - ) - ) - - val tc = TensorTree[ModelParams] - val reconstructed = tc.fromPyTree(tc.toPyTree(params)) - - reconstructed.layer1.w should approxEqual(params.layer1.w) - reconstructed.layer1.b should approxEqual(params.layer1.b) - reconstructed.layer2.w should approxEqual(params.layer2.w) - reconstructed.layer2.b should approxEqual(params.layer2.b) - - it("tuple"): - val myTuple = ( - Tensor1(Axis[A]).fromArray(Array(0.1f, 0.2f, 0.3f)), - Tensor0(0.5f) - ) - - val tc = TensorTree[(Tensor1[A, Float32], Tensor0[Float32])] - val reconstructed = tc.fromPyTree(tc.toPyTree(myTuple)) - - reconstructed._1 should approxEqual(myTuple._1) - reconstructed._2 should approxEqual(myTuple._2) - - it("case class with list"): - case class Params( - val layerWeights: List[Tensor2[A, B, Float32]] - ) - val params = Params( - List( - Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(0.1f, 0.2f), Array(0.3f, 0.4f))), - Tensor2(Axis[A], Axis[B]).fromArray(Array(Array(1.1f, 1.2f), Array(1.3f, 1.4f))) - ) - ) - - val tc = TensorTree[Params] - val reconstructed = tc.fromPyTree(tc.toPyTree(params)) - - reconstructed.layerWeights.size shouldBe params.layerWeights.size - reconstructed.layerWeights(0) should approxEqual(params.layerWeights(0)) - reconstructed.layerWeights(1) should approxEqual(params.layerWeights(1)) From 4db980d16917079eaf1d9db2446ddea0cdfb45cd Mon Sep 17 00:00:00 2001 From: Benjamin Meyer Date: Tue, 4 Aug 2026 09:05:17 +0200 Subject: [PATCH 4/4] Add LearningRateSchedule to other optimizers. --- .../dimwit/optimizer/GradientOptimizer.scala | 48 ++++++++++++------- .../optimizer/LearningRateSchedule.scala | 5 ++ .../optimizer/GradientOptimizerSuite.scala | 27 ++++++++++- .../complex/VariationalAutoencoder.scala | 21 ++++---- 4 files changed, 71 insertions(+), 30 deletions(-) diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index baca4f1..a3f8464 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -44,32 +44,46 @@ trait GradientOptimizer: def iterate[V, Params: IsFloatTree](init: Params)(df: Params => Grad[Params]): Iterator[Params] = iterateWithState(init)(df).map(_._1) -class GradientDescent(learningRate: Double) extends GradientOptimizer: +case class GradientDescentState[P]( + step: Tensor0[Int32] +) + +class GradientDescent(learningRate: Double | LearningRateSchedule) extends GradientOptimizer: + + private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate) + + type State[P] = GradientDescentState[P] - type State[P] = Unit // Stateless optimizer + def init[Params: IsFloatTree](params: Params): GradientDescentState[Params] = + GradientDescentState(step = 1) - def init[Params: IsFloatTree](params: Params): Unit = () + 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] +) - def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: Unit): (Params, Unit) = - val newParams = params -- gradients.value.scale(learningRate) - (newParams, ()) +class Lion(learningRate: Double | LearningRateSchedule, val weightDecay: Double = 0.0f, val beta1: Double = 0.9f, val beta2: Double = 0.99f) extends GradientOptimizer: -class Lion(val learningRate: Double, val weightDecay: Double = 0.0f, val beta1: Double = 0.9f, val beta2: Double = 0.99f) extends GradientOptimizer: + private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate) - type State[P] = P // momentum state has same structure as params + type State[P] = LionState[P] - def init[Params: IsFloatTree](params: Params): Params = - params.fillCopy(0f) + def init[Params: IsFloatTree](params: Params): LionState[Params] = + LionState(params.fillCopy(0f), step = 1) - def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, momentums: Params): (Params, Params) = + 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]( momentums: P, // momentums @@ -88,9 +102,7 @@ class Adam( epsilon: Double = 1e-8 // small constant to prevent division by zero ) extends GradientOptimizer: - private val learningRateF: Int => Double = learningRate match - case f: LearningRateSchedule => f - case d: Double => _ => d + private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate) private val β1 = b1 private val β2 = b2 private val ε = epsilon diff --git a/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala b/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala index bed3e03..982b5b1 100644 --- a/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala +++ b/core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala @@ -5,6 +5,11 @@ 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. diff --git a/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala b/core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala index e62411f..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) @@ -62,7 +84,8 @@ class GradientOptimizerSuite extends DimwitTest: val initMomentum = optimizer.init(initParams) val grad = Grad(Tensor0(6.0f)) - val (nextParams, nextMomentum) = optimizer.update(grad, initParams, initMomentum) + val (nextParams, nextState) = optimizer.update(grad, initParams, initMomentum) nextParams.item shouldBe 1.9f +- 1e-5f - nextMomentum.item shouldBe 0.06f +- 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