README.md

metadata

license: apache-2.0
language:
  - en
tags:
  - uncertainty-quantification
  - span-level
  - probe
  - detr
  - nlp
datasets:
  - DamonDemon/SpanUQ-Benchmark
pipeline_tag: text-classification

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

@misc{zhang2026spanuqspanleveluncertaintyquantification,
      title={SpanUQ: Span-Level Uncertainty Quantification for Large Language Model Generation},
      author={Yimeng Zhang and Yingying Zhuang and Ziyi Wang and Yuxuan Lu and Pei Chen and Aman Gupta and Zhe Su and Ming Tan and Zhilin Zhang and Qun Liu and Manikandarajan Ramanathan and Rajashekar Maragoud and Edward Vul and Jing Huang and Dakuo Wang},
      year={2026},
      eprint={2607.05721},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2607.05721},
}

Related Resources

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

License

Apache License 2.0