# Simple Attention Network in Needle 2: Architecture, Components, and Implementation

> Explore the Simple Attention Network SAN architecture in Needle 2. Learn about its efficient on-device language modeling with engram memory, multi-head attention, and auxiliary heads.

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

---

**The Simple Attention Network (SAN) is a compact transformer-style encoder that combines standard multi-head attention with learned engram memory, multi-token prediction, and auxiliary contrastive and confidence heads for efficient on-device language modeling.**

The Simple Attention Network serves as the core neural model in Needle 2, an open-source inference framework designed for lightweight, high-performance language generation. This architecture extends the conventional transformer stack with several innovative components that reduce memory overhead while maintaining expressive power.

## Core Architecture Components

The SAN implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The `SimpleAttentionNetwork` class composes eight key subsystems that work together during the forward pass.

### Embedding Layer and Input Scaling

The model begins with a learned embedding that maps discrete tokens to dense vectors. Per the initialization in `self.embedding` at line 78, embeddings are **scaled by √d_model** to stabilize early training dynamics.

```python
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

cfg = TransformerConfig(
    vocab_size=8192,
    d_model=768,
    num_heads=12,
    num_kv_heads=6,
    num_layers=27,
)

model = SimpleAttentionNetwork(config=cfg)

```

### Multi-Head Attention Stack

The `Stack` class, instantiated at `self.stack` (line 84), processes sequences through repeated attention blocks. This component supports optional **flash-attention** acceleration and **KV-budget handling** for constrained inference scenarios.

### Engram Memory System

A distinctive feature of the SAN is its **engram memory** (`self.engrams`, line 90). These learned tables store compressed n-gram statistics and inject additional key/value pairs into the attention mechanism. This design enables **long-range retrieval without expanding the KV cache**, addressing a critical bottleneck in transformer inference.

### Rotary Positional Embeddings (RoPE)

Positional information comes from pre-computed sinusoidal frequencies generated by `_rope` at line 104. RoPE applies rotation directly to query and key vectors, eliminating the need for absolute position embeddings or separate positional encoding layers.

### Multi-Token Prediction (MTP) Head

When `return_mtp=True`, the model activates its **two-stage prediction mechanism**. The `mtp_block` (line 96) concatenates the main hidden state with a shifted embedding, projects the result, and processes it through an additional attention block. This improves next-token accuracy with minimal computational overhead.

```python
import jax.numpy as jnp

tokens = jnp.array([[1, 5, 23, 0, 0]])  # batch of token IDs

# Standard inference

logits = model(tokens)  # shape: (1, 5, vocab_size)

# With multi-token prediction

logits, mtp_logits = model(tokens, return_mtp=True)

```

### Auxiliary Heads for Downstream Tasks

The SAN includes two specialized heads beyond standard language modeling:

- **Contrastive head** (`self.contrastive_head`, line 86): Projects pooled representations to a low-dimensional space for self-supervised embedding learning.
- **Confidence head** (`self.confidence_head`, line 88): Outputs scalar uncertainty estimates used for tool selection and calibrated decoding.

Access contrastive embeddings via:

```python
embeddings, log_temp = model.encode_contrastive(tokens)

# embeddings.shape → (batch, contrastive_dim)

# log_temp → scalar temperature for NT-Xent loss

```

## Forward Pass Execution Flow

The `__call__` method implements a strictly ordered computation graph:

1. **Embedding and scaling** — tokens → `embedding(tokens) * sqrt(d_model)`
2. **RoPE table generation** — compute cosine/sine frequencies for current sequence length
3. **Engram KV retrieval** — optionally fetch compressed key/value pairs from engram tables
4. **Main stack processing** — `x, _ = self.stack(x, ...)` executes attention layers
5. **Logit projection** — map hidden states back to vocabulary space via the embedding matrix
6. **Optional MTP computation** — second-stage pass when `return_mtp=True` or during initialization

Causal and padding masks (created by `make_causal_mask` and `make_padding_mask`) enforce autoregressive constraints throughout.

## Configuration and Customization

The `TransformerConfig` dataclass controls all architectural hyperparameters. Key fields include:

| Parameter | Purpose |
|-----------|---------|
| `vocab_size` | Token vocabulary dimension |
| `d_model` | Hidden state dimension |
| `num_heads` / `num_kv_heads` | Attention head configuration (supports GQA) |
| `num_layers` | Depth of the transformer stack |
| `engram_layers` | Which layers receive engram memory (e.g., `(2, 15)`) |
| `contrastive_dim` | Output dimension for contrastive embeddings |

## Source Code Organization

The Needle 2 repository organizes SAN-related code across four primary files:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** — `SimpleAttentionNetwork` definition, `TransformerConfig`, attention blocks, engram memory, and all auxiliary heads
- **[`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py)** — High-level `run` helper for model instantiation and inference execution
- **[`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py)** — Fake-quantization implementation for the `quant=True` codepath
- **[`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_inference.py)** — Unit tests validating forward pass correctness and MTP output shapes

## Summary

- The **Simple Attention Network** in Needle 2 implements a modified transformer encoder with specialized components for efficient inference.
- **Engram memory** provides compressed retrieval capabilities without KV cache explosion.
- **Multi-token prediction** adds a lightweight second stage for improved next-token accuracy.
- **RoPE embeddings** handle positional encoding with no additional parameters.
- **Contrastive and confidence heads** enable self-supervised training and uncertainty-aware decoding.
- All components are configurable through `TransformerConfig` and accessible via clean Python APIs in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

## Frequently Asked Questions

### What makes the Simple Attention Network different from a standard transformer?

The SAN extends the standard transformer with three key innovations: **engram memory** for compressed long-range retrieval, **multi-token prediction** for enhanced next-token accuracy, and **integrated auxiliary heads** for contrastive learning and confidence estimation. These additions require minimal parameter overhead while significantly improving inference efficiency and downstream utility.

### How does engram memory work in the SAN architecture?

Engram memory consists of learned tables that store compressed n-gram statistics. During the forward pass, these tables generate additional key/value pairs that augment the standard attention computation. This allows the model to retrieve long-range dependencies without storing full KV activations for all prior tokens, reducing memory consumption during inference.

### When should I use `return_mtp=True` in Needle 2?

Enable MTP when you need improved next-token prediction accuracy and can tolerate slightly higher latency. The MTP head runs a second attention pass over combined hidden states and shifted embeddings, producing an additional logits tensor. This is particularly beneficial during model evaluation or when maximizing generation quality matters more than raw throughput.

### Where is the Simple Attention Network implemented in the Needle 2 codebase?

The complete implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) starting at line 78 with the `SimpleAttentionNetwork` class definition. Related utilities for inference execution live in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py), while quantization support is implemented in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py).