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

Merged
benikm91 merged 1 commit into from
Aug 10, 2026
Merged

Rework optimizers #139

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 7 additions & 3 deletions AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -870,6 +870,7 @@ println(s"Block Hessian shapes: ${h_x1x1.shape}, ${h_x1x2.shape}, ${h_x2x1.shape

```scala
import dimwit.*
import dimwit.Conversions.given
import dimwit.optimizer.{GradientDescent, GradientOptimizer}
import dimwit.random.Random

Expand Down Expand Up @@ -903,7 +904,7 @@ val lossFunc = mse(trainData, trainLabels)
val gradFunc = Autodiff.grad(lossFunc)

// Create optimizer
val optimizer = GradientDescent(learningRate = 0.01f)
val optimizer = GradientDescent.of(VType[Float32])(learningRate = 0.01f)

// Training loop with iterator
val trained = optimizer.iterate(initModelParams)(gradFunc)
Expand All @@ -918,9 +919,10 @@ val trained = optimizer.iterate(initModelParams)(gradFunc)

```scala
import dimwit.optimizer.Lion
import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V]

// Lion optimizer with momentum
val lionOptimizer = Lion(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)
val lionOptimizer = Lion.of(VType[Float32])(learningRate = 1e-3f, beta1 = 0.9f, beta2 = 0.99f, weightDecay = 0.0f)

// Training with Lion
val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc)
Expand All @@ -932,6 +934,8 @@ val trainedLion = lionOptimizer.iterate(initModelParams)(gradFunc)
### Complete Training Example: Linear Regression

```scala
import dimwit.Conversions.given // enables implicit conversion from Float to Tensor[V]

// Define problem dimensions
trait Sample derives Label
trait InputDim derives Label
Expand Down Expand Up @@ -963,7 +967,7 @@ val initRegressionParams = RegressionParams(initSlope, initIntercept)

// Train
val regressionGrad = Autodiff.grad(regressionLoss(xData, yData))
val gdOptimizer = GradientDescent(learningRate = 0.1f)
val gdOptimizer = GradientDescent.of(VType[Float32])(learningRate = 0.1f)

val finalParams = gdOptimizer.iterate(initRegressionParams)(regressionGrad)
.take(100)
Expand Down
156 changes: 96 additions & 60 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,11 @@
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.*

/** Gradient optimizer interface with functional state management.
*
Expand All @@ -27,85 +26,96 @@ import dimwit.autodiff.*
* optimizer.update(grads, params, state)
* }}}
*/
trait GradientOptimizer:
type State[_, V]
trait GradientOptimizer[V: IsFloating, State0[_]]:

type State[P] = State0[P]

// 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[P](params: P)(using TensorTree[P], FloatTree[P, V]): State[P]
def update[P](gradients: Grad[P], params: P, state: State[P])(using TensorTree[P], FloatTree[P, V]): (P, State[P])

// 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[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[(P, State[P])] =
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[P](init: P)(df: P => Grad[P])(using TensorTree[P], FloatTree[P, V]): Iterator[P] =
iterateWithState(init)(df).map(_._1)

case class GradientDescent(learningRate: Double) extends GradientOptimizer:
object GradientDescent:

def of[V: IsFloating](vtype: VType[V])(learningRate: Tensor0[V]): GradientDescent[V] = new GradientDescent(learningRate)

type GradientDescentState[P] = Unit // empty state

type State[P, V] = Unit // Stateless optimizer
class GradientDescent[V: IsFloating](val learningRate: Tensor0[V]) extends GradientOptimizer[V, GradientDescentState]:

def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): Unit = ()
def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): Unit = ()

def update[V, Params: TensorTree: FloatTreeFor[V]](gradients: Grad[Params], params: Params, state: Unit)(using IsFloating[V]): (Params, Unit) =
def update[P](gradients: Grad[P], params: P, state: Unit)(using TensorTree[P], FloatTree[P, V]): (P, 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:
case class LionState[P](
momentums: P,
step: Tensor0[Int32]
)

object Lion:

def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], weightDecay: Tensor0[V] = Tensor0(vtype)(0.0), beta1: Tensor0[V] = Tensor0(vtype)(0.9), beta2: Tensor0[V] = Tensor0(vtype)(0.99)): Lion[V] = new Lion(learningRate, weightDecay, beta1, beta2)

type State[P, V] = P // momentum state has same structure as params
class Lion[V: IsFloating](val learningRate: Tensor0[V], val weightDecay: Tensor0[V] = Tensor0(0.0), val beta1: Tensor0[V] = Tensor0(0.9), val beta2: Tensor0[V] = Tensor0(0.99)) extends GradientOptimizer[V, LionState]:

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[P](params: P)(using TensorTree[P], FloatTree[P, V]): LionState[P] =
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[P](gradients: Grad[P], params: P, state: LionState[P])(using TensorTree[P], FloatTree[P, V]): (P, LionState[P]) =
// 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 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, V](
momentums: P, // momentums
velocities: P, // velocities
b1: Tensor0[V], // decay rate for momentums mᵗ
b2: Tensor0[V] // decay rate for velocities vᵗ
)

object Adam:

def of[V](vtype: VType[V])(using IsFloating[V])(learningRate: Tensor0[V], b1: Tensor0[V] = Tensor0(vtype)(0.9), b2: Tensor0[V] = Tensor0(vtype)(0.999), epsilon: Tensor0[V] = Tensor0(vtype)(1e-8)): Adam[V] = new Adam(learningRate, b1, b2, epsilon)

/** 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)
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:
class Adam[V: IsFloating](
val learningRate: Tensor0[V],
b1: Tensor0[V] = Tensor0(0.9), // decay rate for momentums mᵗ
b2: Tensor0[V] = Tensor0(0.999), // decay rate for velocities vᵗ
epsilon: Tensor0[V] = Tensor0(1e-8) // small constant to prevent division by zero
) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]:

private val β1 = b1
private val β2 = b2
private val ε = epsilon

type State[P, V] = AdamState[P, V]

def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] =
def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] =
def zeros = params.fillCopy(0f)
AdamState[Params, V](zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f))
AdamState(zeros, zeros, b1 = Tensor0(VType[V])(1f), b2 = Tensor0(VType[V])(1f))

def update[V, Params: TensorTree: FloatTreeFor[V]](
gradients: Grad[Params],
params: Params,
state: State[Params, V]
)(using IsFloating[V]): (Params, State[Params, V]) =
def update[P](
gradients: Grad[P],
params: P,
state: AdamState[P, V]
)(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) =
// rename state variables to last time step for clarity
val `mₜ₋₁` = state.momentums
val `vₜ₋₁` = state.velocities
Expand All @@ -114,17 +124,15 @@ case class Adam(

// rename parameters for internal clarity
val α = learningRate
val ε = epsilon
val `θₜ₋₁` = params

// 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ₜ = `β1ₜ₋₁` * β1
val β2ₜ = `β2ₜ₋₁` * β2
val m̂ = mᵗ `//!` (1f - `β1ₜ`)
val v̂ = vᵗ `//!` (1f - `β2ₜ`)
val θₜ = `θₜ₋₁` -- (α **! m̂) `//` (v̂.sqrt ++! ε)
Expand All @@ -141,24 +149,52 @@ case class Adam(
* @param learningRate The step size.
* @param weightDecayFactor The coefficient for weight decay (lambda).
*/
case class AdamW(
val adam: Adam,
val weightDecayFactor: Double
) extends GradientOptimizer:

type State[P, V] = adam.State[P, V]

def init[V, Params: TensorTree: FloatTreeFor[V]](params: Params)(using IsFloating[V]): State[Params, V] = adam.init(params)

def update[V, Params: TensorTree: FloatTreeFor[V]](
gradients: Grad[Params],
params: Params,
state: State[Params, V]
)(using IsFloating[V]): (Params, State[Params, V]) =
class AdamW[V: IsFloating](
val adam: Adam[V],
val weightDecayFactor: Tensor0[V]
) extends GradientOptimizer[V, [P] =>> AdamState[P, V]]:

def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): AdamState[P, V] = adam.init(params)

