SpanUQ

Span-Level Uncertainty Quantification for LLM Generation

     

Installation • Quick Start • Dataset • Training • Evaluation • Pipeline • Citation

---

Overview

SpanUQ is a lightweight (~25M parameter) probe that estimates uncertainty at the span level — semantically coherent text segments each conveying a single assessable unit of meaning. It distills uncertainty knowledge from expensive multi-sample inference into a single forward pass over LLM hidden states.

Key features:

• DETR-style span decoder that simultaneously detects uncertain spans and estimates their uncertainty
• Mixture of Beta (MoB) distribution for calibrated uncertainty scores
• 10–20× faster than sampling-based methods at inference time
• Generalizes across LLM families — tested on Qwen3 (0.6B–30B) and Mistral-7B

Results

Model	AUROC ↑	MAE ↓	ρ_span ↑	ρ_seq ↑
Qwen3-14B	0.939	0.106	0.790	0.839
Qwen3-8B	0.930	0.110	0.771	0.822
Qwen3-4B	0.944	0.112	0.791	0.826
Qwen3-30B-A3B	0.936	0.114	0.774	0.815
Mistral-7B	0.908	0.129	0.717	0.773

Installation

git clone https://github.com/DamonDemon/SpanUQ.git
cd SpanUQ
pip install -e .

Requirements:

• Python ≥ 3.9
• PyTorch ≥ 2.0
• huggingface_hub (for downloading checkpoints/data)
• transformers (for hidden state extraction)

For GPU training and hidden state extraction:

pip install vllm>=0.4.0

For data construction pipeline (optional):

pip install anthropic>=0.25.0

Quick Start

Step 1: Download Checkpoint & Dataset

# Download SpanUQ probe checkpoint + benchmark dataset for Qwen3-14B
python scripts/download.py --backbone Qwen3-14B --output_dir ./data

# Or download all 5 backbones at once
python scripts/download.py --backbone all --output_dir ./data

This downloads from:

• 🤗 Checkpoints: DamonDemon/SpanUQ (~25M param probe per backbone)
• 🤗 Dataset: DamonDemon/SpanUQ-Benchmark (span-level uncertainty labels)

Available backbones: Qwen3-14B, Qwen3-8B, Qwen3-4B, Qwen3-30B-A3B, Mistral-7B

End-to-End Inference (Fastest Demo)

If you just want to try SpanUQ on a single query without pre-extracting hidden states:

python scripts/inference.py \
    --model Qwen/Qwen3-14B \
    --checkpoint_dir ./data/checkpoints/Qwen3-14B \
    --query "Who invented the telephone?" \
    --device cuda:0

This loads the backbone + probe, generates a response, and outputs uncertainty-annotated spans:

══════════════════════════════════════════════════════════════════════
  Query: Who invented the telephone?
══════════════════════════════════════════════════════════════════════

  Response:
  Alexander Graham Bell is widely credited with inventing the telephone in 1876...

──────────────────────────────────────────────────────────────────────
  Span                                               Uncertainty
──────────────────────────────────────────────────────────────────────
  🟢 Alexander Graham Bell is widely credited...          0.0812
  🟡 He filed his patent on February 14, 1876            0.4231
  🔴 Antonio Meucci had demonstrated a similar...        0.8567
──────────────────────────────────────────────────────────────────────

Step 2: Extract Hidden States (requires GPU)

SpanUQ operates on LLM hidden states. You need to extract them from the target backbone:

python scripts/extract_hidden_states.py \
    --model Qwen/Qwen3-14B \
    --data_dir ./data/dataset/Qwen3-14B \
    --output_dir ./data/hidden_states/Qwen3-14B \
    --layers 22 24 26 \
    --dtype float16

GPU requirement: Qwen3-14B needs ~28GB VRAM in float16. Smaller models (4B, 8B) need less.

Step 3: Evaluate

python scripts/evaluate.py \
    --data_dir ./data/dataset/Qwen3-14B \
    --hidden_states_dir ./data/hidden_states/Qwen3-14B \
    --checkpoint_dir ./data/checkpoints/Qwen3-14B \
    --methods spanuq,pe_span,pe_seq,random \
    --device cuda:0

Expected output for Qwen3-14B:

Method                  auroc         ece         mae     rho_span
spanuq                 0.9388      0.0126      0.1057      0.7896
pe_span                0.7842      0.1523      0.2134      0.5621
pe_seq                 0.6915      0.2041      0.2856      0.4102
random                 0.5012      0.2498      0.3312      0.0023

Alternative: On-the-fly Extraction (no pre-extraction needed)

