# How Confidence Gating Works in Needle 2: Architecture and Implementation

> Discover how confidence gating in Needle 2 works. Learn how scalar confidence scores reject operations below a probability threshold, ensuring model certainty.

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

---

**Needle 2 uses a dedicated transformer head to generate scalar confidence scores that gate tool-invocation calls, rejecting operations when the model's certainty falls below a probability threshold.**

Needle 2 introduces a critical safety layer called **confidence gating** to prevent hallucinated or low-certainty tool calls. This mechanism evaluates the model's internal state before executing external operations, ensuring that only high-confidence actions proceed. According to the cactus-compute/needle source code, the implementation combines a lightweight neural head with explicit threshold logic in the inference pipeline.

## The ConfidenceHead Architecture

The gating mechanism centers on `ConfidenceHead`, a specialized module defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 463-776). This component pools transformer hidden states and projects them to a single confidence logit.

### Pooling and Projection Layer

The head uses an eight-probe pooling strategy to aggregate information across the sequence before final projection:

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

```

The `probe_pool` operation aggregates hidden states using eight learnable probe vectors, producing a single representation that captures the model's overall certainty about the generated content.

### Gradient Isolation in forward_confidence

The model exposes `forward_confidence` (lines 572-777 in the same file) to compute scores without backpropagating through the main transformer body:

```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 `jax.lax.stop_gradient` wrapper ensures that confidence computation does not alter the main model weights during inference.

## The Gating Workflow

When Needle 2 generates a candidate tool call, the system executes a three-step validation pipeline to determine whether to proceed.

1. **Compute confidence** – The inference pipeline calls `model.forward_confidence(tokens)` on the generated sequence to obtain a raw logit score.

2. **Apply threshold** – The system applies a sigmoid activation to convert the logit to a probability value. If the probability is **below a predefined threshold** (commonly 0.5), the tool call is **gated off**, and the model returns a textual response or clarification request instead.

3. **Fallback behavior** – If the confidence head has not been calibrated or was frozen during fine-tuning, the model sets the confidence value to `None` and issues a runtime warning.

## Handling Uncalibrated Models

When users fine-tune Needle 2 without updating the confidence head weights, the system detects the mismatch and warns about uncalibrated scores. The check resides in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (lines 60-124):

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

```

Additionally, [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) registers the head in `HEAD_CODES` to ensure the confidence module is properly serialized during model export. Test suites in [`tests/test_weights.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_weights.py) validate that responses include the confidence field (e.g., `ENVELOPE = {"type":"call","confidence":0.9,...}`), confirming end-to-end integration.

## Practical Implementation

The following pattern demonstrates loading a model, generating a response, and applying confidence gating logic:

```python
from needle.model.run import load_model, run_inference
import jax.numpy as jnp

# Load trained checkpoint

model = load_model("path/to/checkpoint")

# Generate response and compute confidence

tokens, text = run_inference(model, "What is the capital of France?")
conf_score = model.forward_confidence(tokens)

# Apply gating threshold

prob = jnp.sigmoid(conf_score)
if prob < 0.5:
    print("Confidence low – refusing to call external tool.")
else:
    call_external_tool(...)

```

This implementation ensures that only high-certainty operations execute external APIs, reducing costs and preventing errors from tentative model outputs.

## Summary

- **ConfidenceHead** lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and pools transformer hidden states through eight learnable probes to produce a confidence logit.
- **Gradient isolation** via `jax.lax.stop_gradient` prevents confidence computation from affecting model weights during inference.
- **Threshold-based gating** uses a configurable probability cutoff (typically 0.5) to accept or reject tool calls.
- **Calibration warnings** in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) alert users when fine-tuning leaves the confidence head uncalibrated.
- **Export integration** through `HEAD_CODES` in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) ensures the head persists in serialized checkpoints.

## Frequently Asked Questions

### What is the ConfidenceHead in Needle 2?

The **ConfidenceHead** is a lightweight neural module appended to the transformer architecture that aggregates hidden states using eight probe vectors and projects them to a single scalar confidence logit. It operates independently of the main generation pathway to evaluate the model's certainty about potential tool calls.

### How does Needle 2 determine whether to invoke a tool?

Needle 2 calls `forward_confidence` on the generated token sequence to obtain a logit score, converts it to a probability via sigmoid activation, and compares it against a threshold. If the probability exceeds the threshold (e.g., 0.5), the tool executes; otherwise, the system gates the call and returns a text-only response or requests clarification.

### Why does Needle 2 warn about uncalibrated confidence scores?

The warning triggers when users fine-tune the base model without simultaneously training the confidence head. Because the head's weights remain frozen at their original values while the underlying transformer adapts to new data, the confidence scores no longer accurately reflect the model's true certainty, necessitating the `None` value and runtime alert in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py).

### What happens when confidence gating rejects a tool call?

When the confidence probability falls below the threshold, Needle 2 aborts the external API invocation and instead either returns the generated text response directly or prompts the user for additional information to clarify the request. This prevents hallucinated or uncertain tool executions from consuming resources or producing incorrect results.