slm-125m-legal-ppo
A 125.8M-parameter grounded legal/financial Q&A model, tuned with PPO
against a learned reward model on AI-generated feedback (RLAIF).
Honest headline
This is the only model in this project whose RLAIF stage scored above its SFT
baseline — and the result is soft enough that you should not treat it as an
upgrade.
| |
|---|
| win-rate vs slm-125m-legal-sft | 0.570 |
| tally | 22 wins / 8 losses / 70 ties (n=100) |
| two-sided sign test on the 30 decisive pairs | p = 0.016 |
| 95% CI on the overall win-rate | [0.472, 0.663] — straddles 0.5 |
Three reasons to read that 0.570 conservatively:
- The two tests disagree, and the significant one is the permissive one. The
sign test reaches p=0.016 only by discarding the 70 ties. The tie-inclusive
confidence interval still contains 0.5.
- 70% ties. Mean KL from the SFT reference across training was 0.08–0.13 —
the policy barely moved. The judge usually cannot tell the two models apart,
which is what a correctly-anchored PPO run looks like at this scale.
- The evaluation is partly circular. The reward model distilled Gemini 2.5
Flash's preferences, PPO optimized against that reward model, and Gemini 2.5
Flash then judged the result. A positive number under those conditions partly
measures agreement with the judge's taste. Breaking that needs a different
judge, not more prompts.
Independent absolute score (2026-08-15)
Scored by Claude Sonnet on four axes out of 10 — a different model family from the
Gemini judge behind the win-rate above, and one that had no hand in writing this
project's training data. 300 held-out prompts, the same prompts and the same scale used
for all twelve checkpoints, so this number is comparable across models in a way no
win-rate here is.
| |
|---|
| mean score | 4.08 / 10 |
| 95% CI (bootstrap over per-prompt scores) | [3.76, 4.41] |
| paired step | +0.03 against its own SFT (4.05), CI [-0.12, +0.19] - indistinguishable from the SFT. Note this disagrees with the 0.570 win-rate above; prefer this number, which is paired per prompt at n=300 and judged by a model that did not write the training data. |
The full twelve-checkpoint table, with every stage-to-stage interval, is in
MODEL_INDEX.md in the project repository.
Where it sits in the scaling picture
Same pipeline, same prompt pool, same judge, same PPO code across all three
models. PPO's win-rate falls with scale — the mirror image of DPO's rise:
| model | params | PPO win-rate | p | DPO win-rate |
|---|
| 125M (this) | 126M | 0.570 | 0.016 | 0.455 |
| 500M | 518M | 0.520 | 0.572 | 0.480 |
| Gemma-2-2B | 2,614M | 0.330 | 0.0002 | 0.530 |
Only the two ends are significant, and they point in opposite directions. The
Gemma PPO checkpoint is a real degradation and is deliberately not published.
Training
PPO maximizing reward - kl_coef * KL(policy ‖ SFT), 60 rollout steps, 32
prompts per rollout, kl_coef=0.2 on 1×H100. Realized: mean KL 0.080, final
KL 0.028, value loss ~0.005, no KL-stop abort, $0.17.
Three implementation details were load-bearing — without them the run is not
merely worse, it is invalid:
- Reward whitening.
kl_coef is only meaningful relative to reward
magnitude, and a Bradley-Terry reward model only fixes score differences —
its absolute scale is arbitrary. Scores are normalized per rollout batch.
- Gradient accumulation across the rollout. An earlier version took 32
sequential optimizer steps on one sequence each per rollout, which drove KL to
64–99 and value loss to 2260. It now accumulates and takes one step per epoch.
- An adaptive KL controller with a floor at the configured
kl_coef. At step
1 the policy is the reference and KL ≈ 0, so a naive controller relaxes the
anchor exactly when it must hold.
An earlier checkpoint trained before these fixes scored 0.48; it was overwritten
and is not what is published here.
Reward model
Trained from the same SFT backbone with a fresh scalar head and the pairwise
Bradley-Terry loss, reaching
0.837 pairwise accuracy on held-out preference
pairs. Published separately as
slm-125m-legal-rm.
Prompt format
Identical to the SFT model — use the tokenizer's chat template; the assistant
turn ends on <|eos|>.
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3tok = AutoTokenizer.from_pretrained("abhishekai/slm-125m-legal-ppo")
4model = AutoModelForCausalLM.from_pretrained("abhishekai/slm-125m-legal-ppo")
5
6msgs = [
7 {"role": "system", "content": "You are a precise legal and financial assistant. Answer only from the provided context."},
8 {"role": "user", "content": "Context: In a civil negligence action the plaintiff must prove duty, breach, causation, and damages by a preponderance of the evidence.\n\nQuestion: What standard of proof applies to the plaintiff?"},
9]
10text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
11out = model.generate(**tok(text, return_tensors="pt", add_special_tokens=False),
12 max_new_tokens=128, do_sample=False)
Limitations
- No arithmetic reliability. At 125M the model will state a revenue delta and
a percentage that do not follow from the two figures it just produced.
- Grounded QA only — it answers from a supplied passage and confabulates
without one.
- Not legal or financial advice.
- PPO optimizes a proxy for quality. A higher reward-model score is not
evidence of a more correct answer, and at 125M the gap between those two things
is wide.