# How Needle 2 Implements Its Confidence Head: Architecture and Code

> Discover Needle 2's confidence head architecture and JAX code implementation. Learn how it predicts per-token confidence scores for efficient inference gating.

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

---

**The confidence head in Needle 2 is a lightweight JAX module that predicts per-token scalar confidence scores using 8 learned probe vectors, implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and designed for inference-time gating without backpropagation.**

Needle 2 uses a **confidence head** to estimate how much the model should trust its own token-level predictions. This mechanism enables downstream filtering of tool calls, speculative execution, and selective computation. This article walks through the complete implementation—from the core module definition to export serialization and practical usage.

---

## Core ConfidenceHead Module

The `ConfidenceHead` class is a Flax `nn.Module` defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). It consumes the hidden states ("cells") output by the transformer stack and produces a single confidence value per token.

### Architecture overview

```python
class ConfidenceHead(nn.Module):
    dtype: jnp.dtype = jnp.bfloat16

    PROBES = 8  # Number of learned "probe" vectors

    @nn.compact
    def __call__(self, cells, keep=None):
        # 1️⃣ Pool per-token hidden states using the learned probes

        pooled = probe_pool(
            cells,
            self.param("probes", default_init(),
                        (self.PROBES, cells.shape[-1])),
            keep,
            self.dtype,
        )
        # 2️⃣ Project the pooled representation to a single logit

        logit = nn.Dense(
            1,
            dtype=self.dtype,
            use_bias=True,
            kernel_init=default_init(),
            name="proj",
        )(pooled)
        # 3️⃣ Return a float32 confidence score per token

        return logit[..., 0].astype(jnp.float32)

```

Source: [architecture.py L63-L76](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L63-L76)

### Key design decisions

- **Probe-based pooling**: Instead of simple mean-pooling, the head learns **8 probe vectors** that attend to hidden states, allowing it to focus on task-relevant dimensions.
- **Single-logit projection**: A `nn.Dense(1)` layer with bias maps the pooled representation to an unnormalized logit.
- **Float32 output**: The final cast to `jnp.float32` ensures numerical stability for downstream thresholds and comparisons, even when the network runs in **bfloat16**.

---

## Integration with SimpleAttentionNetwork

The confidence head is instantiated during model construction and exposed through a dedicated forward method that deliberately **blocks gradients**.

### Setup and forward method

```python

# Instantiated during model construction

self.confidence_head = ConfidenceHead(cfg.jax_dtype)

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)

```

Source: [architecture.py L72-L77](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L72-L77)

### Gradient isolation

The `jax.lax.stop_gradient` wrapper is intentional: the confidence head operates **only during inference** to gate decisions, not to train the main language model. This prevents confidence predictions from interfering with the primary training objective.

---

## Export and Serialization

The confidence head receives explicit treatment in Needle 2's checkpoint format to enable reconstruction across different environments.

### Head code registration

```python
HEAD_CODES = (("contrastive_head", 1), ("confidence_head", 2))
...
ts += _head_tensors(params)  # pulls out probes, proj and bias tensors

```

Source: [export.py L89-L99](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py#L89-L99)

The numeric code **2** identifies the confidence head uniquely, allowing loaders to distinguish it from other auxiliary heads (such as the contrastive head, coded **1**). The `_head_tensors` function extracts:
- Probe parameters (`probes`: 8 × `d_model`)
- Projection kernel (`proj/kernel`: `d_model` × 1)
- Projection bias (`proj/bias`: scalar)

---

## Practical Usage Examples

### Running confidence inference on token batches

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

cfg = TransformerConfig(
    vocab_size=50257,
    d_model=768,
    num_layers=12,
    num_heads=12,
    jax_dtype=jnp.bfloat16,
)

model = SimpleAttentionNetwork(cfg)

# Example token IDs (batch-size=1, seq-len=5)

tokens = jnp.array([[101, 2003, 1037, 2742, 102]])

# Obtain confidence scores (shape: [batch, seq])

conf_scores = model.forward_confidence(tokens)
print(conf_scores)  # e.g., [[0.92, 0.87, 0.45, 0.78, 0.95]]

```

### Filtering tool calls by confidence threshold

```python
from needle.environments._harness import run_tests

# Run tests with a minimum confidence threshold of 0.4

passed = run_tests(my_environment_module, min_confidence=0.4)
print("All tests passed:", passed)

```

The harness inspects each tool response for a `"confidence"` field and rejects any result below the threshold, preventing low-certainty actions from propagating through the agent loop.

### Loading a checkpoint with confidence head parameters

```python
from needle.model.export import load_checkpoint

params = load_checkpoint("model.ckpt")  # reads probes, proj, bias, etc.

model = SimpleAttentionNetwork(cfg).bind(params)

```

---

## Important Implementation Notes

### Fine-tuning limitations

According to [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), the repository emits a warning that **fine-tuning does not update the confidence head**, causing it to be `None` after tuning. This reinforces its role as a frozen inference-time component.

### Computational footprint

- **Parameters**: 8 × `d_model` (probes) + `d_model` + 1 (projection) ≈ **9d + 1** parameters
- For `d_model=768`: ~6,913 parameters—negligible compared to the base transformer
- **Inference cost**: One attention-style pooling operation + dense projection

---

## Summary

- The **confidence head in Needle 2** lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as a `ConfidenceHead` class with **8 learned probes**
- It produces **per-token float32 scores** via probe pooling and single-logit projection
- The `forward_confidence` method in `SimpleAttentionNetwork` **blocks gradients**, reserving the head for inference-only gating
- Checkpoints serialize head parameters with **code 2** via [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)
- The harness in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) demonstrates threshold-based filtering of tool calls
- **Fine-tuning does not train the confidence head**, per [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py)

---

## Frequently Asked Questions

### Why does the confidence head use 8 probes instead of mean pooling?

The **8 learnable probe vectors** enable attention-weighted aggregation of hidden states. Unlike static mean pooling, probes can focus on dimensions most predictive of model certainty, learned end-to-end from the training data. This design appears in probe_pool's attention mechanism within [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py).

### Why is the confidence head detached from gradients with stop_gradient?

The `jax.lax.stop_gradient` call in `forward_confidence` ensures confidence predictions **do not influence the main language model training**. This isolation prevents the auxiliary head from distorting the primary next-token prediction objective, keeping confidence purely as an inference-time diagnostic and gating signal.

### What happens to the confidence head during fine-tuning?

According to [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), **fine-tuning sets the confidence head to None** and issues a warning. The head is frozen or discarded during training updates, so practitioners must reload original head parameters from base checkpoints if confidence scoring is needed post-tuning.

### How can I interpret confidence score magnitudes?

The confidence head outputs **unbounded logits cast to float32**. In practice, scores typically range 0.0–1.0 when passed through a sigmoid in downstream code, though the raw module returns linear logits. The [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) threshold of 0.4 implies sigmoid-activated interpretation, with higher values indicating greater model certainty.