Selective RAG with Conformal Abstention

Hallucination detection for retrieval-augmented generation that scores its own confidence and abstains when it can't answer reliably — with a distribution-free, finite-sample guarantee on the precision of what it does answer.

Overview

In high-stakes RAG, a wrong-but-confident answer is worse than no answer. This project builds a hallucination detector on RAGTruth (word-level hallucination annotations over QA, summarization, and data-to-text) and then wraps it in a selective-prediction layer: instead of always emitting a label, the system calibrates a decision threshold so that responses it accepts meet a target precision, and everything below threshold is abstained on. The detector is a stacked ensemble of fine-tuned and pretrained grounding signals; the abstention layer is split conformal prediction with a Hoeffding finite-sample correction, so the precision target holds on held-out data rather than just in-sample.

Everything runs CPU-only with no paid APIs — deliberately reproducible on a laptop.

Architecture

RAGTruth  (context, response, word-level hallucination labels)
    │
    ├─► Baseline — sentence-level NLI  (DeBERTa-v3-base-MNLI)
    │     each response sentence vs. each context sentence → entailment
    │     aggregate → min / mean / median support, frac_grounded
    │
    ├─► Fine-tuned classifiers  (CPU, fp32)
    │     DeBERTa-v3-small (7.5k & 15k train) · DeBERTa-v3-large
    │     (response, context) pair → P(grounded)
    │
    └─► Pretrained fact-checkers
          MiniCheck-DeBERTa-v3-L · MiniCheck-RoBERTa-L · Vectara HHEM-2.1
                        │
                        ▼
     Stacking — per-task K-fold out-of-fold logistic regression
     over all signals + task one-hot  →  P(hallucinated)
                        │
                        ▼
     Split-conformal abstention (Hoeffding δ correction)
     pick τ so eval precision ≥ target  →  ANSWER  |  ABSTAIN

Data flow. Each grounding signal scores the full test split independently and writes a per-record CSV. stack_all.py merges them into one out-of-fold logistic-regression score per task, and conformal_calibrate.py turns that score into an accept/abstain threshold at a chosen precision target.

Key Engineering Decisions

• Selective prediction over a bigger classifier. The objective isn't peak AUROC — it's a calibrated gate that guarantees precision on accepted answers. Split conformal gives that guarantee without distributional assumptions; a raw softmax threshold does not.
• Per-task calibration. Hallucination base rate swings 3.6× across tasks (QA 0.18, Summary 0.23, Data2txt 0.64), so one global threshold under-serves low-rate tasks. A K-fold out-of-fold per-task logistic regression lets each task keep its own operating point without leaking labels into its own predictions.
• Hoeffding correction on the conformal target. On a finite calibration split the empirical precision threshold is optimistic. Subtracting a Hoeffding slack (δ = 0.05) shrinks the effective α so the eval-time precision actually holds — trading a little coverage for a guarantee that survives out of sample.
• Diversity beats scale in the ensemble. Fine-tuning DeBERTa-v3-large did not beat the small variant on this data; the gains came from stacking architecturally diverse signals. Ablations show fine-tuned models contribute +0.042 AUROC and the MiniCheck checkers add +0.022 despite weak standalone AUROC (~0.61–0.67). Plain logistic regression beat gradient-boosted stackers at this feature count.
• CPU-only, fail-fast training. CUDA/MPS are hard-disabled and everything runs in float32; a first-step sanity check verifies finite, non-zero loss and gradients before committing hours of CPU time. logging_nan_inf_filter is off so NaNs surface instead of being silently masked.

Results

Test split: RAGTruth, n = 2697 (three empty-response records dropped). "Baseline" is the sentence-level NLI mean-support signal.

Detection (AUROC, predicting hallucination):

Model	Overall	QA	Summary	Data2txt
Baseline (sentence-NLI, mean support)	0.614	0.605	0.610	0.641
Fine-tuned DeBERTa-v3-small (15k)	0.824	0.781	0.653	0.811
Stacked ensemble (final)	0.8515	0.795	0.754	0.804

The stacked ensemble adds +0.24 AUROC over baseline. Summary is the biggest gainer from diversification (+0.09 vs. the best single model).

Selective prediction — conformal coverage (50/50 calibration/eval split, seed 42, Hoeffding δ = 0.05). Coverage = fraction of responses accepted; precision = fraction of accepted responses that are actually grounded:

α	Target precision	Coverage	Eval precision	Target met
0.05	95%	10.9%	95.2%	✅
0.10	90%	35.8%	90.9%	✅
0.15	85%	49.6%	87.3%	✅
0.20	80%	70.5%	81.5%	✅

At the canonical 90% precision target, the system accepts 35.8% of responses vs. the baseline's 8.4% (at 92.0% precision) — a 4.27× gain in reliable coverage for a ~1.2 pp precision cost. It's also the first scheme here to satisfy the strict 95% target with non-trivial coverage.

The full experiment-by-experiment log — every model, hyperparameter, ablation, and the honest failure modes — is in docs/experiment-log.md.

Honest limitations. AUROC plateaued at ~0.85, short of the ~0.88–0.90 API-judge ceiling on RAGTruth — the CPU/no-API constraint capped model scale, and both DeBERTa-v3-large runs saturated near 0.82. At the 90% target, Data2txt coverage is essentially zero: a 64% base hallucination rate makes a high-precision tail hard to find, though loosening to α = 0.20 recovers useful coverage there.

Tech Stack

• Languages: Python
• ML / DL: PyTorch, Hugging Face Transformers & Datasets, DeBERTa-v3, scikit-learn (logistic-regression stacking, K-fold OOF)
• Methods: natural-language inference, selective prediction, split conformal calibration with Hoeffding correction, per-task calibration
• Tooling: matplotlib, NumPy, pandas

Running Locally

Prerequisites: Python 3.10+, ~8 GB RAM. No GPU required (scripts are CPU/fp32).

# 1. Environment
python -m venv venv && source venv/bin/activate
pip install -r requirements.txt

# 2. Data — clone RAGTruth into data/ragtruth (kept out of git)
git clone https://github.com/ParticleMedia/RAGTruth data/ragtruth

# 3. Explore + generate the baseline sentence-NLI signals
#    (notebooks/01_explore_ragtruth.ipynb writes results/test_scores.jsonl)

# 4. Fine-tune a detector (CPU; ~40 min for the 7.5k small run)
python scripts/train_deberta_cpu.py --run_name v2_full_15k --full

# 5. Evaluate one run: score + ensemble + conformal
./scripts/run_phase_evaluation.sh v2_full_15k

# 6. Full signal-stacking pipeline (fine-tuned + MiniCheck + HHEM → stack → conformal)
./scripts/run_full_pipeline.sh

# 7. Regenerate the figures
python scripts/make_figures.py

No API keys, secrets, or .env are required — every model is downloaded from the Hugging Face Hub and runs locally.

Repo Layout

scripts/
  train_deberta_cpu.py     CPU fine-tune of DeBERTa-v3 on RAGTruth
  score_test_set.py        score a fine-tuned model on the test split
  score_minicheck.py       MiniCheck pretrained grounding signal
  score_hhem.py            Vectara HHEM grounding signal
  ensemble_eval.py         fine-tuned + baseline NLI ensembling
  stack_all.py             per-task K-fold OOF logistic-regression stacker
  conformal_calibrate.py   split-conformal abstention with Hoeffding correction
  make_figures.py          ROC, risk-coverage, and coverage figures
  run_*.sh                 orchestration
notebooks/                 EDA + baseline sentence-NLI scoring
docs/experiment-log.md     full experiment record
results/figures/           committed figures

License

MIT — see LICENSE.