abhishekai/gemma-2-2b-legal-raft-v2

gemma-2-2b-legal-raft-v2

Gemma-2-2B trained with RAFT (Retrieval-Augmented Fine-Tuning) on top of `abhishekai/gemma-2-2b-legal-sft-v2`, using QLoRA 4-bit.

RAFT trains one specific competence: read three documents, answer from the one that contains the answer — or decline when none of them does. That is the thing this model does better than its parent, and essentially the only thing.

The result depends entirely on which question you ask

On ordinary grounded QA, where the correct passage is always supplied, it is indistinguishable from its parent:

On the RAFT eval — three documents per question, sometimes none of them sufficient — the pooled number is also flat, and the pooled number is misleading:

The overall +0.16 is a large gain on 13% of rows averaged against a small loss on the other 87%.

The unanswerable subgroup was judged in two passes: 40 rows in the first 300-prompt sample, then the remaining 35 separately. The second pass reproduced the first — +2.22 on rows the model had never been scored on, against +2.24 — which is why the pooled interval tightens rather than moves. The sign test, which disagreed at n=40, agrees at n=75: p = 0.023, 41 W / 22 L / 12 T.

The first two rows stay on the balanced 300-prompt sample on purpose. Pooling the top-up into them would make the judged set 22.5% unanswerable against this eval's real 12.5%, and turn the flat overall row into a spurious +0.38 that excludes zero.

The same result without a judge

Counting refusal phrasings directly in the generations, no model-as-judge involved:

It declines more often when it should and less often when it should not — discrimination improving, rather than a refusal threshold sliding in one direction. This matters because a model that simply refuses more would look identical on the unanswerable rows alone.

Two caveats that travel with the +2.23

1. 1.n = 75, and that is every unanswerable row the eval set has. The interval [+1.03, +3.40] is still wide, and it cannot be narrowed further without building more evaluation data. The answerable arm is 3.5x larger and shows no change, so this is the smaller of the two measurements.
2. 2.The refusal wording is templated. 28 of the 41 correct declines are the identical string "The provided documents do not offer a solution to your inquiry.", and 34 are that sentence's family; sft-v2's 23 declines are 23 distinct sentences, each written from the passage in front of it. The decision is right; the phrasing is memorised, and a retrieval system that wants a reason for the refusal will not get one.

The caveat that did not survive: at n=40 the sign test disagreed with the interval (p = 0.229). With the subgroup complete it agrees — p = 0.023 over 63 decisive pairs.

How it was trained

QLoRA 4-bit, LoRA r=16 on all attention and MLP projections (20.8M trainable, 1.28%). One epoch, 2,087 steps, 1×H100, 72 minutes, $4.76.

Training data is 33,398 rows: each question paired with an oracle document plus a partial distractor (same source, does not contain the answer) and an irrelevant one (different source), at a 12.5% unanswerable share — mixed 50/50 with a general-instruction distribution to prevent forgetting, which two earlier narrow fine-tunes of this model did not survive.

Best validation loss 0.9294, at the final step: it did not converge, so this is a floor for the recipe rather than its ceiling. It is the only stage in this project to complete its learning-rate schedule.

Reading the pooled validation loss alone would have made this run look dead — it moved 0.0119 while the grounded half of the split moved 0.064.

Prompt format

Gemma-2's chat template, with the system instruction folded into the user turn (Gemma-2 raises on a system role). The assistant turn ends on `<end_of_turn>` (107), not <eos> — generation must stop on it. Supply the candidate documents in the user turn as context.

How it was evaluated

Claude Sonnet as an independent judge — a different model family from the Gemini 2.5 Flash that generated the training data. Four axes out of 10 (question answering, instruction following, grounding, appropriate refusal), scored against the prompt, the answer, a golden answer and the corpus evidence. Both models were generated on the same prompts with the same seed, and scored in the same run, so the comparison is paired throughout.

Limitations

• —The gain is confined to the decline behaviour. On answerable questions it is, at best, its parent. If your retrieval always returns the right passage, use `gemma-2-2b-legal-sft-v2` instead.
• —Refusal phrasing is near-canned, so it will not read as a natural hedge.
• —~2% of outputs degenerate into repetition on the RAFT eval.
• —It invents citations when it does answer, in a convincing judicial register.
• —Not legal or financial advice. It is a demonstration of method.
• —No labelled benchmark (CaseHOLD, LexGLUE) has been run, so nothing here establishes absolute quality — only how it compares to its own parent.