# How Contrastive Learning and Tool Retrieval Heads Work in Needle 2: A Deep Dive into the Architecture

> Discover how Needle 2 leverages contrastive learning and tool retrieval heads within its shared transformer encoder for enhanced language understanding and efficient tool selection.

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

---

**Needle 2 uses a shared transformer encoder with two specialized heads—`ContrastiveHead` for scoring query-tool similarity and `ConfidenceHead` for deciding whether a tool call is needed—to enable unified language understanding and tool retrieval.**

The `cactus-compute/needle` repository implements a novel single-model architecture that eliminates the need for separate retrieval and generation systems. By combining **contrastive learning** with **tool retrieval heads**, Needle 2 learns to embed queries and tools into a shared semantic space while simultaneously producing natural language responses. This article breaks down the implementation details found in the source code.

## The Dual-Head Architecture

Needle 2's core innovation lies in its dual-head design atop a shared transformer encoder. Both heads consume the same hidden states but serve distinct purposes during inference.

### ContrastiveHead: Learning Query-Tool Alignment

The `ContrastiveHead` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 443-560) projects encoder outputs into a contrastive embedding space. Key implementation details:

- **Input projection**: Maps `d_model` dimensions down to `cfg.contrastive_dim` (default 128)
- **Temperature scaling**: Uses a learned `log_temp` parameter to scale similarity scores
- **Forward method**: `forward_contrastive(q_emb, t_emb)` computes dot-product similarity

```python

# Conceptual flow based on architecture.py implementation

class ContrastiveHead(nn.Module):
    def __init__(self, d_model: int, contrastive_dim: int = 128):
        super().__init__()
        self.query_proj = nn.Linear(d_model, contrastive_dim)
        self.tool_proj = nn.Linear(d_model, contrastive_dim)
        self.log_temp = nn.Parameter(torch.zeros(1))  # learned temperature

    
    def forward_contrastive(self, q_emb: Tensor, t_emb: Tensor) -> Tensor:
        # q_emb: [batch, contrastive_dim]

        # t_emb: [num_tools, contrastive_dim]

        q_norm = F.normalize(self.query_proj(q_emb), dim=-1)
        t_norm = F.normalize(self.tool_proj(t_emb), dim=-1)
        similarities = torch.matmul(q_norm, t_norm.T)  # [batch, num_tools]

        return similarities * torch.exp(self.log_temp)

```

The temperature parameter is critical—lower temperatures sharpen the probability distribution, making the model more decisive in tool selection during inference.

### ConfidenceHead: The Tool-Call Gatekeeper

Declared in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) (lines 286-291), the `ConfidenceHead` determines whether the current query requires tool intervention. The `HEAD_CODES` registry maps `"confidence_head"` to its output index, allowing the model to emit a scalar confidence score alongside contrastive logits.

```python

# From export.py - HEAD_CODES mapping

HEAD_CODES = {
    "contrastive_head": 0,
    "confidence_head": 1,
    # ...

}

```

This head outputs a single logit passed through sigmoid to produce a **tool-call probability**. If confidence falls below threshold, Needle 2 generates a direct answer; otherwise, it emits a structured `<tool_call>` block.

## Contrastive Learning Training Pipeline

Needle 2's training objective combines next-token prediction with contrastive loss. The workflow operates as follows:

1. **Tokenization with tool annotations**: The tokenizer in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) preserves special tokens (`<tools>`, `</tools>`, `<tool_call>`) that demarcate tool boundaries

2. **Shared encoding**: The transformer processes interleaved query and tool definition sequences

3. **Contrastive loss computation**: Positive query-tool pairs are pulled together while negatives are pushed apart using InfoNCE:

```python

# Training objective structure (inferred from architecture)

def contrastive_loss(query_emb, tool_emb, positive_idx):
    """
    query_emb: [batch, contrastive_dim] - projected query representations
    tool_emb: [num_tools, contrastive_dim] - projected tool representations
    positive_idx: [batch] - index of correct tool for each query
    """
    logits = contrastive_head.forward_contrastive(query_emb, tool_emb)
    # InfoNCE: maximize log probability of positive pair

    return F.cross_entropy(logits, positive_idx)

```

