SpanUQ: Span-Level Uncertainty Quantification for LLM Generation

Pre-trained SpanUQ Checkpoints

SpanUQ is a lightweight (25–35M parameter) DETR-style probe that estimates uncertainty at the span level from LLM hidden states in a single forward pass.

Model Description

SpanUQ attaches to a frozen LLM backbone and reads intermediate hidden states to:

1. Detect uncertain spans (contiguous text segments expressing a single verifiable assertion)
2. Estimate calibrated uncertainty scores via Mixture of Beta (MoB) distributions

The probe is trained with Hungarian matching, UCIR (Uncertainty-Calibrated Importance Reweighting), and a two-phase schedule (span detection warmup → joint training).

Available Checkpoints

Backbone	Params	AUROC ↑	MAE ↓	ρ_span ↑	ρ_seq ↑	Size
Qwen3-14B	29.1M	0.939	0.106	0.790	0.839	111M
Qwen3-8B	28.6M	0.930	0.110	0.771	0.822	109M
Qwen3-4B	25.6M	0.944	0.112	0.791	0.826	98M
Qwen3-30B-A3B	33.9M	0.936	0.114	0.774	0.815	129M
Mistral-7B	34.9M	0.908	0.129	0.717	0.773	133M

Usage

Installation

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

Loading a Checkpoint

import torch
import json
from spanuq.model import SpanUQ
from spanuq.config import SpanUQConfig

# Load model config
with open("checkpoints/Qwen3-14B/model_config.json") as f:
    config_dict = json.load(f)

config = SpanUQConfig(**config_dict)
model = SpanUQ(config)

# Load weights
state_dict = torch.load("checkpoints/Qwen3-14B/best_model.pt", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

Inference Pipeline

# 1. Generate response with target LLM
# 2. Extract hidden states from specified layers
# 3. Run SpanUQ probe

# Example: given hidden states tensor [1, seq_len, d_model]
with torch.no_grad():
    outputs = model(hidden_states, attention_mask)
    # outputs.span_scores: [n_detected_spans] uncertainty in [0, 1]
    # outputs.span_boundaries: [n_detected_spans, 2] start/end positions

Temperature Calibration (Optional)

For models with temperature.json, apply post-hoc calibration:

with open("checkpoints/Qwen3-14B/temperature.json") as f:
    T = json.load(f)["T"]

# Apply: calibrated_logit = raw_logit / T

File Structure

Each model directory contains:

checkpoints/
├── Qwen3-14B/
│   ├── best_model.pt          # Model weights
│   ├── model_config.json      # Architecture parameters (required for loading)
│   ├── training_config.json   # Training hyperparameters (for reproducibility)
│   └── temperature.json       # Calibration temperature T
├── Qwen3-8B/
│   └── ...
├── Qwen3-4B/
│   └── ...
├── Qwen3-30B-A3B/
│   └── ...
└── Mistral-7B/
    ├── best_model.pt
    ├── model_config.json
    └── training_config.json   # (no temperature.json)

Architecture Details

Component	Description
Input projection	Multi-layer hidden states → d_proj=512
Encoder	2-layer Transformer encoder
Decoder	3-layer DETR decoder with n_queries learnable queries
Span head	Regression head predicting (center, width)
Scorer	MoB (K=3) Beta distribution head
Enrichment	Gated span-token attention
Seq aggregation	Importance-weighted span → sequence uncertainty

Training Data

Trained on SpanUQ-Benchmark — ~293K annotated spans across 20K prompts with continuous soft uncertainty labels derived from 20× sampling + cross-sample verification.

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}
}

Related Resources

• 📄 Paper: arXiv:2607.05721
• 💻 Code: github.com/DamonDemon/SpanUQ
• 📊 Dataset: DamonDemon/SpanUQ-Benchmark

License

Apache License 2.0