WIP: New Optimizer interface and Gradient Type

[marcelluethi]

This PR serves as a basis for discussion regarding a new interface to optimisers.
The new interface was made necessary, such that we can accommodate optimizers that have internal state (e.g. momentum). The usage pattern is the following, if we want to work with minibatches:

val optimizer = GradientDescent(lr = 0.1)
val (finalState, finalParams) = batches.foldLeft((optimizer.init(initParams), initParams)):
   case ((state, params), batch) =>
       val grads = Autodiff.grad(loss(batch))(params)
       optimizer.update(grads, state, params) 

For a simple optimization, we can also hide the state completely:

val optimizer = GradientDescent(lr = 0.1)
optimizer.iterate(initParams)(gradientFunction).take(1000).foreach(...)

There is also a second idea in this PR that I am not sure if it is a good idea: the use of the label Grad[Param] to distinguish a gradient from the parameter. It gives some type safety, but makes higher-order differentiation harder. If we decide to go that way, we might also consider labeling Jacobians the same way.

[benikm91]

We can solve this using a StateMonad (holding the train specific state, e.g. momentum).
Advantage: GradientDescent (no state) would be identical code on the user side to Adam (state), both StateMonads (one empty, one with state); this allows for switching between optimizers without having to update the code.
Disadvantage: User may be unfamiliar with Monads or StateMonads, making code harder to understand, steepening the learning curve.

---

An inbetween solution would be that GradientDescent returns a Unit () for state; and update being defined in a common optimizer trait.

[marcelluethi]

Gradient Descent is already returning Unit and the interface is unified.

I am strongly against the use of monads or other advanced functional patterns. Somebody who wants to benefit from more type safety in their tensor computation should not need to read up on advanced functional programming concepts. Especially also, since there is a push towards direct style in Scala.

Do you have any opinion about the use of Grad as a type? I am unsure if it makes it more complicated or more clear.

[benikm91]

I agree with StateMonads. Sorry I didn't realize your implementation already does the Unit. This is perfect.

Grad:

What I like about it:

1. Making Gradients and Values different types makes a lot of sense to me. As these are fundamentally different concepts.

What I don't like about it:

1. The current implementation makes a Grad[...] type, e.g.,
   val delta: Grad[(Tensor[(A, B), Float], Tensor[Tuple1[C], Float])]
   In the user code it is unknown that Grad is itself a tensor, so all TensorOps are not available. So we could not add gradients (e.g., gradient accumulation), or scale parameters differently (fine-tuning), ...
   => I would reject the current implementation

Alternative Suggestion:
what do you think of (Ignoring the complexity of the Implementation for now):
val delta: (Tensor[(A, B), Grad[Float]], Tensor[Tuple1[C], Grad[Float]])]
with Grad[Float] knowing it is a IsFloat (so all TensorOperations for Tensor[_, Float] are available).
And one could not accidentally add a Grad[Float] to a Float Tensor.
Higher derivatives could be Grad[Grad[Float]].

[marcelluethi]

Grad is not really a tensor, but a PyTree. To work with it we would anyway need to use a TreeMap. It is only in the simplest cases that Grad corresponds to a tensor.

  def grad[Input, V](f: Input => Tensor0[V])(using
      inTree: ToPyTree[Input],
      outTree: ToPyTree[Tensor0[V]]
  ): Input => Grad[Input] =

[benikm91]

You are correct. For my suggestion to work, the Params class would have to be Params[V], that is, Input[V] in grad... Feels like the wrong path to take.

We can do the following to "test" gradient accumulation:

  object MLP:
    case class Params(
        layer1: LinearLayer.Params[Height |*| Width, Hidden],
        layer2: LinearLayer.Params[Hidden, Output]
    ):
      def +(other: Params): Params =
        val tree = summon[FloatTensorTree[Params]]
        tree.zipMap(this, other, [T <: Tuple] => (n: Labels[T]) ?=> (a: Tensor[T, Float], b: Tensor[T, Float]) => a + b)

Then we can test gradient accumulation

val grads2 = grads + grads // does not compile

// but with <: T
opaque type Grad[T] <: T = T
val grads2 = grads + grads // compiles
val grads3 = grads + trainState.params // does also compile (should P + G[P] work?)

[benikm91]

Make sure that code works with both:

// val optimizer = Lion(learningRate = Tensor0(learningRate), weightDecay = Tensor0(0f))
val optimizer = GradientDescent(learningRate = Tensor0(learningRate))

