# How Needle Calculates Confidence Scores for Model Responses

> Discover how Needle calculates confidence scores for model responses using its ConfidenceHead module. Learn how probe vectors are pooled and projected to determine raw confidence values.

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

---

**Needle calculates confidence scores using a dedicated ConfidenceHead module that pools probe vectors from hidden cell representations and projects them into a scalar logit, returning raw confidence values alongside model outputs.**

The cactus-compute/needle repository implements confidence estimation through a specialized neural head that operates alongside the main contrastive architecture. Unlike standard probability calibration, Needle's approach extracts internal representations from the model's hidden cells to produce a distinct confidence metric. This mechanism allows developers to assess response reliability through a dedicated API endpoint separate from the generation pathway.

## ConfidenceHead Architecture in architecture.py

The core confidence computation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `ConfidenceHead` class processes intermediate model states. This module aggregates information from hidden representations using learned probe vectors before reducing them to a single confidence value.

```python
class ConfidenceHead(nn.Module):
    PROBES = 8

    @nn.compact
    def __call__(self, cells, keep=None):
        pooled = probe_pool(cells,
                            self.param("probes",
                                       default_init(),
                                       (self.PROBES, cells.shape[-1])),
                            keep, self.dtype)
        logit = nn.Dense(1,
                         dtype=self.dtype,
                         use_bias=True,
                         kernel_init=default_init(),
                         name="proj")(pooled)
        return logit[..., 0].astype(jnp.float32)   # ← scalar confidence logit

```

### Probe Pooling with Learned Vectors

The head utilizes **eight configurable probe vectors** (defined as `PROBES = 8`) that aggregate information from hidden states via the `probe_pool` function. These probes interact with the cell dimensions to distill uncertainty-relevant features from the model's internal representations before passing them to the output layer.

### Scalar Logit Projection

After pooling, the features pass through a single-output dense layer (`nn.Dense(1)`) with bias enabled. This projection converts the aggregated representation into a **raw logit** returned as `jnp.float32`. The output represents an uncalibrated confidence indicator where higher values suggest greater model certainty, though the values do not constitute probabilities.

## Forward Pass Execution and Inference

The `forward_confidence` method in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) orchestrates the confidence computation pipeline during inference. This function extracts hidden states, applies appropriate masking, and feeds processed cells into the confidence head.

```python
def forward_confidence(self, tokens, quant=False, window=0, sink=None):
    cells = jax.lax.stop_gradient(
        self.hidden_cells(tokens, quant=quant, window=window, sink=sink))
    keep = (tokens != self.config.pad_token_id).astype(jnp.float32)
    return self.confidence_head(cells, keep=keep)   # ← returns a float

```

### Gradient-Isolated Hidden States

When invoked, `forward_confidence` extracts hidden cells using `self.hidden_cells()` wrapped in `jax.lax.stop_gradient`. This isolation prevents gradient flow back through the main model parameters during confidence estimation, ensuring that confidence computation does not interfere with the base model's training dynamics.

### Padding Token Masking

The method constructs a binary mask `keep` by comparing tokens against `self.config.pad_token_id`, converting the boolean result to `jnp.float32`. This masking ensures that padding tokens do not contaminate the confidence calculation, focusing the probe pooling exclusively on meaningful content tokens.

## Calibration Limitations and Fine-Tuning

A critical constraint emerges when models undergo fine-tuning without accompanying confidence head updates. In this scenario, which the codebase handles in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), the confidence field is explicitly set to `None` to indicate uncalibrated status:

```python

# When a model has been fine‑tuned but the confidence head was not tuned

response["confidence"] = None

```

This sentinel value alerts consumers that the confidence metric requires external calibration procedures or should be ignored entirely until the head undergoes synchronized training with the base model.

## Production Integration and Confidence Gating

The inference pipeline attaches the confidence logit to the JSON envelope returned to callers, typically structured as `{"type": "call", "confidence": 0.9, ...}`. Production deployments can leverage the **min_confidence gate** implemented in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) to filter low-confidence responses before they reach end users.

## Implementation Examples

The following example demonstrates obtaining a raw confidence score from a trained Needle model:

```python

# Example: obtaining a confidence score from a trained Needle model

import needle

model = needle.load("my-model")                     # load a model

tokens = model.tokenizer.encode("What is the weather?")  # tokenise input

logit = model.forward_confidence(tokens)            # ← float logit

print(f"Confidence logit: {logit:.3f}")

```

For production filtering, implement client-side threshold checks against the JSON response:

```python

# Example: applying a confidence threshold in a client

MIN_CONF = 0.4
response = model.generate("Tell me a joke")
if response.get("confidence", 0.0) >= MIN_CONF:
    print("Answer accepted:", response["output"])
else:
    print("Confidence too low – fallback to another model.")

```

## Summary

- **ConfidenceHead** processes hidden cells through 8 learned probes in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)
- The head outputs **raw logits** (not probabilities) via single-layer projection to enable custom calibration
- `forward_confidence` uses gradient-stopped hidden states with automatic padding masks
- Fine-tuned models without updated confidence heads return `None` in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py)
- Production filtering is available via `min_confidence` gates in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py)

## Frequently Asked Questions

### Why does Needle return raw logits instead of probabilities for confidence?

The ConfidenceHead outputs uncalibrated logits that require separate calibration; returning raw values allows developers to apply their own scaling or thresholding logic based on specific deployment requirements. This design separates the concerns of generation and confidence calibration, enabling flexible post-processing strategies.

### How does fine-tuning affect confidence score accuracy?

When a model undergoes fine-tuning without updating the ConfidenceHead parameters, the head becomes miscalibrated relative to the new model distribution. Needle explicitly sets `response["confidence"] = None` in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) to indicate this uncertainty, requiring developers to either ignore the field or implement recalibration techniques.

### Can I filter responses based on confidence thresholds?

Yes, the `min_confidence` parameter in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) enables automatic server-side rejection of low-confidence responses. Alternatively, you can implement client-side filtering using the confidence field in the JSON envelope, treating the raw logit as a comparable confidence indicator.

### What do the 8 probes in ConfidenceHead represent?

The probes are learnable weight vectors initialized in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) that aggregate information across hidden cell dimensions. They function as attention mechanisms that distill uncertainty-relevant features from the model's internal representations before the final projection to a scalar confidence value.