Optimizers

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 compiled 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)

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

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

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.