Skip to content

Rework optimizers #139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
benikm91 wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

Rework optimizers #139

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7e1c382
move tensor tree related methods to separate package
marcelluethi Aug 1, 2026
d35ad7f
Rework optimizers
benikm91 Aug 4, 2026
aaf493f
Cleanup Optimizers: Remove SequenceFunction (until proven necessary),…
benikm91 Aug 4, 2026
4db980d
Add LearningRateSchedule to other optimizers.
benikm91 Aug 4, 2026
88e36b4
Remove LearningRateSchedule from dimwit (moved to deepwit)
benikm91 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 62 additions & 54 deletions core/src/main/scala/dimwit/optimizer/GradientOptimizer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
package dimwit.optimizer

import dimwit.*

import dimwit.Conversions.given
import dimwit.autodiff.*
import dimwit.autodiff.Grad
import dimwit.tensortree.*
import dimwit.tensortree.FloatTree.*
import dimwit.tensortree.FloatTree.ops.*
import dimwit.autodiff.*
import dimwit.tensortree.IsFloatTree.given

/** Gradient optimizer interface with functional state management.
*
Expand All @@ -28,108 +28,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:

private val learningRateF: Int => Double = LearningRateSchedule.from(learningRate)

type State[P, V] = Unit // Stateless optimizer
type State[P] = GradientDescentState[P]

def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Unit = ()
def init[Params: IsFloatTree](params: Params): GradientDescentState[Params] =
GradientDescentState(step = 1)

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 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 Lion(learningRate: Double, weightDecay: Double = 0.0f, beta1: Double = 0.9f, beta2: Double = 0.99f) extends GradientOptimizer:
case class LionState[P](
momentums: P,
step: Tensor0[Int32]
)

class Lion(learningRate: Double | LearningRateSchedule, 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, V] = P // momentum state has same structure as params
type State[P] = LionState[P]

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): LionState[Params] =
LionState(params.fillCopy(0f), step = 1)

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, 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 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ₜ) = (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).
*
Expand All @@ -141,21 +149,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
Expand Down
11 changes: 11 additions & 0 deletions core/src/main/scala/dimwit/optimizer/LearningRateSchedule.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
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
21 changes: 21 additions & 0 deletions core/src/main/scala/dimwit/tensortree/FloatTree.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:

@marcelluethi marcelluethi Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@benikm91 benikm91 Aug 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@benikm91 benikm91 Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@marcelluethi marcelluethi Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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
44 changes: 33 additions & 11 deletions core/src/test/scala/dimwit/optimizer/GradientOptimizerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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"):
Expand All @@ -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"):
Expand All @@ -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
21 changes: 11 additions & 10 deletions examples/src/main/scala/complex/VariationalAutoencoder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading