---
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](./Qwen3-14B/) | 29.1M | **0.939** | 0.106 | 0.790 | 0.839 | 111M |
| [Qwen3-8B](./Qwen3-8B/) | 28.6M | 0.930 | 0.110 | 0.771 | 0.822 | 109M |
| [Qwen3-4B](./Qwen3-4B/) | 25.6M | **0.944** | 0.112 | 0.791 | 0.826 | 98M |
| [Qwen3-30B-A3B](./Qwen3-30B-A3B/) | 33.9M | 0.936 | 0.114 | 0.774 | 0.815 | 129M |
| [Mistral-7B](./Mistral-7B/) | 34.9M | 0.908 | 0.129 | 0.717 | 0.773 | 133M |

## Usage

### Installation

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

### Loading a Checkpoint

```python
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

```python
# 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:

```python
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](https://huggingface.co/datasets/DamonDemon/SpanUQ-Benchmark) — ~293K annotated spans across 20K prompts with continuous soft uncertainty labels derived from 20× sampling + cross-sample verification.

## Citation

```bibtex
@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](https://arxiv.org/abs/2607.05721)
- 💻 **Code**: [github.com/DamonDemon/SpanUQ](https://github.com/DamonDemon/SpanUQ)
- 📊 **Dataset**: [DamonDemon/SpanUQ-Benchmark](https://huggingface.co/datasets/DamonDemon/SpanUQ-Benchmark)

## License

Apache License 2.0