Score SFT (supervised fine-tuning) samples by gradient similarity to a small set of hand-crafted "gold" examples. Useful for triaging large synthetic or semi-curated datasets: which samples actually push the model in the direction you want, and which pull somewhere else?
Inspired by LESS and MOTIVE-style data selection, simplified for small-scale practical use (a few hundred to a few thousand samples).
- Loads a base LLM and wraps it in a LoRA adapter (so we have a small, well-defined set of parameters to take gradients over).
- For each gold sample and each candidate sample:
- Forward + backward pass, computing loss only on the assistant response tokens.
- Collect the gradient over LoRA parameters.
- Project to a low-dimensional vector (Johnson–Lindenstrauss; see below).
- L2-normalize.
- Score each candidate by cosine similarity to the gold gradients:
score_max— best match against any single gold sample.score_mean— match against the averaged gold direction.
- Rank descending and flag the bottom 25% for human review.
The intuition: two samples that would update the model in the same direction during fine-tuning have similar gradients. So gradient cosine similarity is a proxy for "is this sample teaching the same thing as my gold examples?"
pip install "torch>=2.4" transformers peft accelerate bitsandbytes
# Default: bf16 model, ~41 GB VRAM with Mistral-Nemo-12B
python score_samples.py \
--data your_data.jsonl \
--gold gold.jsonl \
--model mistralai/Mistral-Nemo-Base-2407 \
--output scores.csv
# Tighter VRAM (24 GB):
python score_samples.py ... --load-in-4bit --proj-dim 2048
# Server GPUs (B200) — speedup at the cost of VRAM:
python score_samples.py ... --cache-proj
# Bit-reproducible (slower):
python score_samples.py ... --strict-deterministicBoth --data and --gold expect JSONL with chat-format messages:
{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}scores.csv— all samples ranked byscore_max(descending; rank 1 = most aligned with gold). Bottom 25% flaggedREVIEW.scores_flagged.jsonl— just the flagged samples, full content, for inspection.
Taking full-model gradients on a 12B-parameter model would mean working with ~24 GB vectors per sample. LoRA reduces the trainable parameter count by ~1000×, giving us gradients that are still informative about task direction but small enough to handle. We never actually train the LoRA — we just use it as a structured low-rank slice of the gradient.
The gradient we care about reflects "what does generating this answer given this prompt teach the model?" If we included instruction tokens in the loss, we'd be partly measuring "how to phrase the question," which is noise for our purpose.
Even with LoRA (r=16 on q/k/v/o across 40 layers), the gradient vector is ~10M-dimensional. Comparing 500 such vectors pairwise is feasible but storing them isn't, and the noise floor on cosine similarity in such high dimensions is unhelpfully low.
The Johnson–Lindenstrauss lemma says:
For any set of N points in high-dimensional space and any ε > 0, there exists a linear map to a space of dimension
O(log(N) / ε²)that preserves all pairwise distances within a factor of(1 ± ε).
In practice this means we can project our 10M-dim gradients down to a few thousand dimensions and pairwise cosine similarities are preserved up to small distortion. We use 4096 by default, which gives ε ≈ 0.05 for thousands of samples.
The map is just a random matrix with entries ±1/√k (sign-Rademacher). It satisfies the lemma with high probability without any training. We use a fixed seed so the same projection is applied to every sample, which is what makes the cosines comparable.
For memory reasons, we never materialize the full projection matrix. Instead we project each LoRA parameter tensor with its own small block (mathematically equivalent to one giant block-diagonal projection over the concatenated gradient).
If your gold samples cover diverse skills (e.g. refusals + reasoning + style), the mean gradient can wash out into a vector that resembles none of them individually. score_max (best similarity to any single gold) is more robust in that case. Sort by whichever fits your use case (--sort-by).
| Flag | Default | When to change |
|---|---|---|
--proj-dim |
4096 | Lower (1024–2048) if VRAM-constrained; preserves ranking quality well |
--load-in-4bit |
off | Enable for <40 GB VRAM |
--cache-proj |
off | Enable on B200 for ~2× speedup |
--no-grad-checkpoint |
off | Disable checkpointing if you have huge VRAM headroom |
--seed |
42 | Change to get a different random projection (results should rank similarly) |
--strict-deterministic |
off | Enable for bit-exact reproducibility (papers, audits) |
--sort-by |
score_max |
Use score_mean if your gold set is small and homogeneous |
No. A low score means the sample's gradient points in a different direction than your gold samples' gradients. That's a different claim from "this sample is low quality."
A sample can score low because:
- It teaches a skill your gold set doesn't cover (e.g. coding examples scored against creative-writing gold).
- It's stylistically different but factually correct.
- Its response is short or generic, producing a low-magnitude / noisy gradient.
- It addresses an edge case your gold examples don't represent.
- It's genuinely poor quality.
Only the last category is what most people mean by "bad sample." This script can't distinguish between them — it just tells you "this one doesn't pull in the same direction as your gold."
Treat the flag as "worth a human look," not "delete." A reasonable workflow:
- Skim the flagged JSONL.
- Sort flagged samples into: bad (delete), off-topic but fine (keep, maybe move to a different bucket), covers a skill I forgot to put in gold (keep, and add a gold example for that skill).
- Re-run scoring with the expanded gold set.
The ranking is most useful as a search tool over your dataset, not an automated filter.
By default, the script seeds Python, NumPy, and PyTorch RNGs so the LoRA initialization and projection matrices are identical across runs. This brings run-to-run variance to roughly 1e-3 — small enough not to affect rankings.
If you need bit-exact reproducibility, use --strict-deterministic. This forces deterministic CUDA kernels (cuDNN, math-only SDPA, no TF32) at a real performance cost (~10% on matmul, 2-4× slower attention). Almost never necessary for ranking purposes.
Yes, somewhat. The mean gold gradient (score_mean) becomes meaningless if your gold examples span unrelated skills. Two recommendations:
- Use
--sort-by score_max(the default) so each candidate is matched to its closest single gold example. - Cluster your gold set mentally — if you have e.g. 5 coding + 5 reasoning + 5 refusal examples, expect candidates to score well on whichever cluster they resemble. That's a feature, not a bug.
A handful (5–15) per "skill" you want to catch. The variance of the score estimate goes down with more gold, but past ~20 you're mostly just covering more ground rather than improving precision.
The original LESS paper does this — they fine-tune the LoRA briefly on a small subset before computing gradients, on the theory that gradients from a freshly-initialized LoRA reflect the base model's prior more than task structure. This script does not do that by default for simplicity.
If your scores look like noise (very flat distribution, no obvious good-vs-bad signal), try a 100-step warmup on a random subset of your data before scoring. This often improves the gradient signal substantially, especially when scoring against a base (non-instruct) model.
We want to measure "what does generating this response teach the model" — including the instruction in the loss would partly measure "how to phrase questions," which isn't what we're trying to filter on. This matches the masking used during actual SFT.
Any HuggingFace causal LM with a chat template and standard attention projection names (q_proj, k_proj, v_proj, o_proj). Tested on Mistral-Nemo. For other architectures (e.g. GPT-style with c_attn), update target_modules in LORA_CONFIG.
Not directly — the loss formulation here is plain next-token prediction. You'd need to swap in a preference loss to make the gradients meaningful for preference data. The infrastructure (projection, scoring, ranking) would carry over.
Roughly: (num_gold + num_samples) × time_per_forward_backward. On a single 48 GB GPU with Mistral-Nemo-12B, bf16, gradient checkpointing, and streamed projection, expect ~2-4 seconds per sample at typical SFT lengths (1-2K tokens). 500 samples ≈ 20-30 minutes.
--cache-proj cuts this by roughly half on hardware where the projection matrices fit comfortably in VRAM.