If you prefer not to pre-extract hidden states, the evaluator can extract them on-the-fly:

python scripts/evaluate.py \
    --data_dir ./data/dataset/Qwen3-14B \
    --model_name Qwen/Qwen3-14B \
    --checkpoint_dir ./data/checkpoints/Qwen3-14B \
    --methods spanuq,pe_span,pe_seq,random \
    --device cuda:0

Note: On-the-fly mode is slower (loads LLM into memory) but requires no disk space for hidden states.

Option B: Build Dataset from Scratch

See Data Construction Pipeline below.

Dataset

The SpanUQ-Bench dataset contains ~293K annotated spans across 20K prompts with continuous soft uncertainty labels.

Model	Train	Dev	Test	Total Spans
Qwen3-14B	17,468	499	1,996	344,288
Qwen3-8B	17,468	499	1,996	344,288
Qwen3-4B	17,468	499	1,996	344,288
Qwen3-30B-A3B	17,468	499	1,996	344,288
Mistral-7B	17,458	500	1,995	269,836

Data Format

Each record contains:

Field	Description
query_id	Unique query identifier
prompt	Input prompt
greedy_response	Greedy-decoded LLM response
domain	Source domain (NQ-Open, TriviaQA, etc.)
spans	List of annotated spans (see below)
u_sequence_mean	Sequence-level mean uncertainty

Each span contains:

Field	Description
text	Span text
start / end	Character offsets in the response
claim	Atomic claim expressed by the span
u_consistency	Soft uncertainty label ∈ [0, 1]
supported / contradicted	Verification vote counts
coverage	Fraction of samples addressing this claim

Preparing Data for Training

After downloading via scripts/download.py, extract hidden states:

python scripts/extract_hidden_states.py \
    --model Qwen/Qwen3-14B \
    --data_dir ./data/dataset/Qwen3-14B \
    --output_dir ./data/hidden_states/Qwen3-14B \
    --layers 22 24 26 \
    --dtype float16

The resulting directory structure:

data/
├── checkpoints/
│   └── Qwen3-14B/
│       ├── best_model.pt          # SpanUQ probe weights
│       ├── model_config.json      # Architecture config
│       ├── training_config.json   # Training hyperparams
│       └── temperature.json       # Calibration temperature
├── dataset/
│   └── Qwen3-14B/
│       ├── train.parquet          # Training spans + labels
│       ├── dev.parquet            # Validation split
│       └── test.parquet           # Test split
└── hidden_states/                 # You extract this (Step 2)
    └── Qwen3-14B/
        ├── q_00000/
        │   ├── hidden_states_layer22.pt
        │   ├── hidden_states_layer24.pt
        │   ├── hidden_states_layer26.pt
        │   ├── token_entropy.pt
        │   └── meta.json          # offset_mapping
        └── ...

Training

Train the SpanUQ model with the default configuration:

python scripts/train.py \
    --config configs/default.json \
    --run_name my_experiment \
    --data_root ./data/Qwen3-14B \
    --device cuda:0

Key training arguments:

Argument	Default	Description
--config	—	Path to config JSON
--run_name	required	Experiment name
--data_root	—	Path to data directory
--n_queries	32	Number of DETR queries
--multi_layer_ids	22 24 26	LLM layers to use
--warmup_epochs	5	Span detection warmup
--total_epochs	50	Total training epochs
--lr	1e-4	Learning rate
--batch_size	32	Batch size
--max_seq_len	2048	Max sequence length

Training uses a two-phase schedule:

• Phase 1 (warmup): Train span detection only (pointer + validity losses)
• Phase 2 (joint): Add uncertainty estimation (Beta NLL + contrastive + UCIR losses)

Post-hoc Calibration

After training, apply temperature scaling on the dev set:

python scripts/calibrate.py \
    --checkpoint_dir ./checkpoints/my_experiment \
    --device cuda:0

This learns a single temperature parameter T that improves ECE while preserving AUROC.

Evaluation

Evaluate SpanUQ and baselines:

python scripts/evaluate.py \
    --data_dir ./data/Qwen3-14B \
    --checkpoint_dir ./checkpoints/my_experiment \
    --methods spanuq,pe_span,pe_seq,semantic_entropy,factscore,lexsim \
    --device cuda:0 \
    --output results.json

Available methods:

Method	Type	Description
spanuq	Probe	Our method (DETR + MoB)
pe_span	Sampling	Predictive entropy per GT span
pe_seq	Sampling	Mean token entropy (sequence-level)
semantic_entropy	Sampling	NLI clustering → entropy
factscore	Verification	Claim verification counts
lexsim	Sampling	Pairwise lexical similarity
verbalized	Prompting	LLM self-assessed confidence