def update[P](
gradients: Grad[P],
params: P,
state: AdamState[P, V]
)(using TensorTree[P], FloatTree[P, V]): (P, AdamState[P, V]) =
val α = adam.learningRate
val `θₜ₋₁` = params
val `λ'` = weightDecayFactor
val λ = `λ'` * α // Tie weight decay to learning rate
val decayedParams = `θₜ₋₁` -- λ **! `θₜ₋₁`
val (θₜ, adamState) = adam.update(gradients, decayedParams, state)
(θₜ, adamState)

case class LearningRateScheduleState[P, State[_]](
step: Tensor0[Int32],
optState: State[P]
)
type LearningRateScheduleStateFor[State[_]] = [P] =>> LearningRateScheduleState[P, State]

object LearningRateSchedule:

def of[V: IsFloating, State[_]](
vtype: VType[V]
)(
optF: Tensor0[V] => GradientOptimizer[V, State],
schedule: Tensor0[Int32] => Tensor0[V]
): LearningRateSchedule[V, State] =
new LearningRateSchedule(optF, schedule)

class LearningRateSchedule[V: IsFloating, State[_]](val optF: Tensor0[V] => GradientOptimizer[V, State], schedule: Tensor0[Int32] => Tensor0[V]) extends GradientOptimizer[V, LearningRateScheduleStateFor[State]]:

def init[P](params: P)(using TensorTree[P], FloatTree[P, V]): LearningRateScheduleState[P, State] =
val step = Tensor0(1)
val opt = optF(schedule(step))
LearningRateScheduleState(step, opt.init(params))

def update[P](gradients: Grad[P], params: P, state: LearningRateScheduleState[P, State])(using TensorTree[P], FloatTree[P, V]): (P, LearningRateScheduleState[P, State]) =
val step = state.step
val optState = state.optState
val opt = optF(schedule(step))
val (newParams, newOptState) = opt.update(gradients, params, optState)
(newParams, LearningRateScheduleState(step + 1, newOptState))
3 changes: 1 addition & 2 deletions core/src/main/scala/dimwit/tensor/TensorOps.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,8 @@ object TensorOps:
given [V](using ev1: IsFloating[V]): IsNumber[V] = ev1
given [V](using ev2: IsInteger[V]): IsNumber[V] = ev2

@implicitNotFound("Operation only valid for Floating tensors.")

/** Type class marker for floating point types (Float32, Float64, etc.). */
@implicitNotFound("Operation only valid for Floating tensors.")
trait IsFloating[V] extends IsNumber[V], HasDType[V]:
def dtype: DType

Expand Down
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
Loading
Loading