# How Needle 2’s Simple Attention Network (SAN) Differs from Conventional Transformers

> Discover how Needle 2's Simple Attention Network SAN enhances Transformers with Engram memory, efficient MLPs, and Multi Token Prediction for superior long-range context.

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

---

**Needle 2’s Simple Attention Network extends the standard Transformer architecture by integrating Engram memory for long-range context retrieval, Hadamard-based MLPs for efficient feed-forward processing, Multi-Token Prediction heads, and learnable gating mechanisms, while retaining the core MultiHeadAttention pattern.**

Needle 2 implements the Simple Attention Network (SAN) as a high-level wrapper around a Transformer-like stack in the cactus-compute/needle repository. While it preserves the familiar MultiHeadAttention blocks that define modern Transformers, SAN introduces several architectural novelties—including external memory systems, specialized normalization, and task-specific heads—that differentiate it from vanilla implementations.

## Core Architectural Differences

### External Memory via Engram

Standard Transformers generate key-value pairs on-the-fly from the current sequence without persistent external storage. In contrast, SAN integrates an **Engram** class that functions as a learned KV cache using hashing and convolutional taps. Located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `Engram` module hashes token n-grams into fixed-size slots and supplies additional KV pairs (`engram_kv`) to every block, enabling long-range recall without linearly increasing sequence length.

### HadamardMLP vs. Standard Feed-Forward Networks

Instead of conventional dense layers, SAN employs **HadamardMLP** within its `Block` structure. This component utilizes Walsh-Hadamard transforms for faster, parameter-efficient processing compared to standard linear projections. The `Block` class—which wraps `MultiHeadAttention`—substitutes the typical `nn.Dense` MLP with this specialized feed-forward mechanism, as implemented in the architecture file.

### Multi-Token Prediction and Task-Specific Heads

Vanilla Transformers output only next-token logits through a language modeling head. SAN adds a **Multi-Token Prediction (MTP) block** (`self.mtp_block`) that concatenates hidden states with next-token embeddings and processes them through an additional attention layer to produce `mtp_logits`. Additionally, SAN includes **ContrastiveHead** and **ConfidenceHead** components for downstream tasks, accessible via methods like `encode_contrastive()` for generating embeddings with temperature scaling.

### Normalization and Gating Mechanisms

While traditional Transformers use LayerNorm, SAN implements **ZCRMSNorm** (Root Mean Square Norm with learned scale) for improved stability with bfloat16 precision. Furthermore, SAN introduces explicit **gating mechanisms** (`self._gate`) that learn scalar gates for both attention outputs and MLP contributions, providing fine-grained control over information flow beyond standard residual connections.

## Implementation Structure in the Codebase

The architecture is defined in [[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `SimpleAttentionNetwork` class orchestrates these components. The `setup()` method initializes the embedding layer, main `Stack`, Engram tables, MTP components, and task heads. The `__call__()` method orchestrates the forward pass: creating causal masks, computing RoPE embeddings, fetching Engram KV pairs, running the stacked blocks, and optionally executing the MTP block.

## Practical Usage Examples

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

# Configure the model

cfg = TransformerConfig(
    vocab_size=8192,
    d_model=768,
    num_heads=12,
    num_kv_heads=6,
    num_layers=27,
    max_seq_len=2048,
    dtype="bfloat16",
)

# Initialize SAN

model = SimpleAttentionNetwork(cfg)

# Prepare input tokens

tokens = jnp.arange(10)[None, :]  # Shape: (1, 10)

# Standard forward pass

logits = model(tokens)  # Shape: (1, 10, vocab_size)

# Forward pass with Multi-Token Prediction

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

```

To extract hidden representations for downstream tasks:

```python

# Retrieve hidden cells from all layers

hidden = model.hidden_cells(tokens)  # Shape: (1, 10, num_layers+1, d_model)

# Generate contrastive embeddings

embeddings, log_temperature = model.encode_contrastive(tokens)

```

## Summary

- SAN wraps a Transformer-like stack but replaces standard MLPs with **HadamardMLP** for computational efficiency.
- The **Engram** module provides external hashed memory, unlike the transient KV caches in vanilla Transformers.
- **ZCRMSNorm** replaces LayerNorm for improved numerical stability at lower precision.
- **Multi-Token Prediction** and auxiliary heads (contrastive, confidence) extend capabilities beyond next-token prediction.
- Explicit **gating mechanisms** control attention and MLP contributions with learned parameters.

## Frequently Asked Questions

### What replaces the standard feed-forward network in Needle 2’s SAN?

SAN uses **HadamardMLP**, which applies Walsh-Hadamard transforms instead of traditional dense matrix multiplications. This implementation resides in the `Block` class within [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and offers faster processing with fewer parameters than standard linear layers.

### How does Engram memory differ from a standard Transformer KV cache?

The **Engram** class implements a persistent, learned memory system that hashes token n-grams into fixed slots and applies convolutional taps for retrieval. Unlike standard KV caches that only store representations from the current forward pass, Engram provides long-range context to every layer without increasing sequence length.

### Can SAN predict multiple tokens simultaneously?

Yes. SAN includes a **Multi-Token Prediction (MTP) block** that processes the hidden state concatenated with next-token embeddings through an additional attention layer. When calling `model(tokens, return_mtp=True)`, the forward pass returns both standard logits and `mtp_logits` for multi-step prediction tasks.

### Where are the contrastive and confidence heads implemented?

These task-specific heads are defined in the `SimpleAttentionNetwork` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The `encode_contrastive()` method accesses the `ContrastiveHead`, while the `ConfidenceHead` can be utilized for token-level confidence scoring during inference.