abhishekai/slm-125m-legal-ppo

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.

Three reasons to read that 0.570 conservatively:

1. 1.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.
2. 2.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.
3. 3.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.

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:

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|>.

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("abhishekai/slm-125m-legal-ppo")
model = AutoModelForCausalLM.from_pretrained("abhishekai/slm-125m-legal-ppo")

msgs = [
    {"role": "system", "content": "You are a precise legal and financial assistant. Answer only from the provided context."},
    {"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?"},
]
text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
out = model.generate(**tok(text, return_tensors="pt", add_special_tokens=False),
                     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.