Currently parts of the MLPClassifierMNist code have hardcoded state (for Lion). Optimizers should be plug and play.

Rewriting the code to work with both led me to the following implementation.
The main change is a single TrainState object instead of a tuple (state, params), which makes passing along the two parameters easier. What do you think?

trait GradientOptimizer:
  type State[P]
  // Train State rather than tuple for the user side to pass around.
  case class TrainState[Params](val state: State[Params], val params: Params)

  // Core JAX-style API
  def init[Params: ToPyTree: FloatTensorTree](params: Params): TrainState[Params] = TrainState(initState(params), params)
  def initState[Params: ToPyTree: FloatTensorTree](params: Params): State[Params]
  def update[Params: ToPyTree: FloatTensorTree](gradients: Grad[Params], last: TrainState[Params]): TrainState[Params]

  // Convenience: iterator with fixed gradient function
  def iterate[Params: ToPyTree: FloatTensorTree](initParams: Params)(df: Params => Grad[Params]): Iterator[TrainState[Params]] =
    Iterator.iterate(this.init(initParams)): r =>
      val grads = df(r.params)
      val newResult = this.update(grads, r)
      newResult


case class GradientDescent(learningRate: Tensor0[Float]) extends GradientOptimizer:
  import dimwit.Conversions.given

  type State[P] = Unit // Stateless optimizer

  def initState[Params: ToPyTree: FloatTensorTree](params: Params): State[Params] = ()

  def update[Params: ToPyTree: FloatTensorTree](gradients: Grad[Params], last: TrainState[Params]): TrainState[Params] =
    val paramTree = summon[FloatTensorTree[Params]]
    val newParams = paramTree.zipMap(
      gradients.value,
      last.params,
      [T <: Tuple] => (n: Labels[T]) ?=> (g: Tensor[T, Float], p: Tensor[T, Float]) => p - g.scale(learningRate)
    )
    last.copy(params = newParams)

case class Lion(learningRate: Tensor0[Float], weightDecay: Tensor0[Float] = Tensor0(0.0f), beta1: Tensor0[Float] = Tensor0(0.9f), beta2: Tensor0[Float] = Tensor0(0.99f)) extends GradientOptimizer:
  import dimwit.Conversions.given

  type State[P] = P // momentum state has same structure as params

  def initState[Params: ToPyTree: FloatTensorTree](params: Params): State[Params] =
    val paramTree = summon[FloatTensorTree[Params]]
    paramTree.map(
      params,
      [T <: Tuple] =>
        (n: Labels[T]) ?=>
          (t: Tensor[T, Float]) =>
            Tensor.zeros(t.shape, VType[Float])
    )

  def update[Params: ToPyTree: FloatTensorTree](gradients: Grad[Params], last: TrainState[Params]): TrainState[Params] =
    val (momentums, params) = (last.state, last.params)
    val paramTree = summon[FloatTensorTree[Params]]
    // the direction (1 or -1)
    // is determined by the sign of the momentum + gradient
    val updateDirection = paramTree.zipMap(
      gradients.value,
      momentums,
      [T <: Tuple] =>
        (n: Labels[T]) ?=>
          (grad: Tensor[T, Float], momentum: Tensor[T, Float]) =>
            (momentum *! beta1 + grad *! (1f - beta1)).sign
    )

    val updatedParams = paramTree.zipMap(
      updateDirection,
      params,
      [T <: Tuple] =>
        (n: Labels[T]) ?=>
          (updateDir: Tensor[T, Float], param: Tensor[T, Float]) =>
            param - updateDir *! learningRate - param *! weightDecay
    )

    val newMomentums = paramTree.zipMap(
      gradients.value,
      momentums,
      [T <: Tuple] =>
        (n: Labels[T]) ?=>
          (g: Tensor[T, Float], m: Tensor[T, Float]) =>
            m *! beta2 + g *! (1f - beta2)
    )

    TrainState(newMomentums, updatedParams)

// -----

package examples.basic

import dimwit.*
import dimwit.Conversions.given
import nn.*
import nn.ActivationFunctions.{relu, sigmoid}
import dimwit.random.Random
import dimwit.jax.Jit.jitReduce

import examples.timed
import examples.dataset.MNISTLoader

