# Why Needle Returns `confidence: None` After Fine-Tuning (And How to Restore It)

> Needle confidence returns None after fine-tuning because LoRA only updates the contrastive head. Learn how to restore confidence scores for your fine-tuned models.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-15

---

**Needle deliberately returns `confidence: None` for fine-tuned models because the LoRA fine-tuning process only updates the contrastive head for tool selection, leaving the confidence head untrained and its scores uncalibrated.**

When you load a fine-tuned checkpoint in Needle, the library masks the confidence field to prevent misleading outputs. This design choice protects users from trusting unreliable confidence scores that stem from a frozen module. Understanding this behavior requires examining how the **ConfidenceHead** and **ContrastiveHead** interact during training and inference.

## How Needle's Confidence Scoring Works

The Needle architecture contains two distinct prediction heads in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

- **ContrastiveHead**: Trained during fine-tuning to rank and select tools
- **ConfidenceHead**: Produces a calibrated confidence score (between 0 and 1) for the selected tool

The **ConfidenceHead** pools learned "probe" embeddings and projects them to a single logit:

```python

# needle/model/architecture.py (lines 63-76)

class ConfidenceHead(nn.Module):
    def __init__(self, hidden_dim: int, num_probes: int = 8):
        self.probes = nn.Parameter(torch.randn(num_probes, hidden_dim))
        self.projection = nn.Linear(hidden_dim, 1)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Pool probe activations and output confidence logit

        ...

```

This head is **not** included in the LoRA training configuration, so its parameters remain frozen at their base model values even after fine-tuning completes.

## Why Fine-Tuning Disables Confidence

### The LoRA Training Scope

Needle's fine-tuning command (`needle finetune`) applies **LoRA adapters** to a specific subset of parameters. According to [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), the training loop only targets the contrastive head:

```bash
needle finetune \
    --checkpoint base_model.cact \
    --lora-rank 8 \
    --out my_finetuned.cact \
    --epochs 3

```

This command produces a checkpoint where:
- Tool selection capability is improved through trained LoRA weights
- Confidence estimation remains at base model quality (stale and uncalibrated)

### The Runtime Safeguard

When you load a fine-tuned checkpoint, Needle detects the custom weights and takes protective action. In [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), two critical behaviors occur:

1. **Warning emission** (lines 58-60): The constructor logs that "finetuning does not update the confidence head, so scores are uncalibrated"

2. **Confidence suppression** (lines 15-17): The `complete` method overwrites any confidence value with `None` when custom weights are present

```python

# needle/__init__.py - simplified illustration

class Needle:
    def __init__(self, weights: str | None = None, ...):
        if weights is not None:
            warnings.warn("finetuning does not update the confidence head...")
            self._weights_loaded = True

    def complete(self, prompt: str) -> dict:
        result = self._forward(prompt)
        if self._weights_loaded:
            result["confidence"] = None  # Force None for safety

        return result

```

## Practical Examples

### Example 1: Confidence Returns `None` with Fine-Tuned Weights

```python
from needle import Needle

# Load checkpoint produced by needle finetune

agent = Needle(weights="my_finetuned.cact", tools=[weather_tool, search_tool])

response = agent.complete("What's the weather in Tokyo?")
print(response["tool_selected"])   # → "weather_tool" (fine-tuned selection)

print(response["confidence"])      # → None (deliberately disabled)

```

### Example 2: Obtaining Valid Confidence from Base Model

```python

# Use original base checkpoint without fine-tuning

agent = Needle(weights="base_model.cact", tools=[weather_tool, search_tool])

response = agent.complete("What's the weather in Tokyo?")
print(response["confidence"])      # → 0.87 (calibrated score available)

```

### Example 3: Checking for Confidence Programmatically

```python
def get_confidence_safe(agent: Needle, prompt: str) -> float | None:
    """Handle potentially None confidence values."""
    result = agent.complete(prompt)
    confidence = result.get("confidence")
    
    if confidence is None:
        # Fine-tuned model loaded — confidence unavailable

        return None
    
    # Base model — return calibrated score

    return confidence

```

## Key Files and Their Roles

| File | Purpose | Critical Lines |
|------|---------|--------------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | Defines **ConfidenceHead** with probe pooling logic | 63-76 |
| [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) | Implements confidence suppression and user warning | 15-17, 58-60 |
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | LoRA training loop (excludes confidence head) | ~401 |
| [`tests/test_finetune.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_finetune.py) | Validates confidence `None` behavior after fine-tuning | — |

## Workarounds and Alternatives

If you require confidence scores for a fine-tuned deployment, consider these approaches:

- **Use base model confidence as proxy**: Accept that confidence scores will reflect base model calibration, not fine-tuned behavior
- **Implement custom confidence estimation**: Add a post-hoc confidence module trained on your fine-tuned model's outputs
- **Full fine-tuning** (not LoRA): Modify [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) to include **ConfidenceHead** parameters in the trainable set — note this requires significant compute and risks overfitting

## Summary

- Needle's **ConfidenceHead** remains frozen during LoRA fine-tuning, producing uncalibrated scores
- The library deliberately returns `confidence: None` via logic in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) to prevent misuse
- Valid confidence scores are only available when using base model checkpoints without fine-tuned weights
- The warning message and `None` assignment occur at lines 58-60 and 15-17 of [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) respectively

## Frequently Asked Questions

### Can I enable confidence scoring for my fine-tuned Needle model?

No, not without modifying the source code. The confidence head is explicitly excluded from LoRA training, and the runtime enforces `None` values when custom weights are detected. To obtain confidence scores, you must use the base model checkpoint or implement a custom confidence estimation layer.

### Why does Needle use a separate ConfidenceHead instead of contrastive score magnitude?

The **ConfidenceHead** in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) uses learned probes to model uncertainty in a more nuanced way than raw contrastive logits. However, this sophistication requires proper training data and updates—absent during tool-focused LoRA fine-tuning, the head's outputs become unreliable.

### Will full fine-tuning instead of LoRA restore confidence scores?

Potentially, but this is not the default Needle behavior. You would need to modify [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) to unfreeze **ConfidenceHead** parameters and provide labeled confidence training data. The base implementation intentionally avoids this to reduce compute costs and prevent overfitting on small fine-tuning datasets.

### How can I detect programmatically whether confidence will be `None`?

Check if custom weights were loaded: the `Needle` instance tracks this internally, and the warning emitted at initialization (visible in logs) confirms the condition. There is no public API to query this state directly in current versions.