# How Are Confidence Scores Generated and Used in Needle 2

> Learn how Needle 2 generates confidence scores using its ConfidenceHead module to calibrate logits and gate function execution with configurable thresholds.

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

---

**Needle 2 produces confidence scores through a dedicated ConfidenceHead module that processes transformer hidden representations into calibrated scalar logits, allowing users to gate function execution via configurable thresholds.**

The Needle 2 framework implements a learned confidence estimation mechanism to determine when its transformer-based agent should emit function calls versus abstain. This system relies on a frozen-head architecture that extracts probe features from the model's hidden cells, producing a scalar score that downstream components can threshold to enforce safety constraints.

## Confidence Score Generation in the Architecture

### Extracting Hidden Representations

The generation process begins by capturing the network's internal state. The `hidden_cells` method returns a stack of token embeddings representing the transformer's intermediate activations. These cells feed into `forward_confidence` (defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), lines 72‑77), which first detaches gradients from the backbone to prevent backpropagation through the frozen confidence pathway.

Within this forward pass, the implementation constructs a padding mask to identify valid tokens, then routes the masked cells to the specialized head.

### The ConfidenceHead Module Implementation

The **ConfidenceHead** class (defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), lines 63‑76) handles the actual scoring logic. By default, this module pools **8 probe vectors** from the input cells, aggregating spatial information across the sequence. These pooled features pass through a single dense layer (`nn.Dense(1)`) that outputs a raw logit, which the framework casts to `float32` to produce the final confidence scalar.

This architecture isolates confidence estimation from the main prediction pathway, ensuring that the score reflects the backbone's certainty without being influenced by the primary task gradients.

## Calibration Limitations During Finetuning

A critical constraint affects confidence reliability after model adaptation. According to warnings in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 158‑162), the **confidence head is not updated during finetuning**. Because the head remains frozen while the backbone changes, the calibration between hidden representations and confidence scores breaks.

Consequently, finetuned models set their confidence field to `None`, forcing users to rely on the base pretrained confidence head or retrain the head separately to restore calibrated scores.

## Applying Confidence Scores in Production

### Response Dictionary Integration

During inference, the agent's output dictionary contains a `"confidence"` key populated by `forward_confidence`. This scalar accompanies the model's token predictions and tool call proposals, providing a quantitative measure of certainty for each decision.

### Threshold-Based Gating Mechanism

The primary consumption pattern occurs in the test harness ([`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py), lines 27‑29). Here, the framework compares the confidence score against a user-provided **min_confidence** parameter. If the score falls below this threshold, the harness discards detected function calls, effectively converting the output into a refusal state.

End-users enforce production contracts by invoking environment-specific `run_tests(min_confidence=...)` wrappers (exposed in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py), lines 19‑21), allowing dynamic adjustment of risk tolerance without code changes.

## Implementation Examples

```python

# Compute confidence scores from model hidden states

import needle
from needle.model.architecture import SimpleAttentionNetwork

# Assume `model` is an instantiated SimpleAttentionNetwork 

# and `tokens` is a JAX array of input IDs

conf_logits = model.forward_confidence(tokens)  # → float32 array

print("Raw confidence logits:", conf_logits)

```

```python

# Gate function calls using confidence thresholds in test suites

from needle.environments.smart_home import run_tests

# Execute tests, rejecting any function calls with confidence < 0.4

results = run_tests(min_confidence=0.4, verbose=True)
print("All assertions passed:", results)

```

## Summary

- **ConfidenceHead architecture**: A separate pooling head with 8 probes and a single dense layer generates scalar logits from frozen backbone representations in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).
- **Gradient isolation**: The `forward_confidence` function stops gradients to prevent training-time interference with the confidence estimation pathway.
- **Finetuning limitations**: The confidence head does not update during backbone finetuning (per [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py)), causing scores to become uncalibrated (`None`) after adaptation.
- **Runtime gating**: The test harness in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) filters function calls by comparing scores against `min_confidence` thresholds, treating low-confidence outputs as refusals.

## Frequently Asked Questions

### How does Needle 2 calculate confidence scores?

Needle 2 calculates confidence by passing hidden cell representations through a `ConfidenceHead` module that pools 8 probe vectors and applies a single dense layer (`nn.Dense(1)`), outputting a scalar logit cast to `float32`. This occurs in `forward_confidence` within [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

### Why does confidence become None after finetuning?

The confidence head remains frozen during backbone finetuning, as noted in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 158‑162). Because the head's weights do not adapt to the updated representations, the calibration breaks, and the framework sets confidence values to `None` to prevent misuse of stale scores.

### How do I set a minimum confidence threshold in Needle 2?

Pass the `min_confidence` parameter to environment test runners such as `run_tests(min_confidence=0.4)`. This value is evaluated in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) (lines 27‑29), where scores below the threshold trigger rejection of function calls.

### What happens when confidence falls below the threshold?

When the confidence score is below `min_confidence`, the test harness discards any detected function calls from the model output, effectively treating the response as a refusal. This prevents low-certainty actions from executing in production environments.