def binaryCrossEntropy[L: Label](
    logits: Tensor1[L, Float],
    label: Tensor0[Int]
): Tensor0[Float] =
  val maxLogit = logits.max
  val stableExp = (logits -! maxLogit).exp
  val logSumExp = stableExp.sum.log + maxLogit
  val targetLogit = logits.slice(Axis[L] -> label)
  -(targetLogit - logSumExp)

object MLPClassifierMNist:

  import MNISTLoader.{Sample, TrainSample, Height, Width}
  trait Hidden derives Label
  trait Output derives Label

  object MLP:
    case class Params(
        layer1: LinearLayer.Params[Height |*| Width, Hidden],
        layer2: LinearLayer.Params[Hidden, Output]
    )

    object Params:

      def apply(
          layer1Dim: Dim[Height |*| Width],
          layer2Dim: Dim[Hidden],
          outputDim: Dim[Output]
      )(
          paramKey: Random.Key
      ): Params =
        val (key1, key2) = paramKey.split2()
        Params(
          layer1 = LinearLayer.Params(key1)(layer1Dim, layer2Dim),
          layer2 = LinearLayer.Params(key2)(layer2Dim, outputDim)
        )

  case class MLP(params: MLP.Params) extends Function[Tensor2[Height, Width, Float], Tensor0[Int]]:

    private val layer1 = LinearLayer(params.layer1)
    private val layer2 = LinearLayer(params.layer2)

    def logits(
        image: Tensor2[Height, Width, Float]
    ): Tensor1[Output, Float] =
      val hidden = relu(layer1(image.ravel))
      layer2(hidden)

    override def apply(image: Tensor2[Height, Width, Float]): Tensor0[Int] = logits(image).argmax(Axis[Output])

  def main(args: Array[String]): Unit =

    val learningRate = 1e-4f
    val numSamples = 59904
    val numTestSamples = 9728
    val batchSize = 512
    val numEpochs = 50
    val (dataKey, trainKey) = Random.Key(42).split2()
    val (initKey, restKey) = trainKey.split2()

    val (trainX, trainY) = MNISTLoader.createTrainingDataset(maxSamples = Some(numSamples)).get
    val (testX, testY) = MNISTLoader.createTestDataset(maxSamples = Some(numTestSamples)).get

    def batchLoss(batchImages: Tensor[(TrainSample, Height, Width), Float], batchLabels: Tensor1[TrainSample, Int])(
        params: MLP.Params
    ): Tensor0[Float] =
      val model = MLP(params)
      val losses = zipvmap(Axis[TrainSample])(batchImages, batchLabels):
        case (image, label) =>
          val logits = model.logits(image)
          binaryCrossEntropy(logits, label)
      losses.mean
    val initParams = MLP.Params(
      Axis[Height |*| Width] -> 28 * 28,
      Axis[Hidden] -> 128,
      Axis[Output] -> 10
    )(initKey)

    def accuracy[Sample: Label](
        predictions: Tensor1[Sample, Int],
        targets: Tensor1[Sample, Int]
    ): Tensor0[Float] =
      val matches = zipvmap(Axis[Sample])(predictions, targets)(_ === _)
      matches.asFloat.mean

    // val optimizer = Lion(learningRate = Tensor0(learningRate), weightDecay = Tensor0(0f))
    val optimizer = GradientDescent(learningRate = Tensor0(learningRate))

    def gradientStep(
        imageBatch: Tensor[(TrainSample, Height, Width), Float],
        labelBatch: Tensor1[TrainSample, Int],
        trainState: optimizer.TrainState[MLP.Params],
    ): optimizer.TrainState[MLP.Params] =
      val lossBatch = batchLoss(imageBatch, labelBatch)
      val grads = Autodiff.grad(lossBatch)(trainState.params)
      optimizer.update(grads, trainState)
    val jitStep = jit(gradientStep)

    def miniBatchGradientDescent(
        imageBatches: Seq[Tensor[(TrainSample, Height, Width), Float]],
        labelBatches: Seq[Tensor1[TrainSample, Int]],
    )(
        trainState: optimizer.TrainState[MLP.Params],
    ): optimizer.TrainState[MLP.Params] =
      imageBatches
        .zip(labelBatches)
        .foldLeft(trainState):
          case ((trainState), (imageBatch, labelBatch)) =>
            jitStep(imageBatch, labelBatch, trainState)

    val trainMiniBatchGradientDescent = miniBatchGradientDescent(
      trainX.chunk(Axis[TrainSample], numSamples / batchSize),
      trainY.chunk(Axis[TrainSample], numSamples / batchSize)
    )
    val trainTrajectory = Iterator.iterate((optimizer.init(initParams))): trainState =>
      timed("Training"):
        dimwit.gc()
        trainMiniBatchGradientDescent(trainState)
    def evaluate(
        params: MLP.Params,
        dataX: Tensor[(Sample, Height, Width), Float],
        dataY: Tensor1[Sample, Int]
    ): Tensor0[Float] =
      val model = MLP(params)
      val predictions = dataX.vmap(Axis[Sample])(model)
      accuracy(predictions, dataY)
    val jitEvaluate = jit(evaluate)
    val (finalState, finalParams) = trainTrajectory.zipWithIndex
      .tapEach:
        case (trainState, epoch) =>
          timed("Evaluation"):
            val testAccuracy = jitEvaluate(trainState.params, testX, testY)
            val trainAccuracy = jitEvaluate(trainState.params, trainX, trainY)
            println(
              List(
                s"Epoch $epoch",
                f"Test accuracy: ${testAccuracy.item * 100}%.2f%%",
                f"Train accuracy: ${trainAccuracy.item * 100}%.2f%%"
              ).mkString(", ")
            )
      .drop(numEpochs)
      .next()

    println("\nTraining complete!")