4. **Joint optimization**: The contrastive loss backpropagates through both projection layers while the language modeling loss updates the shared backbone

## End-to-End Tool Retrieval Flow

The complete inference pipeline in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) orchestrates both heads:

```python
from needle import Needle, tool

@tool
def get_weather(city: str) -> str:
    """Return weather information for a city."""
    return f"Weather in {city}: 72°F and sunny."

# Initialize with tool schema

agent = Needle(tools=[get_weather])

response = agent.complete(
    tools_json='[{"name":"get_weather","parameters":{"city":{"type":"string"}}}]',
    query="What's the weather in Tokyo?"
)

# Possible outputs:

# Case 1 (high confidence): "<tool_call>{\"name\":\"get_weather\",\"arguments\":{\"city\":\"Tokyo\"}}</tool_call>"

# Case 2 (low confidence): "I don't have access to current weather data."

```

Under the hood, `agent.complete()` executes:

- **Single forward pass**: Encoder produces hidden states for query + tool schemas
- **Head computation**: Both `ContrastiveHead` and `ConfidenceHead` process encoder outputs
- **Decision logic**: Confidence score gates whether to emit tool call or natural language

## Key Implementation Files

| File | Responsibility | Critical Components |
|------|---------------|---------------------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | Head definitions and forward passes | `ContrastiveHead`, `forward_contrastive`, `log_temp` |
| [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) | Head registration and export formats | `HEAD_CODES` mapping, `"confidence_head"` index |
| [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) | Special token handling | `<tools>`, `<tool_call>` tokens |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Inference orchestration | Tool-call decision logic, response generation |
| [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) | Tool definition and schema generation | `@tool` decorator, JSON schema export |

## Performance Characteristics

The unified architecture provides significant efficiency advantages:

- **Single forward pass**: Both answer generation and tool selection share encoder computation
- **Scalable retrieval**: Tool embeddings are pre-computed and cached; only query projection runs at inference
- **Differentiable selection**: Contrastive scores enable gradient-based optimization of tool retrieval

For deployment with large tool sets (100+ tools), the codebase supports **batch scoring** where multiple tool candidates are evaluated in parallel matrix operations.

## Summary

- **ContrastiveHead** in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) learns a shared embedding space for queries and tools using projected representations and temperature-scaled similarity
- **ConfidenceHead** registered in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) gates tool invocation via a learned scalar confidence score
- Both heads share a single transformer encoder, enabling efficient joint inference for language generation and tool retrieval
- Training employs **InfoNCE loss** to align semantically related query-tool pairs while the language modeling objective maintains generation quality

## Frequently Asked Questions

### How does Needle 2 decide between answering directly and calling a tool?

The **ConfidenceHead** outputs a probability that the query can be answered without external tools. When this probability exceeds a configurable threshold (typically 0.5), Needle 2 generates a `<tool_call>` block; otherwise it produces natural language. This decision happens in a single forward pass alongside tool scoring.

### What embedding dimension does the contrastive head use?

By default, `cfg.contrastive_dim` is set to **128**， significantly smaller than typical hidden dimensions (768-4096). This compression forces the model to learn compact, semantic representations of tool functionality rather than memorizing surface-form patterns.

### Can Needle 2 handle hundreds of tools simultaneously?

Yes. The architecture supports **batched contrastive scoring** where all tool embeddings are computed in parallel matrix operations. The computational cost scales linearly with tool count for the final similarity computation, while the expensive encoder pass remains constant regardless of tool set size.

### How is the temperature parameter `log_temp` initialized?

The parameter is initialized to **zeros** and learned during training. This starting point yields an initial temperature of 1.0 (since exp(0) = 1), with the model learning to sharpen or soften its confidence distribution based on validation performance.