From ade62e0390519f4d706d4c39eedcddece2bae8e1 Mon Sep 17 00:00:00 2001 From: Marcel Luethi Date: Fri, 17 Jul 2026 17:05:22 +0200 Subject: [PATCH] change optimizer hyperparameters to double As hyperparameters are not optimized, they don't need to be differentiable. --- AGENTS.md | 6 +-- .../dimwit/optimizer/GradientOptimizer.scala | 44 +++++++++++-------- docs/quickstart.md | 2 +- .../main/scala/basic/LogisticRegression.scala | 2 +- .../complex/VariationalAutoencoder.scala | 2 +- mdocs/AGENTS.md | 6 +-- mdocs/docs/quickstart.md | 2 +- 7 files changed, 36 insertions(+), 28 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index dd3dcd45..7721cc5b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -902,7 +902,7 @@ val lossFunc = mse(trainData, trainLabels) val gradFunc = Autodiff.grad(lossFunc) // Create optimizer -val optimizer = GradientDescent(learningRate = Tensor0(0.01f)) +val optimizer = GradientDescent(learningRate = 0.01) // Training loop with iterator val trained = optimizer.iterate(initModelParams)(gradFunc) @@ -919,7 +919,7 @@ val trained = optimizer.iterate(initModelParams)(gradFunc) import dimwit.optimizer.Lion // Lion optimizer with momentum -val lionOptimizer = Lion(learningRate = Tensor0(1e-3f), beta1 = Tensor0(0.9f), beta2 = Tensor0(0.99f), weightDecay = Tensor0(0.0f)) +val lionOptimizer = Lion(learningRate = 1e-3, beta1 = 0.9, beta2 = 0.99, weightDecay = 0.0) // Training with Lion val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) @@ -962,7 +962,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept) // Train val regressionGrad = Autodiff.grad(regressionLoss(xData, yData)) -val gdOptimizer = GradientDescent(learningRate = Tensor0(0.1f)) +val gdOptimizer = GradientDescent(learningRate = 0.1) val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad) .take(100) diff --git a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala index 17a2d535..2f6e6983 100644 --- a/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala +++ b/core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala @@ -40,17 +40,23 @@ trait GradientOptimizer: def iterate[Params: TensorTree: FloatTreeFor[Float32]](init: Params)(df: Params => Grad[Params]): Iterator[Params] = iterateWithState(init)(df).map(_._1) -case class GradientDescent(learningRate: Tensor0[Float32]) extends GradientOptimizer: +case class GradientDescent(learningRate: Double) extends GradientOptimizer: + + private val lr = Tensor0(learningRate.toFloat) type State[P] = Unit // Stateless optimizer def init[Params: TensorTree: FloatTreeFor[Float32]](params: Params): Unit = () def update[Params: TensorTree: FloatTreeFor[Float32]](gradients: Grad[Params], params: Params, state: Unit): (Params, Unit) = - val newParams = params -- gradients.value.scale(learningRate) + val newParams = params -- gradients.value.scale(lr) (newParams, ()) -case class Lion(learningRate: Tensor0[Float32], weightDecay: Tensor0[Float32] = Tensor0(0.0f), beta1: Tensor0[Float32] = Tensor0(0.9f), beta2: Tensor0[Float32] = Tensor0(0.99f)) extends GradientOptimizer: +case class Lion(learningRate: Double, weightDecay: Double = 0.0, beta1: Double = 0.9, beta2: Double = 0.99) extends GradientOptimizer: + + val beta1f = Tensor0(beta1.toFloat) + val beta2f = Tensor0(beta2.toFloat) + val lr = Tensor0(learningRate.toFloat) type State[P] = P // momentum state has same structure as params @@ -62,12 +68,13 @@ case class Lion(learningRate: Tensor0[Float32], weightDecay: Tensor0[Float32] = ) def update[Params: TensorTree: FloatTreeFor[Float32]](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 + val updateDirection = (momentums **! beta1f ++ gradients.value **! (1f - beta1f)).sign - val updatedParams = params -- updateDirection.scale(learningRate) -- params.scale(weightDecay) - val newMomentums = momentums **! beta2 ++ gradients.value **! (1f - beta2) + val updatedParams = params -- updateDirection.scale(lr) -- params.scale(Tensor0(weightDecay.toFloat)) + val newMomentums = momentums **! beta2f ++ gradients.value **! (1f - beta2f) (updatedParams, newMomentums) @@ -83,20 +90,21 @@ case class AdamState[P]( * @see [[https://arxiv.org/abs/1412.6980 Adam: A Method for Stochastic Optimization]] */ case class Adam( - learningRate: Tensor0[Float32], // step size (learning rate) - b1: Tensor0[Float32] = Tensor0(0.9f), // decay rate for momentums mᵗ - b2: Tensor0[Float32] = Tensor0(0.999f), // decay rate for velocities vᵗ - epsilon: Tensor0[Float32] = Tensor0(1e-8f) // small constant to prevent division by zero + learningRate: Double, // step size (learning rate) + 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 β1 = b1 - private val β2 = b2 + private val β1 = Tensor0(b1.toFloat) + private val β2 = Tensor0(b2.toFloat) + private val ε = Tensor0(epsilon.toFloat) type State[P] = AdamState[P] def init[Params: TensorTree: FloatTreeFor[Float32]](params: Params): State[Params] = def zeros = params.fillCopy(0f) - AdamState(zeros, zeros, b1 = Tensor0(1f), b2 = Tensor0(1f)) + AdamState(zeros, zeros, b1 = Tensor0(1.0f), b2 = Tensor0(1.0f)) def update[Params: TensorTree: FloatTreeFor[Float32]]( gradients: Grad[Params], @@ -110,8 +118,8 @@ case class Adam( val `β2ₜ₋₁` = state.b2 // rename parameters for internal clarity - val α = learningRate - val ε = epsilon + val α = Tensor0(learningRate.toFloat) + val ε = Tensor0(epsilon.toFloat) val `θₜ₋₁` = params // update moments for bias correction @@ -140,7 +148,7 @@ case class Adam( */ case class AdamW( val adam: Adam, - val weightDecayFactor: Tensor0[Float32] + val weightDecayFactor: Double ) extends GradientOptimizer: type State[P] = adam.State[P] @@ -152,9 +160,9 @@ case class AdamW( params: Params, state: State[Params] ): (Params, State[Params]) = - val α = adam.learningRate + val α = Tensor0(adam.learningRate.toFloat) val `θₜ₋₁` = params - val `λ'` = weightDecayFactor + val `λ'` = Tensor0(weightDecayFactor.toFloat) val λ = `λ'` * α // Tie weight decay to learning rate val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁` val (θₜ, adamState) = adam.update(gradients, decayedParams, state) diff --git a/docs/quickstart.md b/docs/quickstart.md index c75886c4..6988fa70 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -43,7 +43,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat val gradFn = grad(loss(x, y)) // gradient based optimization - val gd = GradientDescent(learningRate = Tensor0(0.1f)) // this is wrong, should be 0.1f not Tensor0 + val gd = GradientDescent(learningRate = 0.1) gd.iterate(p0)(gradFn) ``` diff --git a/examples/src/main/scala/basic/LogisticRegression.scala b/examples/src/main/scala/basic/LogisticRegression.scala index 1b35cf1c..aabda6d6 100644 --- a/examples/src/main/scala/basic/LogisticRegression.scala +++ b/examples/src/main/scala/basic/LogisticRegression.scala @@ -125,7 +125,7 @@ object LogisticRegression: val trainLoss = jit(BinaryLogisticRegression.loss(trainingData, trainLabels)) val valLoss = jit(BinaryLogisticRegression.loss(valData, valLabels)) val learningRate = 5e-1f - val gd = GradientDescent(Tensor0(learningRate)) + val gd = GradientDescent(learningRate) // Training loop val numiterations = 1000 diff --git a/examples/src/main/scala/complex/VariationalAutoencoder.scala b/examples/src/main/scala/complex/VariationalAutoencoder.scala index a075f996..8fba39ef 100644 --- a/examples/src/main/scala/complex/VariationalAutoencoder.scala +++ b/examples/src/main/scala/complex/VariationalAutoencoder.scala @@ -209,7 +209,7 @@ object VariationalAutoencoderExample: losses.sum / batchSize.toFloat val batches = trainImages.chunk(Axis[TrainSample], numSamples / batchSize) - val optimizer = GradientDescent(learningRate = Tensor0(learningRate)) + val optimizer = GradientDescent(learningRate = learningRate) def trainBatch(trainKey: Random.Key, batch: Tensor3[TrainSample, Height, Width, Float32], params: Params): Params = val grads = grad(batchLoss(trainKey, batch))(params) val (newParams, _) = optimizer.update(grads, params, ()) diff --git a/mdocs/AGENTS.md b/mdocs/AGENTS.md index 0cff8cf7..d74afb7f 100644 --- a/mdocs/AGENTS.md +++ b/mdocs/AGENTS.md @@ -709,7 +709,7 @@ val lossFunc = mse(trainData, trainLabels) val gradFunc = Autodiff.grad(lossFunc) // Create optimizer -val optimizer = GradientDescent(learningRate = Tensor0(0.01f)) +val optimizer = GradientDescent(learningRate = 0.01) // Training loop with iterator val trained = optimizer.iterate(initModelParams)(gradFunc) @@ -726,7 +726,7 @@ val trained = optimizer.iterate(initModelParams)(gradFunc) import dimwit.optimizer.Lion // Lion optimizer with momentum -val lionOptimizer = Lion(learningRate = Tensor0(1e-3f), beta1 = Tensor0(0.9f), beta2 = Tensor0(0.99f), weightDecay = Tensor0(0.0f)) +val lionOptimizer = Lion(learningRate = 1e-3, beta1 = 0.9, beta2 = 0.99, weightDecay = 0.0) // Training with Lion val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc) @@ -769,7 +769,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept) // Train val regressionGrad = Autodiff.grad(regressionLoss(xData, yData)) -val gdOptimizer = GradientDescent(learningRate = Tensor0(0.1f)) +val gdOptimizer = GradientDescent(learningRate = 0.1) val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad) .take(100) diff --git a/mdocs/docs/quickstart.md b/mdocs/docs/quickstart.md index 231768e8..7cf1dc6a 100644 --- a/mdocs/docs/quickstart.md +++ b/mdocs/docs/quickstart.md @@ -43,7 +43,7 @@ def fit(x: Tensor2[Batch, Feature, Float32], y: Tensor1[Batch, Float32]): Iterat val gradFn = grad(loss(x, y)) // gradient based optimization - val gd = GradientDescent(learningRate = Tensor0(0.1f)) // this is wrong, should be 0.1f not Tensor0 + val gd = GradientDescent(learningRate = 0.1) gd.iterate(p0)(gradFn) ```