# Why Fine-Tuned Models Produce Uncalibrated Confidence Scores in Needle

> Discover why fine-tuned Needle models have uncalibrated confidence scores. Learn how excluding the ConfidenceHead causes weight mismatch with updated transformer parameters.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-20

---

**Uncalibrated confidence scores in fine-tuned Needle models occur because the `ConfidenceHead` is excluded from the fine-tuning process, leaving its weights mismatched to the updated transformer parameters.**

Needle is an open-source language model framework that separates core transformer layers from task-specific prediction heads. When you fine-tune a Needle model—whether through LoRA adapters or full-model updates—the **confidence head remains frozen**, creating a distribution mismatch that invalidates any confidence logits it produces.

## How Needle's Confidence Head Works

The confidence mechanism lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as a lightweight feed-forward network.

### The ConfidenceHead Architecture

```python

# Located at needle/model/architecture.py#L63-L76

class ConfidenceHead(nn.Module):
    """
    Maps pooled hidden states to a single scalar logit
    representing model certainty about next-token predictions.
    """

```

This `ConfidenceHead` takes pooled hidden states from the final transformer layer and outputs a logit that gets converted to a probability in [0, 1]. During **pre-training**, this head learns to correlate internal representations with prediction accuracy on the pre-training distribution.

## Why Fine-Tuning Breaks Calibration

When you run fine-tuning via [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), the training loop follows this parameter update pattern:

1. **Transformer weights** — updated (or LoRA adapters added)
2. **Contrastive head** — updated if retrieval tasks are included
3. **Confidence head** — **completely excluded** from gradients and optimizer steps

The result is a fundamental mismatch: the transformer now operates on a new data distribution, but the `ConfidenceHead` still applies weights calibrated to the old distribution. As stated in `needle/model/finetune.py#L101-L102`:

> *"Note: Confidence scores are None for tuned weights."*

## How Needle Handles the Mismatch

The Python wrapper explicitly prevents misleading confidence values from reaching users.

### Initialization Warning

At `needle/__init__.py#L60-L62`, loading a fine-tuned checkpoint triggers:

```python
if is_tuned_weights(weights_path):
    warnings.warn(
        "Fine-tuned weights detected. Confidence scores will be disabled."
    )

```

### Response Nullification

After generation completes, the wrapper enforces null confidence at `needle/__init__.py#L23-L25`:

```python
if is_tuned_checkpoint:
    response["confidence"] = None  # Force null to prevent miscalibration

```

## Practical Demonstration

The behavioral difference between base and fine-tuned models:

```python

# Load a base (pre-trained) model — confidence is available

from needle import Needle

base = Needle(weights="model.cact")          # pre-trained weights

resp = base.complete("What is the capital of France?")
print(resp["confidence"])   # → float in [0, 1], e.g., 0.92

# Load a fine-tuned model — confidence is forced to None

tuned = Needle(weights="fine_tuned.cact")   # fine-tuned checkpoint

resp = tuned.complete("What is the capital of France?")
print(resp["confidence"])   # → None (warning issued on initialization)

```

## Solutions for Calibrated Confidence After Fine-Tuning

If your application requires reliable uncertainty estimates post fine-tuning, you have two architectural options:

**Retrain the confidence head** on your fine-tuned data distribution. Add a regression loss term that teaches the head to predict whether the fine-tuned model's tokens are correct. This requires:
- Forward passes through the frozen fine-tuned transformer
- Gradient updates confined to `ConfidenceHead` parameters
- A calibration target (e.g., exact-match accuracy on held-out data)

**Use base model confidence as a proxy** and treat fine-tuned predictions as uncalibrated point estimates. This preserves the pre-training calibration but loses task-specific uncertainty modeling.

## Summary

- The `ConfidenceHead` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) is a frozen component during fine-tuning
- Fine-tuning updates transformer weights without touching confidence parameters, creating distribution mismatch
- [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) actively guards users by nullifying confidence fields and emitting warnings
- Recovering calibration requires explicit post-hoc training of the confidence head or accepting uncalibrated outputs

## Frequently Asked Questions

### Can I force Needle to return confidence scores for fine-tuned models?

No—the library explicitly prevents this. Even if you bypass the `None` assignment in `needle/__init__.py#L23-L25`, the returned logit would be **statistically meaningless** due to the distribution shift. You would need to retrain `ConfidenceHead` on your fine-tuning data first.

### Does LoRA fine-tuning preserve confidence calibration?

No. LoRA adapters in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) update transformer attention layers while the `ConfidenceHead` remains untouched. The calibration degradation is identical to full-model fine-tuning because the root cause—frozen confidence parameters—persists in both cases.

### How can I verify my model has uncalibrated confidence?

Check for the runtime warning emitted at `needle/__init__.py#L60-L62` on initialization. Additionally, any response where `response["confidence"] is None` indicates the safety mechanism has activated for a fine-tuned checkpoint.

### Is the contrastive head also affected by fine-tuning?

No—the contrastive head **is** included in the fine-tuning objective when retrieval tasks are specified. Only the `ConfidenceHead` is deliberately excluded from gradient updates, as confirmed by the implementation in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py).