[marcelluethi]

Looks good. The only issue is that iterator now iterate exposes the state, whereas before it exposed only the parameters. I think for simple use cases iterating only over the parameters would be sufficient. Maybe we can have an iterate and and iterateWithState method?

I also like your suggestion with opaque type Grad[T] <: T = T which shows the right behavior.

[benikm91]

Oh, the iterate change was an accidental one...

Maybe we have to better understand what the internal state is. So far I was thinking of it as only a state of the optimizer, not really as a part of the bigger picture (like params). However, for checkpointing (storing the model each epoch in case of an error - a common operation), we must store the internal training state along the parameters to successfully resume training. This suggests to me that the internal optimizer state should maybe be made more explicit (like parameters). So the (accidental) iterator change is actually the right one. Having an iterator over parameters makes the state implicit.

Very interesting :)

[marcelluethi]

I just pushed another update. After many failed attempts, I gave up with the unified interface. It might be possible for special cases, but even when the state is only the parameters signature and usage becomes clunky. My suggestion is to keep it bare bones, and let the user build his/her own abstraction if swapping out optimizers is needed. Maybe with experience we find something that abstracts this well.

As a data point, Jax/Stax do not seem to provide a unified interface.

[benikm91]

I struggled myself yesterday evening with unification, as it is not possible without making the user code harder to understand (general TrainState class, or dependent variable) - dropping this seems the right path forward to me.

We also have to make sure we do not get lost in NN stuff as DimWit we want to design well, and NN is just experimentation... So this PR is already good enough. Approved.

---

Here are some more thoughts, how we could further improve the PR, but feel free to ignore. Lets focus more on DimWit and examples:

Currently, I don't like that the Optimizer State is only a type alias. I suggest:

case class LionState[P](momentums: P)

Making things clearer:

   def batchStep(
        imageBatch: Tensor[(TrainSample, Height, Width), Float],
        labelBatch: Tensor1[TrainSample, Int],
        params: MLP.Params,
        state: LionState[MLP.Params]  // HERE
    ):

---

Some ideas I had yesterday, what do you think?

1. An optimizer object is unaware of its Parameters as the type is passed into its methods. We could move the type to the optimizer case class Lion[Params]. An optimizer object can then only be applied to these kinds of parameters.

val optimizer: Lion[Params] = Lion[MLP.Params](learningRate = learningRate, weightDecay = 0f)

For me this makes more sense, but it is technically unnecessary...

2. I started to separate updateState and updateParams. Without unification this would make concepts clearer as parameters would document dependencies. E.g. for Lion just the type of updateParams documents that the gradients are not used...

val newState = optimizer.updateState(state, grads)
val newParams = optimizer.updateParams(params, newState)

[marcelluethi]

Thanks for your comments. I think these are good ideas, but also come with trade offs. For example having Type Parameters in the class itself forces you to explicitly write the Type parameter in construction, something we don't have anywhere else in the library. Having LionState as a case class, make generalization more difficult, etc.

I like your suggestion that we continuing iterating on the design, but only after the core has stabilised.