Rework optimizers by benikm91 · Pull Request #139 · dimwit-dev/dimwit

benikm91 · 2026-08-03T17:22:38Z

This PR builds on #137 and should be merged afterward.

I am updating the optimizers to run with deepwit GPT branch. By doing this I found an issue with the optimizers:

The precision I added recently was suboptimal, as the optimizers took two generics [Params, V]; however, Params could be of a different precision than V, resulting in potential errors. I fixed this with the abstraction IsFloatTree that allows removing V from optimizers; additionally,
I removed beta1 and beta2 from AdamState, as these are implementation details and not the actual state of Adam. Storing the iteration is mathematically enough, as beta1^n allows for bias correction in exponential decay. However, beta1^n is an inefficient implementation. I implemented a stateful abstraction, SequenceFunction, to handle this. Please check if this is okay.
I added a primitive learning rate schedule from the lr-schedule branch of my fork, used in GPT-2 training. Check the design. Extension is possible later. Currently implemented only in Adam.

marcelluethi

Thanks for starting the implementation of these ideal. I am not an expert in optimizer and I am not aware of all the practical issues. Therefore my review is mostly a set of questions I was wondering about when reading the code.

marcelluethi · 2026-08-03T18:02:33Z

+/** 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]:


I don't quite understand how IsFloatTree is different from FloatTree. Wouldn't it be possible to enforce the IsFloating[V] constraint there already?

FloatTree[P, V] is for a specific V.
IsFloatTree[P] marks any possible FloatTree

I can write

def x[P[_], V](p: P[V])(using FloatTree[P, V])

But then we can't use for hard-coded precision in e.g. params:

// x not applicable to as Params1 is not a P[_] type (takes no generic) case class Params1(weights: Tensor2[Row, Col, Float32] // x only applicable to as Params2 is a P[_] type case class Params2[V: IsFloating](weights: Tensor2[Row, Col, V]

If I do

def x2[P](p: P)(using IsFloatTree) // x2 applicable to Params1[Float32] and applicable to Params2.

So far to the motivation. I don't know if there is a better solution :) Best I came up with.

To make it concrete for the optimizers. IsFloatTree was here necessary to make the VAE example run that has hard coded Params precision. I think we should support hard coding precision.

Do we need the higher kinded type here? Maybe something like would be easier to work with?

def x[P, V](p: P)(using FloatTree[P, V])

marcelluethi · 2026-08-03T18:06:40Z

+  *
+  * @tparam R The type of the state.
+  */
+trait SequenceFunction[R] extends (Int => R):


I am not an expert on optimizers and learning rate adaption in deep learning. But my naive guess would be that way more work is done within a computation step than for the learning rate adaption. If this is true, why would we need a "high performance" abstraction, rather than having a simple function?

Yes, it will not matter computation-wise. It felt more like bad code if exponential decay re-evaluates each iteration with .pow; however, mathematically, this is what happens. I changed to .pow in the last commit, until SequenceFunction will show to be necessary.

marcelluethi · 2026-08-03T18:14:13Z

-case class Adam(
-    learningRate: Double, // step size (learning rate)
+class Adam(
+    learningRate: Double | (Int => Double),


I don't understand this interface. I assume (Int => Double) allows for adapting the learning rate. But isn't this the same as the learning rate scheduler? How does serialization (for a checkpoint) work? I guess there should be more than a function for that. Should the learning rate somehow be stored with the state?

I changed (Int => Double) to LearningRateSchedule type alias to make it more explicit.

marcelluethi · 2026-08-03T18:16:22Z

+    def doStep(): Unit = β1ₜ *= β1; β2ₜ *= β2
+    def state(): (Double, Double) = (β1ₜ, β2ₜ)
+
+    override protected def outOfOrderErrorMessage(step: Long, current: Long): String =


Is this a consequence of the stateful abstraction? If yes, is the price worth paying?

Yes, it is a consequence. No, probably not worth paying until shown to be necessary due to slowdown.

benikm91 · 2026-08-04T06:27:20Z

Something I am unsure about is the LearningRateSchedule. My abstraction is a function from Int => Double. When composing two schedules, we do it with pointwiseMin, taking the lower value of all schedules.

// in deepwit
val schedule = pointwiseMin(linearWarmup(lr, 1000), cosineDecay(lr, minLr, 20_000).delay(1000))

linearWarmup ramps up from 0 to lr in 1000 steps.
cosineDecay is lr the first 1000 steps (delay(1000)) and then decays to minLr in the next 20k steps and stays at minLr.

If we take the pointwiseMin, we take warmup for the first 1000 steps and cosineDecay for the rest of training.

+) All schedules are valid schedules from [1, inf]
-) Confusing?

An alternative design would be to make a more bounded function, e.g.,

linearWarmup(lr, 1000) would ramp up to lr in 1000 steps and throw an out-of-range exception afterward.

and we would compose these functions:

val schedule = linearWarmup(lr, 1000).compose(cosineDecay(lr, minLr, 20_000))

+) Schedules are function compositions
-) Error-prone, e.g. "linearWarmup(lr, 1000)" results in error after 1000 steps if not composed with another schedule.

* Remove beta1 and beta2 from AdamState * Remove generic V from all optimizers, replace with IsFloatingTree abstraction. * Add learning rate schedule

