Optimizers

A training step

let sched =
  Vega.Schedule.warmup_cosine_decay ~init_value:0.0 ~peak_value:1e-3
    ~warmup_steps:1000 ~decay_steps:9000 ()

let step (params, st) =
  let loss, grads = Rune.value_and_grad model objective params in
  let grads = Vega.clip_by_global_norm model ~max_norm:1.0 grads in
  let params, st =
    Vega.adamw_step model ~lr:sched ~weight_decay:0.01 st ~params
      ~grads
  in
  (loss, (params, st))

The gradients flow through each stage in order:

1. Rune.value_and_grad differentiates the loss, giving gradients with the parameters' structure
2. clip_by_global_norm rescales all gradients together so that their joint L2 norm is at most 1
3. adamw_step updates Adam's moments, derives the rate from the schedule at the state's step counter, and applies the update with decoupled weight decay

Everything in step is tensor arithmetic, so the same function runs eagerly and compiles under Rune.jit (States Are Structures).

Gradient transformations

A gradient transformation maps a gradient structure to another, and goes between the backward pass and the step:

Your own transformation is a function over the structure, usually one Nx.Ptree.map. To centralize the gradients of every matrix, for instance:

let centralize p grads =
  Nx.Ptree.map p
    (fun _ g ->
      if Nx.ndim g < 2 then g
      else
        let axes = List.init (Nx.ndim g - 1) (fun i -> i + 1) in
        Nx.sub g (Nx.mean ~axes ~keepdims:true g))
    grads

Choosing an optimizer

Weight decay in adamw_step, lamb_step, lars_step and adan_step is decoupled: it is added to the update, where the adaptive scaling does not touch it, and applies to every float leaf.

States are structures

An optimizer state is a record whose parameter-shaped fields have the parameters' structure, plus a step counter as a scalar int32 tensor. Each state type has its structure over the parameters' one: sgd_ptree, adam_ptree, rmsprop_ptree, adagrad_ptree, adan_ptree, lion_ptree, adafactor_ptree and lbfgs_ptree. A state's leaves are named by the field and the parameter's path:

Nx.Ptree.visits (Vega.adam_ptree model) st
(* mu.w, mu.b, nu.w, nu.b, step *)

Those names are the ones a saved file stores, and the structure is what a compiled step takes. With the state an argument of the step, donated at each call, a whole training step is one program that replays on every call:

let state = Nx.Ptree.pair model (Vega.adam_ptree model)

let step =
  Rune.jit
    Nx.Ptree.(tensor @-> tensor @-> state @-> returns (pair tensor state))
    (fun inputs targets (params, opt) ->
      let loss, grads =
        Rune.value_and_grad model (objective inputs targets) params
      in
      let grads = Vega.clip_by_global_norm model ~max_norm:1.0 grads in
      let params, opt = Vega.adamw_step model ~lr:sched opt ~params ~grads in
      (loss, (params, opt)))

let loss, s = step inputs targets (Nx.Ptree.donate state s)

A structure may carry tensors that are not parameters, such as an RNG key or a batch of indices. A step passes every leaf whose dtype is neither a float nor a complex one through unchanged, in the parameters and in the state.

A complex parameter steps as its two real coordinates, the real and imaginary parts, each with its own state, against the conjugate of Rune.grad's gradient: sgd_step moves z to z - lr * conj g, which descends the real objective.

Skipping overflowed steps

In float16 training, the loss is multiplied by a scale before the backward pass and the gradients are divided by it in the step. Vega.Loss_scale.step wraps a step: it divides the gradients, applies the step if they are all finite, keeps the parameters and the whole optimizer state if not, and returns the scale for the next step:

let model_state = Nx.Ptree.pair model (Vega.adam_ptree model)

let (params, st), ls =
  Vega.Loss_scale.step model model_state ls ~grads
    (fun grads -> Vega.adam_step model ~lr st ~params ~grads)
    (params, st)

The step is tensor arithmetic, so it compiles. A scale is made where the parameters lie, Vega.Loss_scale.dynamic Nx.Host.on on the host. With Loss_scale.static at 1.0 it only skips the steps whose gradients are not finite.

Next steps

• Learning Rate Schedules — decay, warmup, restarts, and composition
• Getting Started — parameters as a structure, your first optimizer
• Optax Comparison — mapping from Python's Optax to Vega

this page is tested with mdxedit this page ↗