# Optimizers

posttrainllm ships five optimizers, selectable on `posttrainllm train`, `sft`, and
`dpo` via `--optimizer {adamw|lion|sophia|muon|adafactor}`. Default is
`adamw` (preserves backward compatibility — pre-existing scripts work
unchanged).

The optimizer interface is the standard MLX-Swift `Optimizer` protocol;
each implementation lives in `native-mac/Sources/TinyGPTModel/Optimizers.swift`.
The drop-in surface — same flags, same trainer, same `compile`d step —
is what lets you A/B optimisers without touching the train loop.

## Mechanism, in one paragraph each

### AdamW (Loshchilov & Hutter, 2019) — the baseline
Two EMAs per parameter — first moment `m` and second moment `v` of the
gradient — plus *decoupled* weight decay (the AdamW-vs-Adam fix).
Update is `lr · m / (√v + ε)`; weight decay is applied separately to
the parameter, not folded into the gradient. State footprint: `2× |θ|`.
Robust default for transformer pre-training.

### Lion (Chen et al., 2023) — sign-based, smaller state
Tracks only the first moment `m`. Update is `lr · sign(b1·m + (1−b1)·g)`
— the gradient is `sign`-ed away, so the per-step move has unit
magnitude in every coordinate. Sensitivity to LR is sharper (the paper
recommends LR 3-10× smaller than AdamW, weight decay 3-10× larger).
State footprint: `1× |θ|`. Sometimes beats AdamW on transformer LM at
~½ optimizer memory.

### Sophia (Liu et al., 2023) — second-order with clipping
Stores `(m, h)` where `h` is an EMA approximation of the Hessian
diagonal. Update is `lr · sign(m) · clip(|m| / (ρ·h + ε), 1)` — the
per-coordinate clip caps the per-step move at `lr`, which is what
gives Sophia robustness to bad-curvature directions. This
implementation uses the EMA-of-squared-gradient ("Sophia-light")
proxy for `h` rather than the paper's full Gauss-Newton Hessian
estimator — same memory pattern, similar dynamics on transformer LM,
no extra forward passes. State footprint: `2× |θ|` (same as AdamW).

### Muon (Jordan et al., 2024) — orthogonalised matrix update
For 2D weights (attention and MLP matrices), the momentum buffer
`m_t = β·m_{t-1} + g` is *orthogonalised* via a fixed 5-step Newton-
Schulz quintic polynomial — the result is approximately `U·Vᵀ` from
the SVD of `m`. The update is then `lr · scale · NS(m)` where
`scale = max(1, √(d_out/d_in))` keeps update norm comparable to AdamW.
For 1D weights (LayerNorm γ/β, biases) and embedding tables, Muon
falls back to AdamW internally — orthogonalisation is meaningless for
vectors. State footprint: nominally `1× |θ|` for 2D leaves (our impl
keeps the second slot present but zero for shape uniformity; a follow-
up could halve this — see "Open work" below).

### Adafactor (Shazeer & Stern, 2018) — sublinear-memory Adam
For 2D weights, stores only the row-sum and column-sum of the second
moment instead of the full matrix — the rank-1 reconstruction
`v ≈ outer(row, col) / mean(row)` is what the per-step update divides
by. For 1D weights, falls back to a regular second-moment vector. By
default beta1 = nil → no first-moment tracking either (set
`beta1 = 0.9` to add it back). Configured here with
`relativeStep=false, scaleParameter=false` so the LR scheduler
(`--lr-schedule cosine`) keeps working. State footprint: `~½ × |θ|` —
the headline claim.

## Smoke results

50-step run on `/tmp/eval-holdout-tail.txt` (5 MB raw text). LR per
optimizer follows paper recommendations. Build:

```
DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer xcodebuild \
  -scheme posttrainllm -destination "platform=macOS" \
  -derivedDataPath /tmp/posttrainllm-smoke-opts -configuration Release build
```

### Tiny preset (4L · d=128 · ctx=128 · batch=8 · 842k params)

| Optimizer  | LR    | step 1 loss | step 50 loss | step/s | peak MB | active MB |
|------------|-------|-------------|--------------|--------|---------|-----------|
| AdamW      | 3e-4  | 6.345       | 3.046        | 46.2   | 160.6   | 10.1      |
| Lion       | 3e-5  | 5.528       | 3.714        | 25.4   | 153.3   |  6.7      |
| Sophia     | 3e-4  | 6.217       | 2.927        | 30.9   | 161.2   | 10.1      |
| Muon       | 2e-3  | 5.925       | 3.450        | 18.2   | 163.5   | 10.1      |
| Adafactor  | 3e-4  | 6.185       | 3.029        | 12.1   | 148.6   |  3.4      |

