# How Needle 2's Confidence-Gating System Works: A Technical Deep Dive

> Explore Needle 2's confidence-gating system. Learn how its confidence head adds scalar scores to transformer outputs for downstream filtering based on thresholds. Dive into the technical details.

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

---

**Needle 2's confidence-gating system adds a dedicated confidence head to the transformer architecture that produces scalar confidence scores for every output, enabling downstream filtering based on user-specified thresholds.**

Needle 2's confidence-gating system is a safety mechanism designed to help developers filter unreliable model outputs. This article explains how the system generates confidence scores from hidden cell representations and how environment test harnesses use these scores to accept or reject assistant responses.

## The Confidence Head Architecture

The confidence head is implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as a lightweight neural network module that pools transformer hidden states and projects them to a single confidence logit.

### ConfidenceHead Module Structure

The `ConfidenceHead` class uses **probe-based pooling**—the same mechanism employed by the contrastive head—to aggregate information across token positions:

```python
class ConfidenceHead(nn.Module):
    dtype: jnp.dtype = jnp.bfloat16
    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)

```

Key implementation details:

- **8 learned probes** aggregate information across the hidden cell dimensions
- **Bfloat16 computation** for efficiency, with float32 output
- **Optional `keep` mask** to ignore padding tokens during pooling
- Single dense projection with bias to produce the final logit

### Forward Confidence Inference

The model exposes confidence computation through `forward_confidence`, which extracts hidden cells with **gradient stopping**—ensuring confidence inference doesn't affect main model training:

```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 returned logit is passed through a **sigmoid function downstream**, yielding a confidence score in `[0, 1]`.

## Confidence Gating in Environment Testing

Once generated, confidence scores integrate with Needle 2's evaluation infrastructure in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py).

### The run_tests Gate

The environment harness implements threshold-based filtering:

```python
def run_tests(module, min_confidence=0.0, verbose=True):
    ...
    # Inside the test loop:

    if got and response.get("confidence", 0.0) < min_confidence:
        # Treat as failure (or ignore) because confidence is below the gate

        if verbose:
            print(f"(confidence gate {min_confidence})")
        got = False
    ...

```

This mechanism allows **per-environment configuration** of acceptance criteria. Each domain-specific environment (`wearable`, `smart_home`, `productivity`) exposes a convenience wrapper:

```python
def run_tests(min_confidence=0.0, verbose=True):
    return _harness.run_tests(sys.modules[__name__], min_confidence, verbose)

```

### Practical Gating Example

Setting `min_confidence=0.4` means any tool or action with model-generated confidence below **0.4** is ignored or marked as failure—enforcing conservative behavior where uncertain predictions are suppressed.

## Important Limitation: Fine-Tuning and Calibration

A critical caveat exists in the current implementation. As noted in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py):

```python
warnings.warn(
    "finetuning does not update the confidence head, so scores are "
    "uncalibrated for tuned weights; this agent reports confidence as None",
    RuntimeWarning)

```

**Fine-tuned models do not update the confidence head**, meaning:

- Confidence scores become **uncalibrated** after supervised fine-tuning
- The system may report `None` rather than potentially misleading scores
- This is an intentional safety measure rather than a bug

## End-to-End Confidence Flow

The complete confidence-gating pipeline operates in four stages:

1. **Token processing** — Input prompts are encoded and processed through the transformer to produce hidden cells
2. **Confidence inference** — `forward_confidence` runs pooled representations through `ConfidenceHead`
3. **Response packaging** — The agent attaches the confidence field (`response["confidence"]`) alongside textual output
4. **Threshold evaluation** — Environment harnesses or custom gating logic compare against `min_confidence` and filter accordingly

## Code Examples

### Obtaining Confidence Scores

```python

# Encode input and run confidence inference

tokens = tokenizer.encode("What is the weather today?")
logits, confidence = model.forward_confidence(tokens)

# confidence is a float in [0, 1]

print(f"Model confidence: {confidence:.2f}")

```

### Applying Confidence Gates in Environment Tests

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

# Only accept high-confidence actions

run_tests(min_confidence=0.5)

# More permissive setting for exploratory testing

run_tests(min_confidence=0.3, verbose=True)

```

## Summary

- **ConfidenceHead** in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) generates scalar confidence scores from pooled hidden cell representations using 8 learned probes
- **Gradient stopping** in `forward_confidence` isolates confidence inference from main model training dynamics
- **Environment harnesses** in [`needle/environments/_harness.py`](https://github.com/cactus-compute/needle/blob/main/needle/environments/_harness.py) implement configurable threshold-based gating via `min_confidence` parameter
- **Fine-tuned models** do not calibrate the confidence head, resulting in `None` scores to prevent unreliable confidence estimates
- The system enables **tunable safety-reliability tradeoffs** across diverse deployment environments

## Frequently Asked Questions

### Why does fine-tuning disable confidence scores?

The confidence head receives no gradient updates during fine-tuning, as implemented in the training loop. The system deliberately reports `None` because uncalibrated confidence scores could provide false assurance about model reliability. This is a conservative safety choice documented in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py).

### Can I use confidence gating outside environment tests?

Yes. The confidence score is available through `model.forward_confidence()` and can be integrated into any downstream application logic. The environment harness pattern in [`_harness.py`](https://github.com/cactus-compute/needle/blob/main/_harness.py) demonstrates one implementation, but the underlying `min_confidence` comparison can be replicated in custom pipelines.

### What confidence threshold should I use?

Threshold selection depends on your application's cost of false positives versus false negatives. The examples in Needle 2's codebase use `0.3` to `0.5` as reasonable starting points. Higher thresholds increase precision at the cost of recall; evaluate on your specific task distribution to optimize.

### How does probe pooling in ConfidenceHead work?

The `probe_pool` function uses 8 learned probe vectors to compute attention-weighted aggregations of hidden cell representations. This mechanism, shared with the contrastive head, allows the model to flexibly attend to different token positions when forming a sequence-level confidence estimate.