[NeurIPS 2026] 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:
- Detect uncertain spans (contiguous text segments expressing a single verifiable assertion)
- 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/damon-demon/SpanUQ
- 📊 Dataset: DamonDemon/SpanUQ-Benchmark
License
Apache License 2.0