gemma-2-2b-legal-dpo

A 2.61B-parameter Gemma-2 model for grounded legal/financial Q&A, preference-tuned with DPO on AI-generated feedback (RLAIF).

⚠️ Gemma license — read before use

This is a derivative of Google's Gemma-2. It is governed by the Gemma Terms of Use and the Gemma Prohibited Use Policy, which apply to this model and to anything you derive from it. It is not Apache-2.0 or MIT.

Provenance note: the immediate parent, thesreedath/slm-gemma-2b-qa, declares no license of its own. The upstream Gemma terms attach regardless, so this derivative is published under them.

Honest headline

This model performs at parity with the SFT model it was trained from, with a positive lean that is not statistically significant.

win-rate vs thesreedath/slm-gemma-2b-qa	0.530
decisive comparisons	37 wins / 31 losses / 32 ties (n=100)
two-sided binomial p	0.545
95% CI on decisive win-rate	[0.426, 0.662]

The interval straddles 0.5. Judged by Gemini 2.5 Flash on 100 held-out grounded prompts, order randomized to cancel position bias.

The actual finding: RLAIF damage shrinks with scale

Identical pipeline across three models — same prompt pool, same judge, same preference thresholds, same DPO recipe:

model	params	DPO win-rate vs its own SFT
125M	126M	0.455 (degraded)
500M	518M	0.480 (parity)
Gemma-2-2B (this)	2,614M	0.530 (parity, trending positive)

DPO rises monotonically with scale, crossing from harmful to slightly positive. At 125M, distilling toward a stronger teacher pushed the model into fluent prose form it could not fill with substance. At 2.6B that failure mode disappears. No single point is significant at n=100; the trend is.

PPO, by contrast, failed. A PPO run from the same reward model scored 0.33 (24/58/18, p < 0.001) — a real degradation, even after the run itself was made numerically healthy (KL 5.5 vs a target of 6, stable value function). That PPO checkpoint is deliberately not published. Fixing the optimization fixed the optimization, not the model.

Preference data

628 triplets (528 train / 100 eval), mined on-policy from the frozen SFT model: 4 samples per prompt → Gemini scores 1–10 → best/worst kept when the gap is ≥2 → an independent pairwise judge confirms (94.4% agreement) → embedding dedup.

Training

DPO with an RPO-style NLL anchor. beta 0.1, LR 5e-6, 2 epochs, effective batch 16 (micro-batch 1 × accum 16 — Gemma's 256k vocab makes one sequence's fp32 logits ~1GB), 58 steps, 1×H100, 3.8 min, $0.25. Final eval preference accuracy 0.782, margin 1.449.

Prompt format — NOT the standard Gemma template usage

Gemma-2's chat template raises an exception on a system role, so the system instruction is folded into the user turn. The assistant turn ends with <end_of_turn> (id 107), not <eos> — generation must stop on that.

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

m = "abhishekai/gemma-2-2b-legal-dpo"
tok = AutoTokenizer.from_pretrained(m)
model = AutoModelForCausalLM.from_pretrained(m).eval()
model.config.use_cache = True   # the published config ships use_cache=false

system = ("You are a precise legal and financial assistant. Answer using only the "
          "provided context. If the context does not contain the answer, say you "
          "cannot answer from the context.")
question = "What standard of proof applies to the plaintiff?"
context = ("In a civil negligence action the plaintiff must prove duty, breach, "
           "causation, and damages by a preponderance of the evidence.")

# system folded into the user turn — Gemma-2 has no system role
messages = [{"role": "user",
             "content": f"{system}\n\n{question}\n\nContext:\n{context}"}]
text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
ids = tok(text, add_special_tokens=False, return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=160, do_sample=False,
                     eos_token_id=tok.convert_tokens_to_ids("<end_of_turn>"),
                     pad_token_id=tok.convert_tokens_to_ids("<pad>"))
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

Limitations

• Grounded only. The user turn must contain the passage.
• Parity, not an upgrade. See the headline; the SFT parent is an equally reasonable choice.
• LLM-judged preferences. Labels come from Gemini 2.5 Flash, not humans.
• Number fidelity is not guaranteed. Verify any figure it produces.
• Not legal or financial advice.
• Trained and evaluated at a 1,024-token context.

License

Gemma Terms of Use. Use is subject to the Gemma Prohibited Use Policy. Gemma is a trademark of Google LLC.