# Writing optimizers

An optimizer is a function from a parameter, its gradient, its states and
`ctx` to the new parameter and the new states. The body is traced once per
parameter into that parameter's update graph.

```lisp
(require "sexpgpu/optim" schedule-value update-diagnostics)

(defoptimizer momentum-sgd (param grad ctx &key (lr 0.01) (momentum 0.9))
  (setq lr (schedule-value lr ctx))
  (defstate velocity (zeros-like param))
  (let* ((v (+ (* momentum velocity) grad))
         (next (- param (* lr v))))
    (update-diagnostics param grad next)
    (optimizer-update next :velocity v)))
```

## defoptimizer

`(defoptimizer name (param grad ctx &key ...) body...)` defines an
optimizer. Calling it with keyword values configures it:
`(momentum-sgd :lr 0.05)` is the value the contract's `optimizer` holds,
and what `group :optimizer` accepts; see
[optimizer groups](https://sexpgpu.041.io/docs/optimizer-groups.md).

- Traced once per parameter, so a shape-dependent choice is a plain
  compile-time `if`: Muon transposes a tall matrix by asking
  `(> (dim x 0) (dim x 1))`.
- Each update graph owns one parameter; there is no cross-parameter state.
- A body reports with `diagnostic`, never `metric` (`E-CONTRACT-011`); see
  [optimizer diagnostics](https://sexpgpu.041.io/docs/writing-diagnostics.md#optimizer-diagnostics).

## States

`(defstate name init)` declares one per-parameter state.

- `init` evaluates normally and must be a zero constant: `zeros`,
  `zeros-like` or a zero `full`, including casts, reshapes, transposes and
  broadcasts of those. Helpers, conditionals and macros may return it.
- Nonzero or computed values (even ones that happen to be zero),
  duplicate states and a `defstate` outside an optimizer are `E-OPT-003`.
- Ordinary states are named in the IR and checkpoints by their names;
  gensym states get stable `state.N` names in declaration order.

## optimizer-update

`(optimizer-update new-param :state new-value ...)` is the body's result.

- Every declared state is returned exactly once. A body that does not end
  in `optimizer-update`, or drops a state, is `E-OPT-004`.
- A state is named by keyword, `:velocity`, or by its quoted declaration
  symbol, `',velocity`, which is how a macro returns a gensym state. Never
  turn a gensym into a keyword to name its state.
- It is an ordinary function, so `apply` can pass a computed list of
  updates.

## Schedules

A schedule is a function of `ctx` returning a scalar. Called while the
update graph is traced, it builds ordinary nodes in it.

```lisp
(defun stable-then-decay (peak &key (decay-fraction 0.7))
  (lambda (ctx)
    (let ((progress (getf ctx :progress)))
      (* peak (minimum 1.0 (/ (- 1.0 progress) decay-fraction))))))
```

To accept a number or a schedule, call `(schedule-value lr ctx)` from
`sexpgpu/optim` in the body: it calls a function with `ctx` and returns
anything else unchanged. The standard optimizers do this for every declared
hyperparameter. Other function arguments, a gradient transform for example,
stay callable by the body. The `ctx` keys are in
[curriculum and ctx](https://sexpgpu.041.io/docs/curriculum.md#ctx).

Related: [sexpgpu/optim](https://sexpgpu.041.io/docs/optim.md), [optimizer groups](https://sexpgpu.041.io/docs/optimizer-groups.md),
[macros](https://sexpgpu.041.io/docs/macros.md#a-worked-example).

---

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