gemma-2-2b-legal-rm

gemma-2-2b-legal-rm

A 2.61B-parameter scalar reward model for grounded legal/financial Q&A, trained on AI-generated preferences (RLAIF), derived from Gemma-2-2B.

This is not a chat model. It emits a single scalar per sequence and cannot generate text. Load it with AutoModelForSequenceClassification, not AutoModelForCausalLM.

License — read this first

This model is a derivative of Google's Gemma-2-2B. Google's Gemma Terms of Use and Prohibited Use Policy travel with every derivative and apply to your use of this model. It is published under license: gemma and may not be redistributed as Apache-2.0 or as unlicensed.

Note on the absolute-score matrix (2026-08-15)

Every generating checkpoint in this project was re-scored by a Claude Sonnet judge on four axes out of 10 over a shared set of 300 held-out prompts. A reward model has no cell in that table: it emits a scalar preference score rather than an answer, so there is nothing for an answer-quality rubric to grade. Its pairwise accuracy above remains the relevant measure. The full table is in MODEL_INDEX.md in the project repository.

What it does

Given a rendered prompt concatenated with a candidate answer, it returns one number. Higher means "closer to what Gemini 2.5 Flash preferred." Only differences between scores are meaningful — the Bradley-Terry objective it was trained with fixes score gaps, not absolute magnitude, so the raw value is arbitrary and not comparable across reward models. (This one happens to emit scores on the order of ~28 where the 125M's emit ~1; that difference carries no information about quality, and mistaking it for signal is precisely what broke the first PPO run built on it.)

Prompt format — different from the SLM reward models

Gemma-2's chat template raises on a system role, so the system prompt is folded into the user turn, and an assistant turn ends with <end_of_turn>, not <eos>.

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

name = "abhishekai/gemma-2-2b-legal-rm"
tok = AutoTokenizer.from_pretrained(name)
rm = AutoModelForSequenceClassification.from_pretrained(
    name, num_labels=1, torch_dtype=torch.bfloat16).eval()
rm.config.pad_token_id = tok.convert_tokens_to_ids("<pad>")

SYSTEM = "You are a precise legal and financial assistant. Answer only from the provided context."

def score(question, context, answer):
    user = f"{SYSTEM}\n\nContext: {context}\n\nQuestion: {question}"
    prompt = tok.apply_chat_template([{"role": "user", "content": user}],
                                     tokenize=False, add_generation_prompt=True)
    enc = tok(prompt + answer.strip() + "<end_of_turn>", add_special_tokens=False,
              return_tensors="pt", truncation=True, max_length=1024)
    # input_ids/attention_mask only — the tokenizer also emits token_type_ids,
    # which the SequenceClassification forward() rejects.
    with torch.no_grad():
        return rm(input_ids=enc["input_ids"],
                  attention_mask=enc["attention_mask"]).logits.squeeze().item()

Scoring a differently-formatted string puts the model out of distribution and the number becomes meaningless without any error being raised.

Training data

628 preference triplets built on-policy from the frozen Gemma-2-2B SFT model:

1. Sample 4 candidates per prompt (temp 0.9, top-p 0.95)
2. Gemini 2.5 Flash scores each 1–10 on correctness and grounding
3. Keep the best/worst pair only if the gap is ≥ 2
4. An independent pairwise judge re-checks the ordering, A/B randomized
5. Embedding dedup on prompts, then split

Limitations

• The most accurate reward model here produced the worst policy. PPO against this reward model scored 0.330 against its own SFT baseline (24 wins / 58 losses, p = 0.0002) — a real, statistically significant degradation. That checkpoint is deliberately not published. A better-fitting reward model is not a safer optimization target.
• It is a proxy, not a quality oracle, and it inherits Gemini 2.5 Flash's preferences, including any bias toward fluent prose over terse correctness.
• In-domain only — grounded legal/financial QA against a supplied passage.
• Optimizing hard against it will find its failure modes. If you use it, anchor the policy with an explicit KL term to a reference model.