Layer Probing

Analyze which LLM layers are most informative for uncertainty:

python scripts/layer_probe.py \
    --data_root ./data/Qwen3-14B \
    --output_dir ./results/layer_probe \
    --device cuda:0

Data Construction Pipeline

Note: You can skip this entirely by using the pre-built HuggingFace dataset.

The pipeline constructs span-level uncertainty labels in 4 steps:

Step 1: LLM Generation     →  1 greedy + 20 sampled responses per query
Step 2: Claim Decomposition →  Atomic claims with span alignment (via Claude)
Step 3: Claim Verification  →  20× cross-sample verification (via Claude)
Step 4: Soft Labels         →  u = 1 - SUPPORTED / (SUPPORTED + CONTRADICTED)

Prerequisites

• GPU server with vLLM for Step 1
• Anthropic API key for Steps 2–3: export ANTHROPIC_API_KEY=your_key
• Query set in JSONL format (we provide ours on HuggingFace)

Step 1: Generate Responses

python pipeline/step1_generate.py \
    --model Qwen/Qwen3-14B \
    --queries_file data/queries.jsonl \
    --output_dir data/generations \
    --n_samples 20 \
    --tensor_parallel_size 8

Step 2: Claim Decomposition

python pipeline/step2_claim_decomposition.py \
    --input_file data/generations/greedy_responses.jsonl \
    --output_file data/claims.jsonl \
    --model claude-sonnet-4-20250514 \
    --max_workers 50

Step 3: Claim Verification

python pipeline/step3_verify_claims.py \
    --claims_file data/claims.jsonl \
    --samples_dir data/generations \
    --output_file data/verification.jsonl \
    --model claude-sonnet-4-20250514 \
    --max_workers 50

Step 4: Compute Soft Labels

python pipeline/step4_soft_labels.py \
    --claims_file data/claims.jsonl \
    --verification_file data/verification.jsonl \
    --queries_file data/queries.jsonl \
    --output_dir data/labels

See pipeline/README.md for detailed documentation.

Project Structure

SpanUQ/
├── spanuq/                      # Core library
│   ├── config.py                # Model configuration
│   ├── model.py                 # DETR-style span decoder
│   ├── losses.py                # Hungarian matching + MoB + UCIR losses
│   ├── dataset.py               # Data loading and collation
│   ├── span_pool.py             # Differentiable span pooling
│   └── trainer.py               # Two-phase training loop
├── scripts/                     # Entry points
│   ├── download.py              # Download checkpoints & dataset from HF
│   ├── inference.py             # End-to-end inference (query → spans + uncertainty)
│   ├── train.py                 # Training
│   ├── evaluate.py              # Evaluation (SpanUQ + baselines)
│   ├── extract_hidden_states.py # Hidden state extraction
│   ├── calibrate.py             # Temperature scaling
│   └── layer_probe.py           # Layer importance analysis
├── pipeline/                    # Data construction (optional)
│   ├── step1_generate.py        # LLM response generation
│   ├── step2_claim_decomposition.py  # Claim extraction
│   ├── step3_verify_claims.py   # Cross-sample verification
│   └── step4_soft_labels.py     # Soft label computation
└── configs/                     # Per-backbone configurations
    ├── qwen3_14b.json           # Qwen3-14B (best model)
    ├── qwen3_8b.json            # Qwen3-8B
    ├── qwen3_4b.json            # Qwen3-4B
    ├── qwen3_30b_a3b.json       # Qwen3-30B-A3B (MoE)
    ├── mistral_7b.json          # Mistral-7B
    └── default.json             # Alias for qwen3_14b

Resources

• 📄 Paper: arXiv:2607.05721
• 🤗 Model Checkpoints: DamonDemon/SpanUQ
• 🤗 Benchmark Dataset: DamonDemon/SpanUQ-Benchmark

Citation

@article{zhang2026spanuq,
  title={SpanUQ: Span-Level Uncertainty Quantification for Large Language Model Generation},
  author={Zhang, Yimeng and Zhuang, Yingying and Wang, Ziyi and Lu, Yuxuan and Chen, Pei and Gupta, Aman and Su, Zhe and Tan, Ming and Zhang, Zhilin and Liu, Qun and others},
  journal={arXiv preprint arXiv:2607.05721},
  year={2026}
}

License

This project is licensed under the Apache License 2.0 — see LICENSE for details.