### Small preset (6L · d=192 · ctx=256 · batch=4 · 3.2M params)

| Optimizer  | step 1 loss | step 50 loss | step/s | peak MB | active MB |
|------------|-------------|--------------|--------|---------|-----------|
| AdamW      | 5.861       | 2.779        | 30.8   | 404.0   | 33.2      |
| Lion       | 5.877       | 3.343        | 32.0   | 392.9   | 22.2      |
| Sophia     | 5.757       | 2.785        | 30.6   | 404.0   | 33.2      |
| Muon       | 6.102       | 3.139        | 18.8   | 404.0   | 33.2      |
| Adafactor  | 6.448       | 2.894        | 13.1   | 382.0   | 11.2      |

### Huge preset (12L · d=512 · ctx=512 · batch=2 · ~38M params)

| Optimizer  | step 1 loss | step 50 loss | step/s | peak MB | active MB |
|------------|-------------|--------------|--------|---------|-----------|
| AdamW      | 5.905       | 2.820        | 16.3   | 523.6   | 115.3     |
| Lion       | 6.332       | 3.024        | 16.3   | 485.2   |  76.9     |
| Sophia     | 6.076       | 3.058        | 11.5   | 523.6   | 115.3     |
| Muon       | 5.914       | 3.037        |  5.2   | 523.6   | 115.3     |
| Adafactor  | 5.982       | 3.033        |  7.8   | 447.1   |  38.8     |

`active MB` is MLX's "live tensor memory" snapshot at the end of the
50-step run — the most diagnostic number for *optimizer state* memory
(peak includes activations, which dominate at long context). At the
`huge` scale:

- AdamW ≡ Sophia ≡ Muon: 115.3 MB (2-state per param)
- Lion: 76.9 MB ≈ ⅔ AdamW
- Adafactor: 38.8 MB ≈ ⅓ AdamW

