mlx-bun 0.0.4 → 0.0.6

package/docs/design/orpo-dynamic-lambda.md

@@ -0,0 +1,179 @@
+# Dynamic λ controller for ORPO (adaptive preference pressure)
+
+**Status:** proposed (design only — not built). **Opened:** 2026-06-20.
+**Surface:** `src/train/loss.ts` (orpoMetrics/orpoLoss), `src/train/trainer.ts`
+(orpoLoop, the λ config), the segmented ORPO classes in `src/train/segmented.ts`.
+
+## Motivation
+
+ORPO's loss is `L = NLL(chosen) + λ·OR`. The two terms learn on different
+timescales, and a **static λ** can't serve both:
+
+- **The SFT/NLL half learns easily** — the model quickly raises `logp(chosen)`
+  (fits the response *style*). Observed live: val loss falls steadily while
+  preference accuracy/margin barely move.
+- **The preference/OR half is hard** — on subtle data (UltraFeedback on a 1B) the
+  model doesn't generalize "which of two plausible answers is better."
+- **Over-weighting OR backfires** via **likelihood displacement**: the model widens
+  the *relative* margin by **crushing `logp(rejected)`** (and sometimes dropping
+  `logp(chosen)` too) rather than genuinely preferring chosen — degrading fluency.
+
+A fixed λ either under-pressures the preference (λ too low → never leans on OR) or
+over-pressures it (λ too high → displacement/degeneration). The hypothesis: an
+**adaptive, displacement-aware λ** — low while SFT learns, ramped once SFT
+saturates, backed off the moment displacement appears — can find a generalization
+regime a static λ sails past. (Static λ tuning was the blunt version we rejected;
+this is the targeted version.)
+
+## Control signals (what we measure)
+
+`orpoMetrics` (`loss.ts:1249`) already computes everything; we just don't surface it:
+
+- **`lw`** = mean `logp(chosen)` (currently folded into `nll = mean(−lw)`).
+- **`lr`** = mean `logp(rejected)` (currently only used inside the odds ratio).
+- **`margin`** = mean log-odds (chosen vs rejected) — already emitted.
+- **`accuracy`** = fraction `lw > lr` — already emitted.
+
+The controller keys on the **decomposition of margin growth**:
+- margin ↑ because **`lw` ↑** → *healthy* (genuinely preferring chosen).
+- margin ↑ because **`lr` ↓** while `lw` flat/falling → *displacement* (the cheat).
+
+So the load-bearing new instrumentation is **surfacing `lw` and `lr` separately**
+per eval (not just `nll`/`margin`).
+
+## The controller (state machine)
+
+Runs at **each val eval** (every `stepsPerEval` steps) — val (`lw`, `lr`) is the
+stable signal; per-step B=1 train metrics are too noisy to control on. Maintains a
+short window of recent val (`lw`, `lr`, margin, accuracy).
+
+```
+state: { phase, lambda, window[] }   # lambda starts at LAMBDA_MIN
+on each val eval (lw, lr, margin, acc):
+  dLw = lw - lw_prev ; dLr = lr - lr_prev
+
+  PHASE 1 — SFT warmup (lambda = LAMBDA_MIN, low):
+    stay while the NLL is still improving (dLw > EPS_SFT, i.e. chosen logp rising).
+    when dLw plateaus (|dLw| < EPS_SFT over K evals) → SFT saturated → PHASE 2.
+
+  PHASE 2 — preference ramp:
+    each eval, lambda += RAMP (cap LAMBDA_MAX), applying OR pressure —
+    AS LONG AS chosen holds: dLw >= -EPS_DISP.
+    if accuracy/margin respond (margin widening with dLw >= 0) → keep ramping.
+    if displacement detected (dLw < -EPS_DISP, i.e. chosen logp dropping, or
+       dLr << 0 with dLw <= 0) → PHASE 3.
+
+  PHASE 3 — displacement backoff:
+    lambda -= BACKOFF (floor LAMBDA_MIN) until chosen recovers (dLw >= 0),
+    then hold (hysteresis) before considering another ramp.
+```
+
+Knobs (all tunable): `LAMBDA_MIN≈0.05`, `LAMBDA_MAX≈0.5`, `RAMP/BACKOFF≈0.05`,
+window `K≈3` evals, `EPS_SFT`/`EPS_DISP` set from the observed val-`lw` noise scale.
+Hysteresis (require K consecutive evals before a phase flip) prevents oscillation.
+
+**Degeneration guard (hard stop):** if `lw` falls below its phase-1 starting value
+by more than a threshold, freeze λ at `LAMBDA_MIN` and log — the model is paying for
+margin with chosen's likelihood, which a backoff alone isn't recovering.
+
+## Control-theory framing (PID / AIMD) — added 2026-06-23
+
+The phase state-machine above is one instance of a more general and more honest
+frame: **this is closed-loop control.** The **process variable** is the SFT signal
+(NLL / `lw`), the **actuator** is the preference weight (ORPO's λ, *or* SimPO's
+β/γ — the controller is objective-agnostic), and the **control law** is what maps the
+measured degradation back to the actuator. Two laws worth naming:
+
+- **AIMD** (additive-increase / multiplicative-decrease — TCP congestion control):
+  ramp the preference up *gently* while the SFT is healthy, **slam it down hard** the
+  moment degradation ("congestion") appears. Famously robust and forgiving of its own
+  knobs. The PHASE-2 ramp / PHASE-3 backoff is essentially AIMD.
+- **PID** decomposes the response into three terms, each fixing a failure the others
+  can't:
+  - **P (proportional · the level):** how degraded is `lw` *right now* → back off in proportion. (Leaves a steady-state offset.)
+  - **D (derivative · the slope):** is degradation *accelerating* → a **leading indicator** that catches the helping→hurting **inflection before the level rises**, so you pull back *earlier* and leave less damage to heal. **Gotcha:** differentiating a noisy signal amplifies noise → use a *robust* slope (slope of the EMA, or a windowed regression), never raw step deltas. The slope is where most of the value is.
+  - **I (integral · accumulated error):** the "missing letter." P leaves a chronic offset and D ignores it (flat slope at steady-state) → the model can settle hovering with a *small persistent* NLL elevation = quiet permanent SFT damage. I accumulates that residual and drives `lw` **back to its floor**. **Gotcha — integral windup:** if the actuator is pinned (λ at floor) while error persists, the integral piles up and **overshoots** on recovery → needs **anti-windup** (clamp/stop accumulating when saturated).
+
+**Build it incrementally — P → PD → PID** — each term earning its place against an
+*observed* failure of the simpler controller (P oscillates/offsets → add D; chronic
+offset remains → add I with anti-windup). That's both good control practice and an
+**ablation that writes the experimental section**, and it stops us from tuning three
+gains blind.
+
+**Why a controller beats a sweep (the real argument):** the optimal preference
+pressure is **non-stationary** — fragile early (push less), robust late (push more) —
+so **no fixed λ is optimal.** A sweep finds the best *constant*; the controller tracks
+the moving target and finds the best *trajectory*, which strictly dominates. The pitch
+isn't just "one run instead of six," it's "**beats any fixed value** a sweep could
+find." (The displacement backoff is also why this is *more* than auto-tuning — it
+targets the degradation signal directly.)
+
+**Degenerate failure to guard (beyond the §degeneration hard-stop):** if the law backs
+off too eagerly it **collapses the preference term to ~0 → you've trained pure SFT.**
+The AIMD/PID balance must **favor pushing** (aggressive-ish increase, retreat only on
+*confirmed* degradation), and target a small **safety margin below the cliff** rather
+than the absolute max-bearable (the edge is fragile).
+
+**Synergy with the diffusion frontier:** all of this is noise-robustness machinery for
+controlling off a noisy process variable — so it transfers *especially* well to a
+diffusion LM, where the likelihood is a stochastic ELBO (even noisier). See
+[orpo-base-uf-experiment-and-directions.md](../investigations/orpo-base-uf-experiment-and-directions.md) §6.7.
+
+## Wiring
+
+1. **Instrument (small):** add `lw`, `lr` to `orpoMetrics`'s return + the trainer's
+   `metric` emit (train and val). `pref-control.ts` / `metrics.json` log them so the
+   controller (and our plots) can see the decomposition.
+2. **Make λ dynamic:** today `cfg.orpoLambda` is a constant threaded into
+   `orpoLoss(model, batch, cfg.orpoLambda, …)` and **baked into the segmented ORPO
+   classes' constructors** (`new SegmentedBackwardOrpo(model, lora, ranges,
+   cfg.orpoLambda, …)`). Change to a **mutable λ source**:
+   - simplest: a `{ value: number }` box (or a getter) passed instead of the scalar;
+     the loss / segmented step reads `lambdaRef.value` each step.
+   - the segmented classes expose a `setLambda(x)` (or read the ref) so the
+     constructor-baked constant becomes a live field.
+3. **The controller** lives in `orpoLoop` (`trainer.ts`): after each val eval, feed
+   the val (`lw`, `lr`, margin, acc) to `updateLambda(state)` and write the new λ
+   into the ref. Log every λ change (`· λ 0.10 → 0.15 (SFT saturated, ramping)`).
+4. **Off by default:** gate behind a config flag (`orpoLambdaSchedule: "static" |
+   "adaptive"`); static = today's behavior, exact.
+
+## Experiment design (how we validate it)
+
+Head-to-head, **same data / seed / everything else**, on the from-base UF setup:
+- **Arm A:** static λ=0.1 (the current run — the baseline).
+- **Arm B:** adaptive λ (this controller).
+
+Compare:
+- **val accuracy + margin** trajectories (does B clear the noise floor / widen
+  margin where A stalls?);
+- **`lw` trajectory** (does B avoid displacement — `lw` non-decreasing — where a
+  high static λ would crater it?);
+- the **downstream** numbers: Exp-1 capability x2 and Exp-2 win-rate (does B's model
+  actually win more head-to-head?).
+
+Success = B generalizes measurably better (val acc clears ~0.55 with margin lifting,
+and/or Exp-2 win-rate > A's) **without** degeneration (`lw` holds). Null result is
+also informative: "even adaptive λ can't beat the 1B capacity ceiling on UF."
+
+## Open questions / risks
+
+- **Cadence:** controlling on val every `stepsPerEval` may be too coarse (few control
+  steps per run). Option: a cheaper, more frequent "control eval" on a small fixed
+  probe set distinct from the reported val set.
+- **Confounds the replication claim:** an adaptive-λ result is no longer "standard
+  ORPO" — it's our method. Keep Arm A (static 0.1) as the honest paper-replication;
+  Arm B is the "can we do better" research arm. Report both.
+- **Tuning the controller's own knobs** risks the same trap (we'd be tuning a tuner).
+  Mitigate: set EPS thresholds from observed noise scale, not by hand-fitting to one
+  run; validate the controller on the *positive control* (prefer-uppercase) first —
+  it should ramp λ early (SFT saturates fast) and never trip displacement.
+
+## References
+
+- `src/train/loss.ts:1249` `orpoMetrics` (computes `lw`,`lr`,`nll`,`or`,`accuracy`,`margin`).
+- `src/train/trainer.ts` `orpoLoop` (λ config, val emit), `DEFAULT_TRAIN_CONFIG.orpoLambda`.
+- `src/train/segmented.ts` segmented ORPO classes (λ baked in constructor).
+- `scripts/experiments/pref-control.ts` (trajectory logging — extend with `lw`,`lr`).
+- Context: [[opssdpa-dk-vjp-bug]] (validate grads vs autograd), the dynamic-λ
+  discussion + the live from-base UF val trajectory (2026-06-20).

package/docs/design/orpo-future-enhancements.md

@@ -0,0 +1,340 @@
+# ORPO future enhancements
+
+Performance-first follow-on ideas for mlx-bun's ORPO trainer, plus a short
+appendix of objective variants that may be useful once the runtime path is
+comfortable.
+
+> **LANDED since this was written (2026-06-19):** the top items here are now in
+> production — **Priority 1 (Cut Cross Entropy / fused linear CE)** and **§1
+> (vocab-axis chunked CE)** shipped as the `[M,vocab]`-free **flash-CCE Metal head**
+> (verbatim MLX steel GEMM, fwd+bwd; `orpo_flash_ce`), and **§3 (prefix sharing in the
+> trainer)** shipped as `orpo_prefix_shared`, composed with the flash head AND the
+> segmented backward for MiniCPM5 **and** e4b. See [training.md](../reference/training.md),
+> [orpo-quickstart.md](../reference/orpo-quickstart.md), and the live record
+> [steel-flash-cce-handoff.md](../investigations/steel-flash-cce-handoff.md). The
+> remaining sections (§2, §4–§9, the objective variants) are still genuinely future.
+
+## Deferred levers — queued, do NOT start mid-run (2026-06-19)
+
+Two optimizations identified during the MiniCPM5 UltraFeedback ORPO run. **Both are
+deferred until the active run finishes** — the GPU is fully occupied, so these get
+*measured* later, not "quickly tested" now (a quick probe steals the training GPU).
+
+1. **Expose `--grad-accum` (gradient accumulation).** The machinery already exists —
+   `accumulateStep` (`src/train/trainer.ts:975`) loops `gradAccumSteps` B=1 micro-steps,
+   **averages** their grads (`scale = 1/accumSteps`, line 989), and evals+frees each
+   micro-graph so **peak memory stays flat** (line 1015). It's wired into the ORPO loop
+   (`trainer.ts:837`) but `gradAccumSteps` defaults to 1 and has **no CLI flag**.
+   - **Why it helps us:** ORPO here is hard **B=1** (segmented + prefix-share assume a
+     single stream — `segmented.ts:636` throws on B>1), so every optimizer update is
+     computed from exactly **one** preference pair — the noisiest possible gradient (the
+     per-step loss swinging 0.4↔3.3 is that noise). Grad-accum gives an effective batch of
+     G — averaging G pairs/update, variance ↓ ~1/G, steadier descent, tolerates a higher
+     LR — at **B=1 memory**. It's the *only* way to a bigger effective batch when B>1 is
+     impossible.
+   - **What it is NOT:** not a speedup (G× passes per optimizer update → same total
+     passes, same wall-clock, same data seen); and not a guaranteed win (the prior run
+     converged fine at B=1, val 1.66→1.50). It's a **try-and-measure** lever — expose the
+     flag, A/B G=4/8 judged on **val-margin**, don't assume.
+   - **Note on prefix-share:** the existing `orpo_prefix_shared` is *already* a hand-built
+     batch-of-2-with-shared-prefix (1F + 1B over `[P; Rc; Rr]`, prompt deduped in **both**
+     directions — strictly cheaper than a generic 2-batch's `2F + 2B`). Generic N-way
+     batching would re-duplicate per-pair prefixes and need a block-diagonal mask layer
+     *on top of* the prefix-share mask — high friction, and the payoff is dominated by
+     grad-accum (effective batch, free) + W4A16 (the real GPU bottleneck). Park batching
+     unless a profile ever shows real short-row GPU idle.
+
+2. **W4A16 compute-precision spike.** The head GEMM dequants 4-bit weights into **fp32**
+   and matmuls in fp32 (W4A32); dropping the operands to `half` (fp32 accumulate) attacks
+   the actual bottleneck. Confirmed **compute-bound live** (GPU 100% / BW ~⅓ of 400 on
+   M1 Max). Full plan + parity gating: [w4a16-compute-precision-spike.md](w4a16-compute-precision-spike.md).
+
+**Context — seg vs no-segment A/B (MiniCPM5, seq 4096, same-machine M1 Max):**
+`--no-segment` measured **~1.25 s/step steady @ peak 9.26 GB** vs `--seg 2`
+**~1.58 s/step @ 2.62 GB** — ~20% faster for 3.5× peak memory, **bit-identical output**
+(segmenting is just gradient checkpointing). Dropping the segs removed the recompute (real
+GPU work) → GPU busy but no longer pinned. Fine when memory is spare; the current 3-epoch
+run uses `--no-segment`.
+
+Product relevance: ORPO is not only a researcher feature. It is the path to
+small, downloadable task adapters that make the default local assistant better
+without forcing every user onto a larger base model. The performance work below
+matters when it turns adapter training from a fragile lab run into a repeatable
+part of the mlx-bun product loop.
+
+## Priority 1: Cut Cross Entropy / Liger Fused Linear CE
+
+The top ORPO performance priority is a Cut Cross Entropy / Liger-style Fused
+Linear Cross Entropy path for the LM head. This should be treated as the next
+major kernel project before lower-impact objective variants.
+
+Why it is first:
+
+- Gemma's vocab is huge, so materializing `[tokens, vocab]` logits dominates
+  peak memory in ORPO.
+- ORPO needs chosen and rejected response log-probs, so the head cost is paid
+  twice unless the trainer can share or fuse more work.
+- Token-axis chunking already helps fit the run, but vocab-axis chunking cuts
+  the largest remaining logits allocation at the source.
+- The same fused engine can support SFT, DPO, ORPO, SimPO, IPO, and hinge-style
+  preference losses once it exposes mean log-prob with a custom upstream
+  cotangent.
+- A Metal fused version can avoid extra MLX dispatches and intermediate memory
+  traffic in the hot training path.
+
+The implementation target is not just "chunk logits"; it is a fused log-prob
+head that computes the linear projection, online vocab logsumexp, target-logit
+gather, loss reduction, and backward replay without ever storing full logits.
+
+## Performance focus
+
+### 1. Vocab-axis chunked cross entropy
+
+The current ORPO head optimization is **token-axis chunking**: split the response
+positions into chunks, run the LM head on `[chunk, vocab]`, and rematerialize
+those chunks in the backward. That bounds peak logits by `chunk_tokens * vocab`.
+
+The complementary trick is **vocab-axis chunking**: split the vocabulary instead,
+run the head on `[tokens, vocab_chunk]`, and combine chunks with an
+online-logsumexp recurrence. This is the idea behind **Cut Cross Entropy** and
+Liger-style **Fused Linear Cross Entropy**. This is priority 1 for the ORPO
+performance roadmap.
+
+Exact forward sketch:
+
+1. For each vocab chunk, compute partial logits `h @ W_chunk`.
+2. Gather the target logit only if the target id falls inside this chunk.
+3. Maintain per-token online `logsumexp` across vocab chunks.
+4. At the end, `logp = target_logit - lse`.
+
+Backward reconstructs the softmax chunk-by-chunk from the saved per-token LSE
+and never materializes full `[tokens, vocab]` logits. For Gemma's 262k vocab,
+this is likely the next real moat after segmented backward.
+
+Two implementation levels:
+
+- Pure MLX chunked VJP first: exact, easier to validate, probably enough to fit
+  longer contexts.
+- Metal fused kernel later: fuse quantized head matmul + online LSE + target
+  gather + backward. This is the high-performance version.
+
+### 2. Fold ORPO gradients into the fused CE path
+
+ORPO needs two mean log-probs, `lw` and `lr`, then applies the odds-ratio loss.
+For a fused CE/log-prob kernel, the clean design is to feed the upstream scalar
+gradient for each branch back into the chunked head backward. In other words,
+the fused head should expose "mean logp with custom upstream cotangent", not
+only cross-entropy.
+
+This avoids materializing logits while still letting ORPO, SimPO, and future
+preference losses share the same head engine.
+
+### 3. Prefix sharing in the trainer
+
+`src/train/prefix-shared.ts` already proves the MiniCPM5 shared-prefix idea:
+one forward over `[prompt; chosen; rejected]` with block-sparse masking and
+block-wise RoPE. It is not yet wired into `orpoLoop`.
+
+Next useful path:
+
+- Add an opt-in flag, for example `orpo_prefix_sharing`.
+- Gate it by prompt/response ratio so short-prompt data stays on the simple
+  two-forward path.
+- Wire MiniCPM5 first.
+- Port e4b later, where sliding-window masks, donor-KV reuse, and per-layer
+  input make it trickier.
+
+### 4. Sequence-chunk the MLP sub-block
+
+The MLP is per-token independent. Inside segmented backward, run MLP blocks over
+sequence chunks so the GeGLU/SwiGLU intermediate is `chunk_tokens * d_ff` instead
+of `seq_len * d_ff`.
+
+This stacks with depth segmentation and vocab-axis CE:
+
+`layer segment -> layer -> MLP seq chunk -> vocab chunk`
+
+Attention cannot be chunked this naively, but the MLP can.
+
+### 5. Rotor / budget-aware segment planner
+
+Replace fixed `segment_size` with a planner that chooses segment boundaries from
+estimated activation bytes and a target memory budget. The planner should isolate
+expensive full-attention layers and larger MLP regions instead of treating every
+layer equally.
+
+This is the practical version of rotor/checkpoint scheduling for this codebase:
+less hand tuning, better fit at 8K+.
+
+### 6. Faster flash-attention backward
+
+The custom flash backward is now correctness-useful but too slow. A serious
+performance version should mimic MLX's tuned SDPA style: simdgroup tiling,
+larger threadgroups, and FA2-style backward that stores only per-row LSE and
+recomputes attention probabilities.
+
+This matters most when `ops.sdpa` backward's O(L^2) memory is the blocker.
+
+### 7. Fused LoRA epilogue
+
+Today LoRA adds multiple ops after the quantized base matmul:
+
+`base(x) + scale * ((x @ A) @ B)`
+
+A fused dequant-GEMM + LoRA epilogue would reduce dispatches and intermediate
+traffic. Backward can recompute the cheap rank projection:
+
+- `h = x @ A`
+- `dB = scale * h.T @ dy`
+- `dA = scale * x.T @ (dy @ B.T)`
+- `dx = dy @ W.T + scale * (dy @ B.T) @ A.T`
+
+This is more kernel work, but it benefits SFT, DPO, ORPO, and serving with
+adapters.
+
+### 8. Shape-specialized compile for training graphs
+
+Shapeless compile is mainly required when a dimension changes every replay, such
+as autoregressive decode with a growing KV cache:
+
+`active_kv_len = prompt_len + generated_tokens`
+
+ORPO training is different. Each dataset row already has concrete lengths:
+
+- `P`: prompt length
+- `Rc`: chosen response length
+- `Rr`: rejected response length
+- `T`: prefix-shared length, `P + Rc + Rr`
+
+That means ORPO does not need shapeless compile in the same way decode does. A
+training graph can be compiled for an exact `(P, Rc, Rr)` shape, or for a padded
+length bucket, and cached under a key that includes the model, flags, segment
+plan, token/vocab chunk sizes, and objective mode.
+
+This is especially relevant for e4b prefix-sharing: the proof path should stay
+eager first, but a later L3 speed tier can use shape-specialized compiled graphs
+for known training shapes instead of requiring every primitive to support
+shapeless output-shape inference.
+
+The deeper L3 framing is **model-checksum-specific specialization**. For a fixed
+model snapshot, we know the architecture, layer count, hidden sizes, quantization
+layout, RoPE geometry, vocab, layer types, KV-sharing pattern, and supported
+context budget in advance. It is acceptable to spend much more work up front if
+that buys a faster hot path later.
+
+Practical version:
+
+- Treat the model snapshot checksum as part of the compilation key.
+- Generate/compile/cache a table of exact graph variants for the shapes and
+  modes we care about.
+- Bucket or pad examples into those graph variants when helpful.
+- Fall back to eager only for shapes/modes not in the table.
+
+This is effectively a lookup table of specialized execution plans: more compile
+and analysis cost up front, less branching, shape handling, and dispatch overhead
+at runtime.
+
+### 9. Fixed-shape decode cache as a compile route
+
+Decode is harder because the active KV length grows by one token per step. The
+current shapeless-compile path exists to avoid recompiling for:
+
+`L = 17, 18, 19, 20, ...`
+
+A separate route is to make decode shapes fixed:
+
+1. Preallocate K/V buffers to a known maximum context length.
+2. Write the new token into a dynamic position.
+3. Carry `active_length` / `write_position` as scalar array inputs.
+4. Mask out inactive cache positions, or preferably use an
+   active-length-aware attention kernel that simply does not scan them.
+
+The correctness version is "fixed cache + mask". The performance version is
+"fixed cache + active-length-aware attention"; otherwise early decode steps may
+pay for the full max context even when only a short prefix is active.
+
+There is also a more aggressive specialization route: prebuild a graph table for
+specific active lengths or length buckets. Even if that means many compiled graph
+variants, the trade can be valid for a fixed model/checksum if compile cost is
+paid once and reused. In that regime, the important design question is not "can
+one shapeless graph infer every shape?" but "which finite set of exact shapes and
+runtime states should we specialize for?"
+
+If every op in a chosen route is compile-safe, the whole decode graph can become
+a stable compiled graph. Custom Metal kernels still need to stay outside compiled
+regions unless the route is intentionally eager/custom-kernel-only or MLX compile
+can accept the primitive's output contract.
+
+## Quality and robustness ideas
+
+These are lower priority if the current goal is performance, but they are useful
+once ORPO can run comfortably.
+
+### SimPO
+
+SimPO is a reference-free preference loss that uses the **average log-probability
+of each response** as the reward and adds a target reward margin. It reuses the
+same `lw` / `lr` values ORPO already computes.
+
+Rough shape:
+
+`loss = softplus(-beta * ((lw - lr) - gamma))`
+
+where `gamma` is the desired chosen-vs-rejected margin. It should be cheap to add
+as a sibling objective because the expensive part, response mean log-probs, is
+already implemented.
+
+### IPO
+
+IPO means **Identity Preference Optimization**. It was introduced as a more
+overfit-resistant alternative to the original DPO logistic objective. Instead of
+driving preference probabilities toward hard 0/1 labels indefinitely, IPO uses a
+smoother identity-style objective over preference gaps. The practical point:
+when preference labels are noisy, IPO is often less eager to push margins to
+extremes.
+
+For mlx-bun, IPO is interesting mainly as a loss-family option after the shared
+logp engine exists.
+
+### Hinge / SLiC-style loss
+
+The hinge loss is a margin-ranking objective:
+
+`loss = max(0, margin - reward_gap)`
+
+For preference training, `reward_gap` is usually the chosen score minus rejected
+score. Once the gap exceeds the margin, that pair contributes no gradient. This
+can be useful when you want to stop over-optimizing already-solved pairs and
+focus compute on pairs whose preference margin is still too small.
+
+### Label smoothing and robust weighting
+
+Preference pairs are noisy. Cheap robustness knobs:
+
+- `label_smoothing`: treat a pair as, for example, 95% chosen / 5% rejected
+  rather than a hard label.
+- Per-example `weight`: let curated data carry confidence or score gap.
+- Loss clipping / winsorization: cap extreme high-loss pairs so noisy examples
+  do not dominate gradients.
+- Ambiguity filtering: skip pairs with tiny initial margin or near-identical
+  chosen/rejected responses.
+
+### Better curation
+
+The current UltraFeedback curation mostly extracts rows and filters by length.
+Quality can likely improve by filtering for:
+
+- Large enough chosen-vs-rejected score gap when available.
+- Similar response length, or at least bounded length ratio.
+- High semantic contrast but not unrelated answers.
+- Instruction-following categories if IFEval is the target metric.
+
+Data quality may beat objective cleverness on small ORPO runs.
+
+## Short answer to "token vs vocab"
+
+The thing to look for is **vocab-axis chunked / fused cross entropy**, especially
+under the names **Cut Cross Entropy** or **Liger Fused Linear Cross Entropy**.
+Token chunking slices `[tokens, vocab]` by rows. Vocab chunking slices it by
+columns and combines chunks with online logsumexp so the full logits matrix is
+never live.