… move LRSchedule in custom file

benikm91 · 2026-08-04T07:13:04Z

Moving params inside optimizer states? Something worth considering for this PR?

trait GradientOptimizer:
  type State[_]

  // Core API
  def init[Params: IsFloatTree](params: Params): State[Params]
  def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: State[Params]): (Params, State[Params])

We could change this to

case class OptimizerState[P](val params: Params)

  */
trait GradientOptimizer:
  type State[P] <: OptimizerState[P]

  // Core API
  def init[Params: IsFloatTree](params: Params): State[Params]
  def update[Params: IsFloatTree](gradients: Grad[Params], state: State[Params]): State[Params]

Move the params into the optimizer state:

+) Reduces two parameters to one in user side.

    def trainBatch(..., params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) =
      val grads = grad(...)
      val (newParams, newState) = optimizer.update(grads, params, state)
      (newParams, newState)

would become:

    def trainBatch(..., state: optimizer.State[Params]): optimizer.State[Params] =
      val grads = grad(...)
      val newState = optimizer.update(grads, state)
      newState

-) Conceptually misaligned? Parameters are "hidden" within the state... Or does this make sense? Optimizer state has the current position/params as state.

marcelluethi · 2026-08-04T18:26:30Z

Something I am unsure about is the LearningRateSchedule. My abstraction is a function from Int => Double. When composing two schedules, we do it with pointwiseMin, taking the lower value of all schedules.
// in deepwit
val schedule = pointwiseMin(linearWarmup(lr, 1000), cosineDecay(lr, minLr, 20_000).delay(1000))
linearWarmup ramps up from 0 to lr in 1000 steps. cosineDecay is lr the first 1000 steps (delay(1000)) and then decays to minLr in the next 20k steps and stays at minLr.

If we take the pointwiseMin, we take warmup for the first 1000 steps and cosineDecay for the rest of training.

+) All schedules are valid schedules from [1, inf] -) Confusing?

An alternative design would be to make a more bounded function, e.g.,

linearWarmup(lr, 1000) would ramp up to lr in 1000 steps and throw an out-of-range exception afterward.

and we would compose these functions:
val schedule = linearWarmup(lr, 1000).compose(cosineDecay(lr, minLr, 20_000))
+) Schedules are function compositions -) Error-prone, e.g. "linearWarmup(lr, 1000)" results in error after 1000 steps if not composed with another schedule.

Maybe this should not be a component in dimwit, but rather in deepwit. If we move it there, we can experiment. In dimwit we would have to hand roll our own schedule (without compositional api)

marcelluethi · 2026-08-04T18:34:27Z

Moving params inside optimizer states? Something worth considering for this PR?

trait GradientOptimizer:
  type State[_]

  // Core API
  def init[Params: IsFloatTree](params: Params): State[Params]
  def update[Params: IsFloatTree](gradients: Grad[Params], params: Params, state: State[Params]): (Params, State[Params])

We could change this to

case class OptimizerState[P](val params: Params)

  */
trait GradientOptimizer:
  type State[P] <: OptimizerState[P]

  // Core API
  def init[Params: IsFloatTree](params: Params): State[Params]
  def update[Params: IsFloatTree](gradients: Grad[Params], state: State[Params]): State[Params]

Move the params into the optimizer state:

+) Reduces two parameters to one in user side.

    def trainBatch(..., params: Params, state: optimizer.State[Params]): (Params, optimizer.State[Params]) =
      val grads = grad(...)
      val (newParams, newState) = optimizer.update(grads, params, state)
      (newParams, newState)

would become:

    def trainBatch(..., state: optimizer.State[Params]): optimizer.State[Params] =
      val grads = grad(...)
      val newState = optimizer.update(grads, state)
      newState

-) Conceptually misaligned? Parameters are "hidden" within the state... Or does this make sense? Optimizer state has the current position/params as state.

To me it feels wrong to have the Parameters be part of the Optimization state, even if it would make it a bit more user friendly. I haven't thought through the consequences. But it just doesn't seem to belong there conceptually.

marcelluethi reviewed Aug 3, 2026

View reviewed changes

marcelluethi and others added 2 commits August 4, 2026 08:40

move tensor tree related methods to separate package

7e1c382

Rework optimizers

d35ad7f

* Remove beta1 and beta2 from AdamState * Remove generic V from all optimizers, replace with IsFloatingTree abstraction. * Add learning rate schedule

benikm91 force-pushed the rework-optimizers branch from 8e10342 to 299781b Compare August 4, 2026 06:42

Cleanup Optimizers: Remove SequenceFunction (until proven necessary),…

aaf493f

… move LRSchedule in custom file

benikm91 force-pushed the rework-optimizers branch from 299781b to aaf493f Compare August 4, 2026 06:43

Add LearningRateSchedule to other optimizers.

4db980d

Uh oh!

Conversation

benikm91 commented Aug 3, 2026

Uh oh!

marcelluethi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

benikm91 commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

marcelluethi commented Aug 4, 2026

Uh oh!

marcelluethi commented Aug 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

benikm91 Aug 4, 2026 •

edited

Loading

benikm91 commented Aug 4, 2026 •

edited

Loading

benikm91 commented Aug 4, 2026 •

edited

Loading