The Adafactor headline ("½ optimizer memory") shows up *more aggressively*
than the paper's prediction at this scale because we also drop the
first-moment buffer (Adafactor's `beta1 = nil` default). With `beta1 =
0.9`, expect Adafactor to land near 75 MB — still well under Lion.

### Tiny-preset 200-step trace

To verify all five continue converging, not just take a single drop:

```
adamw     5.748 → 2.832 → 2.729 → 2.552 → 2.624   (200 steps · 61.5 step/s · final 2.624)
lion      5.809 → 3.871 → 3.609 → 3.497 → 3.178   (200 steps · 77.4 step/s · final 3.178)
sophia    5.792 → 2.699 → 2.701 → 2.532 → 2.655   (200 steps · 72.7 step/s · final 2.655)
muon      6.159 → 3.397 → 2.775 → 2.896 → 2.631   (200 steps · 50.2 step/s · final 2.631)
adafactor 6.238 → 2.903 → 2.853 → 2.700 → 2.638   (200 steps · 74.1 step/s · final 2.638)
```

All five descend monotonically (with the usual noise band). At 200
steps on this small problem, AdamW / Sophia / Adafactor / Muon are
within 0.05 nats of each other; Lion lags by 0.5 nats because the
sign-based update needs more steps to find the right scale (the paper
notes Lion typically catches up by 1k-2k steps on LM tasks).

## When to pick which

### Use AdamW when
- You have no specific reason to deviate. AdamW is the default for a
  reason — robust, well-understood, predictable hyperparameter
  behaviour across model scales.

### Use Lion when
- You're memory-bound by optimizer state at fp32, and want a quick win
  (~33% off optimizer memory) without changing model architecture.
- Your run is long (>5k steps) so Lion has time to catch up to AdamW
  on the loss curve. On <500-step smoke runs Lion often looks worse —
  this is expected.
- Pair with a 3-10× smaller LR than AdamW (Lion's `sign`-based update
  has unit magnitude, so the same LR translates to a 10×-100× bigger
  move).

### Use Sophia when
- The loss surface has very different curvatures across coordinates
  (the per-coordinate clip in Sophia's update is the part that matters
  here). On standard LM pre-training, expected wins are 10-20% fewer
  steps to the same loss vs AdamW.
- Compute budget allows ~10-15% slower steps (Sophia is more arithmetic
  per parameter than AdamW). On our tiny-preset trace Sophia is 25-30%
  slower because the small parameter count makes the overhead more
  visible; on larger models the gap shrinks toward 10%.
- Note: this is the EMA-of-squared-gradient ("Sophia-light") variant,
  not the full Gauss-Newton Hessian estimator. The paper-recommended
  Sophia-G is a future enhancement (would need an extra forward pass
  every k steps).

### Use Muon when
- Your model is dominated by 2D matrix weights (attention/MLP). The
  orthogonalisation pre-conditions the update step to look like an
  SVD-truncated update, which trains substantially better on those
  matrices in the Muon paper's settings.
- Step time is acceptable — the Newton-Schulz iteration is 5 matmuls
  per 2D param per step. On our `huge` preset Muon was 3× slower than
  AdamW (5.2 vs 16.3 step/s); on smaller models the relative slowdown
  is smaller (1.5-2× on tiny/small).
- LR can be tuned generously larger (paper uses 2e-3, an order of
  magnitude above AdamW's typical 1e-4 to 3e-4).

### Use Adafactor when
- The model is large enough that fp32 optimizer state dominates
  per-step memory and you can't afford bf16 (numerical stability) or
  ZeRO-style sharding (single-machine). At 38M params and fp32,
  Adafactor saves ~80 MB; at 7B fp32 it would save ~15 GB.
- bf16/fp16 training: Adafactor's relative-rms scaling makes it
  particularly bf16-friendly — the second-moment is already noisy, so
  the factorised approximation introduces less additional error than
  it would on fp32. Pair `--optimizer adafactor` with `--dtype
  bfloat16` for the most aggressive memory profile in this repo.

## Implementation notes

### LearningRateMutable

The `Trainer` schedule code does `trainer.optimizer.learningRate = ...`
each step to drive cosine warmup/decay. To preserve this across
optimizers (AdamW, Lion, Adafactor, Sophia, Muon), we added the
`LearningRateMutable` protocol with a single `var learningRate: Float`
requirement. Adafactor's stored property is `Float?` (supports
relative-step mode), so it's wrapped in `AdafactorAdapter` that
round-trips through the optional. We also lock Adafactor into
`relativeStep: false, scaleParameter: false` mode in the factory so
the optional is always populated.

### Compile compatibility

All optimizers conform to MLX-Swift's `Optimizer: Updatable,
Evaluatable` protocol, so the compiled train step
`compile(inputs: [m, optimizer], outputs: [m, optimizer]) { ... }`
works unchanged. The training closure's `optimizer.update(model:,
gradients:)` call is dispatched dynamically through the protocol; MLX
traces through it the same way it traces AdamW.

### Why we re-implement Sophia/Muon instead of inheriting OptimizerBase

`MLXOptimizers.OptimizerBase`'s synthesised initializer has internal
access, so cross-module subclasses fail to compile. Sophia and Muon
implement the `Optimizer` protocol directly, maintaining their own
`NestedDictionary<String, PairState>` state storage and a
`mapValues`-driven `update(...)` that mirrors `OptimizerBase.apply()`.

## Open work

- **Muon's state**: our `PairState` has two MLXArrays per param even
  though the 2D-Muon path only uses one. Reclaiming that half (custom
  `Updatable` enum state) would bring Muon's memory profile in line
  with Lion's. Punted for clarity; the headline "Newton-Schulz update"
  behaviour is unchanged.

- **Sophia-G**: the Gauss-Newton Hessian estimator (the paper's
  preferred variant) requires an extra forward + per-token logit
  sampling every `k` steps. Plumbable through `Trainer.step` as a
  callback; not done here to keep the drop-in story simple.

- **Lion LR autotune**: Lion's "right" LR is typically 3-10× smaller
  than AdamW's. A helper that scales the user's `--max-lr` when
  `--optimizer lion` is chosen would smooth the swap, at the cost of
  hiding paper-recommended behaviour.

## Caveats from the 50-step smoke

- **50 steps is too few to rank optimisers on a loss curve.** Lion in
  particular benefits from 2k+ steps. The smoke results here verify
  *correctness* (loss decreases, memory measurements are stable), not
  end-of-training quality.

- The `huge`-preset Muon step rate (5.2 step/s) was measured against
  GPU contention from another training process (the mega-v2 flagship
  run on the same machine). Solo numbers will be 1.5-2× higher.

- Adafactor's step rate appears low at 50 steps because the first few
  steps trigger MLX trace recompilation more aggressively (the
  factorised second-moment has different shapes per leaf, and our
  config sets `relativeStep=false` so the per-step graph is more
  uniform after warmup). At 200 steps Adafactor reaches 77.8 step/s on
  tiny — the same neighbourhood as AdamW/Sophia/Lion.