gemma-2-2b-raft

Gemma 2B · RAFT — answers from the right document among distractors, and abstains when the answer is absent.

One of 13 models in a controlled study of how far a small language model can be pushed on US legal and financial text. Every version was trained on the same data, evaluated on the same frozen held-out set, and scored by the same blind LLM judge, so the stages are directly comparable. Compare them side by side in the SLM Arena.

Trained and aligned by Harman Sandhu (Vizuara AI Labs).

This is a LoRA adapter, not a standalone model. It must be loaded on top of google/gemma-2-2b-it, which is gated — you need to accept Google's licence on that repo before this adapter can be used. Usage of the adapter is governed by the Gemma Terms of Use.

At a glance

Base model	google/gemma-2-2b-it
Method	QLoRA — 4-bit NF4 base + rank-16 LoRA
Parameters	2.6B (frozen base)
Trainable parameters	20.8M LoRA (0.8%)
Training data	10,000 retrieval-augmented examples (golden passage + distractors)
Schedule	2 epochs · 1,250
Compute	140.0 min on A100-40GB (Modal)
Cost	$2.47 total lineage
Judge score	9.74 / 10

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE = "google/gemma-2-2b-it"      # gated — accept the licence first
ADAPTER = "Ace-2504/gemma-2-2b-raft"

bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type='nf4',
                         bnb_4bit_compute_dtype=torch.bfloat16,
                         bnb_4bit_use_double_quant=True)
base = AutoModelForCausalLM.from_pretrained(
    BASE, quantization_config=bnb, attn_implementation='eager',
    torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(base, ADAPTER).eval()
tok = AutoTokenizer.from_pretrained(BASE)

SYS = ('You are a precise legal and financial assistant. Answer clearly using the '
       'provided context; do not invent facts.')
user = 'Context:\n<passage>\n\nQuestion: <your question>'
text = tok.apply_chat_template([{'role': 'user', 'content': f'{SYS}\n\n{user}'}],
                               tokenize=False, add_generation_prompt=True)
ids = tok(text, return_tensors='pt').to(model.device)
# Gemma-2's hybrid cache misbehaves here — use_cache=False is required.
out = model.generate(**ids, max_new_tokens=160, use_cache=False)
print(tok.decode(out[0, ids['input_ids'].shape[1]:], skip_special_tokens=True))

Evaluation

Scored on 500 decontaminated held-out questions (chunk-level dedup against public legal benchmarks, so no evaluation passage was trained on). Every model in the study answered the identical questions with deterministic greedy decoding, and a Gemini judge blind to the model graded each answer with the gold answer in hand.

Metric	Value
Judge score (0–10 rubric)	9.74
Judged correctness (0–1)	0.975
Groundedness	99.4%
Fabrication ↓	0.8%
Token-F1	0.145
n	500

By source:

US case law	SEC filings	Educational web
9.61	9.81	9.78

The 0–10 figure is a four-dimension rubric (correctness 0–5 + completeness 0–2 + groundedness 0–2 + clarity 0–1). The 0–1 figure is the stricter correctness-only scale used in the experiment reports. Same questions, same answers, same judge — different scale, so the two numbers differ.

Token-F1 is reported for completeness and should not be read as quality: it punishes correct paraphrase heavily, which is exactly why the judge carries the headline.

Architecture

Class	Gemma2ForCausalLM
Layers	26
Hidden size	2,304
Attention	8 heads / 4 KV · head dim 256 · GQA
Feed-forward	GeGLU · inner 9,216
Attention window	sliding 4,096 (alternating) · logit soft-capping
Norm	RMSNorm
Context	8,192 tokens
Vocabulary	256,128
Embeddings	tied input/output

Training

• Initialised from google/gemma-2-2b-it
• Method — QLoRA — 4-bit NF4 base + rank-16 LoRA
• Data — 10,000 retrieval-augmented examples (golden passage + distractors), generated from a legal/financial corpus and gated by an LLM judge for faithfulness
• Schedule — 2 epochs · 1,250, 140.0 min on A100-40GB
• Loss is masked to the answer tokens only, so the model is never trained to reproduce the prompt

Cost — $2.47. The base is Google's and free to build on, so only this project's own work is counted: $2.47.

Limitations and intended use

• Not legal or financial advice. This is a research artefact for studying small-model training, not a professional tool. Do not rely on its output for real decisions.
• Emits its evidence wrapped in ##begin_quote## … ##end_quote## markers, and the sampler mangles these fairly often — strip or parse them before display.
• The judge behind these scores has not been calibrated against human labels, so treat small differences between models cautiously.
• English only; the corpus is US case law, SEC filings and educational web text.

The rest of the family

Size	Base	QA SFT	RAFT	DPO	RLAIF
125M	slm-125m-e4	slm-125m-sft	slm-125m-raft	slm-125m-sft-dpo	slm-125m-sft-rlaif
500M	slm-500m-base	slm-500m-sft	slm-500m-raft	slm-500m-sft-dpo	slm-500m-sft-rlaif
Gemma 2B	gemma-2-2b-it	gemma-2-2b-sft	gemma-2-2b-raft	gemma-2-2b-sft-dpo	gemma-2-2b-sft-rlaif

All under Ace-2504 except the two imported bases.

This model has its own write-up — training details, cost breakdown and live demo — at https://slm-gemma-raft-harman.vercel.app.

Citation

@misc{sandhu2026slm,
  title  = {Small Language Models for Legal and Financial Text: a controlled study of
            pretraining, instruction tuning, retrieval augmentation and alignment},
  author = {Harman Sandhu},
  year   = {2026},
  note   = {https://slm-arena-harman.vercel.app}
}