README.md

---
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
@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](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