# sexpgpu/optim

The standard optimizer module, compiled into the binary:
`(require "sexpgpu/optim" name ...)`. Every optimizer accepts a number or a
schedule for each numeric hyperparameter and calls `update-diagnostics`.

## Optimizers

| definition | keywords | states |
|---|---|---|
| `sgd` | `(lr 0.01) (weight-decay 0.0)` | none |
| `momentum-sgd` | `(lr 0.01) (momentum 0.9) (nesterov false)` | `velocity` |
| `adamw` | `(lr 1e-3) (betas [0.9 0.95]) (eps 1e-8) (weight-decay 0.0)` | `m`, `v` |
| `muon` | `(lr 0.02) (weight-decay 0.0) (momentum 0.95) (nesterov true) (steps 12)` | `momentum-buffer` |

- `sgd`: `param - lr * (grad + weight-decay * param)`.
- `adamw`: bias-corrected with `t = step + 1`; weight decay is added to the
  normalized step, then scaled by `lr`.
- `muon` follows modded-nanogpt: heavy-ball momentum, then the
  orthogonalization of that momentum by `newton-schulz`, scaled by
  `sqrt(max(1, out/in))` in PyTorch's layout, with decoupled weight decay.
  Matrices only: `(assert (= 2 (rank param)))` fails otherwise. Give the
  rest to `adamw` with a group; see
  [optimizer groups](https://sexpgpu.041.io/docs/optimizer-groups.md#a-different-optimizer-per-group).

## Helpers

| definition | signature and meaning |
|---|---|
| `schedule-value` | `(schedule-value value ctx)`: `(value ctx)` when it is a function, otherwise `value`; see [writing optimizers](https://sexpgpu.041.io/docs/optimizers.md#schedules) |
| `update-diagnostics` | `(update-diagnostics param grad next)`: the diagnostics `optimizer/grad-rms` and `optimizer/update-ratio`, both `:max` |
| `newton-schulz` | `(newton-schulz x &key (steps 12))`: the orthogonal factor of a matrix by Newton-Schulz with coefficients (2, -1.5, 0.5) from `x / ‖x‖`, in `bf16`; a tall matrix is transposed in and back out, decided at compile time |

## Schedules

| definition | signature and meaning |
|---|---|
| `stable-then-decay` | `(stable-then-decay peak &key (decay-fraction 0.7))`: `peak` until `1 - decay-fraction` of the run, then linear to zero |
| `warmup-stable-decay` | `(warmup-stable-decay peak &key (warmup-steps 0) (decay-fraction 0.7))`: the same times a linear warmup `min(1, (step + 1) / warmup-steps)`, so the first update already moves |

```lisp
(adamw :lr (warmup-stable-decay 0.0015 :warmup-steps 25) :weight-decay 0.1)
```

## Layout

Weights are stored `[in, out]`, the transpose of PyTorch's `[out, in]`.
`newton-schulz` commutes with transposition, so the `muon` update is
exactly the transpose of the modded-nanogpt update; only the scale is read
in PyTorch's layout: `out/in` is `(dim param 1)` over `(dim param 0)`.

## Diagnostics

| name | reducer | what it is |
|---|---|---|
| `optimizer/grad-rms` | max | the gradient's root mean square, the group's largest |
| `optimizer/update-ratio` | max | the step's root mean square over the parameter's: how far the update moves the parameter against its own size |

See [optimizer diagnostics](https://sexpgpu.041.io/docs/writing-diagnostics.md#optimizer-diagnostics)
for how they fold and what makes one noisy.

Related: [writing optimizers](https://sexpgpu.041.io/docs/optimizers.md), [sexpgpu/nn](https://sexpgpu.041.io/docs/nn.md).

---

SexpGPU documentation. Every page: https://sexpgpu.041.io/llms.txt