# How Needle 2's Confidence Head Works and What It Scores: A Technical Deep Dive

> Explore Needle 2's confidence head. Understand how this Flax module projects Transformer states into a certainty score for base models, preventing misleading values for tuned weights.

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

---

**Needle 2's confidence head is a lightweight Flax neural network module that projects pooled Transformer hidden states to a single float32 scalar representing the model's self-estimated certainty, returning active scores for base models and None for tuned weights to prevent misleading confidence values.**

Needle 2 introduces a specialized confidence estimation mechanism to assess output reliability in the cactus-compute/needle codebase. The confidence head operates as an auxiliary component alongside the primary Transformer architecture, providing a quantitative certainty measure for generated responses.

## Architecture of the Needle 2 Confidence Head

The confidence head in Needle 2 extends the Transformer architecture as a dedicated Flax module. According to the source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the implementation processes hidden representations through learned probe vectors to produce calibrated uncertainty estimates.

### The ConfidenceHead Module Implementation

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 63-76), the confidence head is implemented as a Flax `nn.Module` that receives the same **pooled hidden cells** utilized by the contrastive head. The architecture processes these representations through **8 learned probe vectors** defined by the constant `PROBES = 8`. A dense projection layer `nn.Dense(1)` maps the pooled representation to a single logit value, which the module returns as a **float32 confidence score**.

### Input Specifications and Output Format

The head consumes the pooled hidden states from the Transformer backbone and applies its probe-based transformation. Unlike the primary prediction heads, this component focuses exclusively on meta-cognitive assessment, outputting a scalar value rather than token distributions or function call probabilities.

## What the Confidence Score Measures

The confidence logit represents the model's **self-estimated certainty** that the generated reply or function-call envelope contains correct information. This scalar value quantifies the model's internal assessment of its own prediction reliability, distinct from the prediction itself.

### Base Models vs. Tuned Weights Behavior

The scoring behavior differs significantly based on the model loading configuration. For **base (un-tuned) models**, the engine populates the confidence field directly from the head's output, providing active certainty estimates. However, for models loaded with **tuned weights**, the confidence head remains **not fine-tuned** during the adaptation process. To avoid propagating misleading or miscalibrated scores, the system deliberately sets the confidence field to `None`, as implemented in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 23-25).

## Implementation Details in the Source Code

Three critical files define the confidence head's behavior across the Needle 2 codebase.

**Architecture Definition**

The core implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `ConfidenceHead` class and `forward_confidence` method produce the scalar confidence logit. This method handles the transformation from pooled representations to the final certainty estimate.

**Response Integration**

In [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 23-25), the high-level API wraps engine responses and conditionally injects the `confidence` field into the JSON envelope. This logic branch determines whether to include the computed score or assign `None` based on whether tuned weights are active.

**Export and Serialization**

The [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) file defines `HEAD_CODES` mapping, assigning **code 2** to the confidence head (lines 289-290). This identifier enables downstream tooling to correctly locate and serialize the head during model export and import operations.

## Accessing Confidence Scores in Practice

Developers can retrieve confidence estimates through the high-level API or by directly invoking the underlying model methods.

### Using the Needle Agent API

When initializing a base agent without tuned weights, the `complete()` method returns a response dictionary containing the confidence score:

```python

# Example: Query the model and read the confidence score

from needle import Needle

# Initialise a base agent (no tuned weights → confidence is computed)

agent = Needle(tools="[]")           # default base weights

response = agent.complete("What is the capital of France?")

# The engine returns a JSON envelope that now includes a confidence value

print("Answer:", response.get("function_calls"))
print("Confidence:", response.get("confidence"))  # e.g. 0.87

```

### Direct Model Invocation

For low-level access, instantiate the `SimpleAttentionNetwork` and call `forward_confidence` directly:

```python

# Example: Directly invoke the confidence head on raw tokens

import jax.numpy as jnp
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

# Build a config (use the same config you would pass to Needle)

cfg = TransformerConfig(...)
model = SimpleAttentionNetwork(cfg)

# Convert a prompt to token IDs (assume `tokenizer` is available)

tokens = jnp.array(tokenizer.encode("Is it raining today?"))[None, :]

# Forward‑pass the tokens and fetch the confidence score

conf_score = model.forward_confidence(tokens)   # → float32 scalar

print("Confidence score:", conf_score)

```

## Summary

- Needle 2's confidence head is a Flax `nn.Module` processing pooled hidden states through 8 learned probes to produce a float32 certainty scalar.
- The head resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and exposes logits via the `forward_confidence` method.
- Confidence scores represent self-estimated correctness certainty for generated outputs.
- Base models return active confidence values in the JSON response envelope, while tuned models return `None` to prevent misleading estimates.
- The export system identifies the confidence head using `HEAD_CODES` value 2 in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py).

## Frequently Asked Questions

### What neural network architecture does the Needle 2 confidence head use?

The confidence head is implemented as a Flax `nn.Module` that processes pooled Transformer hidden states through 8 learned probe vectors (`PROBES = 8`), projecting the result through a dense layer (`nn.Dense(1)`) to generate a single float32 logit representing the confidence score.

### Why does the confidence score return None for tuned models?

The confidence head is intentionally excluded from fine-tuning when adapting models to specific tasks. Since the head's weights remain frozen at base model values while other parameters adapt, its outputs would provide miscalibrated or misleading certainty estimates. The system therefore sets the confidence field to `None` in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) to prevent erroneous reliability assessments.

### How can I access the raw confidence score without using the high-level API?

You can directly invoke the `forward_confidence` method on a `SimpleAttentionNetwork` instance after tokenizing your input. This method, defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), accepts token arrays and returns the scalar confidence estimate without processing through the response wrapper.

### What does the confidence score actually measure in Needle 2?

The confidence logit quantifies the model's **self-estimated certainty** regarding the correctness of its generated response or function call envelope. It serves as a meta-cognitive metric indicating how reliable the model believes its current output to be, distinct from the content of the prediction itself.