change optimizer hyperparameters to double

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)

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)

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)
 ```
 

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

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, ())

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)

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)
 ```