Grounded-RAG via RL
Post-trained a small language model to keep factual mental-health QA faithful to its retrieved sources — not just plausible, but traceable, sentence by sentence, back to the passages it was given.
Small LMs used for retrieval-augmented QA drift. They will happily paraphrase a source into something almost-right, or answer from parametric memory when the passage doesn't cover the question. In a mental-health domain, "almost-right" is the failure mode that matters most — a confident, ungrounded claim is worse than a hedge.
The goal wasn't a bigger model. It was a small one that knows what it's allowed to say.
Retrieval is Cohere Embed + Rerank over a curated mental-health corpus. The policy model is Qwen2.5-1.5B. The signal driving post-training is a claim-decomposable groundedness verifier — a Cohere Command-R+ LLM-judge that breaks each generated response into atomic claims and scores whether each claim is supported by the retrieved context.
Before the verifier drove any training, I calibrated it against human judgment on a held-out set. That step gets skipped a lot and it's the difference between optimizing a real signal and optimizing noise.
The same verifier feeds two tracks. Track one: auto-constructed preference pairs used for DPO — a stable floor that reliably lifts groundedness. Track two: the verifier score used as a scalar reward for GRPO, targeting the harder metric of per-sentence citation compliance.
DPO moved groundedness from 0.86 to 0.94 without hurting fluency or answer coverage. That's the stable win.
GRPO is where it got interesting. The first reward formulation lifted the aggregate score but the model was hacking it — producing shorter, hedgier answers that were trivially grounded because they claimed less. Diagnosed by inspecting length and claim-count distributions alongside the reward curve.
Fix: reshape the reward around per-sentence citation compliance rather than aggregate groundedness. Citation coverage moved from 48% to 72% at no measurable cost to groundedness. The lesson isn't that GRPO works — it's that the verifier has to reward the behavior you actually want, and the first version of a reward almost never does.