# How Needle 2 Calculates Confidence Scores for Generated Responses

> Discover how Needle 2 calculates confidence scores for generated responses. Learn about the ConfidenceHead, learned probes, and scalar logit projection.

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

---

**Needle 2 computes confidence scores by passing token-level hidden states through a dedicated `ConfidenceHead` that pools representations with learned probes and projects them to a scalar logit.**

Needle 2, an open-source transformer inference library from **cactus-compute/needle**, attaches a lightweight confidence head to its `SimpleAttentionNetwork` architecture. This head produces a numeric confidence value for every generated token—except when the model has been fine-tuned, in which case the score is disabled for safety. Understanding this mechanism helps developers interpret model certainty and debug generation quality.

## Confidence Score Pipeline Overview

The confidence calculation flows through four distinct stages in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

1. **Hidden state extraction** via `hidden_cells`
2. **Mask construction** for valid tokens
3. **Probe-based pooling** and projection in `ConfidenceHead`
4. **Gradient-stopped forward pass** via `forward_confidence`

## Hidden State Extraction and Masking

### Extracting Token Representations

The `SimpleAttentionNetwork.hidden_cells` method builds a stack of embeddings and returns intermediate hidden states for every token position. These `cells` serve as the foundation for confidence estimation.

```python
cells = self.hidden_cells(tokens, quant=quant, window=window, sink=sink)

```

The tensor contains one vector per token, preserving both semantic and positional information from the model's forward pass.

### Building the Validity Mask

To exclude padding tokens from confidence computation, Needle 2 constructs a boolean mask called `keep`:

```python
keep = (tokens != self.config.pad_token_id).astype(jnp.float32)

```

This mask is passed directly to the confidence head, ensuring probes only attend to real tokens in the sequence.

## ConfidenceHead Architecture

The **confidence head** in Needle 2 combines probe-based pooling with a single-layer projection. Located at lines 68-75 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the implementation follows this pattern:

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

```

### Probe Pooling Mechanism

The head uses **8 learned probes** (defined by `self.PROBES`) that pool across token positions. This multi-probe design captures diverse aspects of uncertainty—similar to ensemble methods but with fixed, learned aggregation weights.

### Final Projection

A dense layer with single-unit output converts the pooled representation to a scalar logit. This logit represents the raw confidence score before any sigmoid or normalization, giving developers flexibility in threshold selection.

## Forward Confidence Wrapper

The `forward_confidence` method (lines 72-76) orchestrates the full pipeline with one critical safeguard: **gradient stopping**.

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

```

The `stop_gradient` call prevents backpropagation through the main network when training or calibrating the confidence head independently. This architectural choice enables modular fine-tuning without destabilizing the base model.

## Fine-Tuned Model Handling and Safety

Needle 2 explicitly **disables confidence scores for fine-tuned weights**. In [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) at lines 104-105, the agent wrapper nullifies the value:

```python
if self._weights:
    response["confidence"] = None

```

### Rationale for Disabling

The confidence head is **not fine-tuned** during standard Needle 2 training runs. Presenting uncalibrated scores on a fine-tuned base model would mislead users about actual generation quality. The library opts for explicit `None` values rather than potentially deceptive numerics.

This safety behavior is also warned about in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) during the fine-tuning entry point.

## Practical Code Examples

### Querying Confidence Scores Directly

For research or debugging, access raw confidence logits via the network API:

```python
import needle as nd

model = nd.SimpleAttentionNetwork(cfg)          # cfg is a TransformerConfig

tokens = nd.tokenizer.encode("What is the capital of France?")

# Get confidence scores (float per token)

conf_scores = model.forward_confidence(tokens)
print(conf_scores)   # → array([...]) of confidence logits

```

### Using the High-Level Agent

The `Needle` agent interface handles confidence automatically:

```python
import needle as nd

agent = nd.Needle(model_path="path/to/needle-2-weights")
result = agent.run("Translate 'hello' to French.")
print(result["response"])        # generated text

print(result["confidence"])      # None (weights present → un-calibrated)

```

Note the `None` return when weights are loaded—this is intentional and documented behavior.

## Key Source Files

| File | Purpose |
|------|---------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | Core implementation: `ConfidenceHead`, `forward_confidence`, and `hidden_cells` |
| [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) | Agent wrapper with safety nullification at line 104 |
| [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) | Head code registry (`confidence_head` = 2) for model serialization |
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | Warns that confidence heads are excluded from fine-tuning |

## Summary

- **Hidden states** drive confidence: `SimpleAttentionNetwork.hidden_cells` extracts token-level representations for downstream processing.
- **Probe pooling** captures uncertainty: 8 learned probes aggregate information across the sequence before final projection.
- **Gradient isolation** enables modular training: `forward_confidence` stops gradients to preserve base model stability.
- **Safety override** prevents misuse: Fine-tuned models return `None` for confidence to avoid uncalibrated scores.

## Frequently Asked Questions

### What does a higher confidence score indicate in Needle 2?

Higher confidence values indicate that the `ConfidenceHead` assigns greater certainty to the model's hidden representation for that token. The score is a learned logit—not a calibrated probability—so developers should establish task-specific thresholds rather than interpret values as literal percentages.

### Why does my Needle 2 agent return `None` for confidence?

The `Needle` agent returns `None` when `self._weights` is present, as implemented in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) lines 104-105. This occurs because the confidence head is not fine-tuned alongside the base model; presenting raw scores would be misleading. Use the base `SimpleAttentionNetwork` directly if you need confidence values on custom weights.

### How many probes does the Needle 2 confidence head use?

The confidence head uses **8 probes** defined by the class constant `PROBES`. These probes are learned parameters initialized with `default_init()` and trained to pool hidden states effectively for uncertainty estimation.

### Can I fine-tune the confidence head separately in Needle 2?

The architecture supports this via the gradient-stopped `forward_confidence` method, but the default [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) pipeline excludes confidence head updates. You would need to implement a custom training loop that allows gradients through `confidence_head` while keeping base parameters frozen or separately scheduled.