# How Needle 2's Simple Attention Network (SAN) Differs from Traditional Transformers

> Explore Needle 2's Simple Attention Network (SAN) differences from traditional transformers. Discover its Hadamard-based FFNs, Engram memory, and auxiliary heads for enhanced retrieval.

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

---

**Needle 2's Simple Attention Network (SAN) retains multi-head self-attention but replaces standard MLP layers with Hadamard-based feed-forward networks, adds persistent Engram memory for learned n-gram retrieval, and introduces auxiliary contrastive and confidence heads for retrieval-augmented generation.**

The `cactus-compute/needle` repository implements a novel language model architecture that modernizes the classic transformer while maintaining its core attention mechanism. Needle 2's Simple Attention Network (SAN) introduces several architectural innovations designed to improve computational efficiency and enable retrieval-augmented capabilities. This analysis examines the specific technical differences between SAN and traditional transformer implementations based on the source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

## Core Architectural Modifications

### Block Composition and Self-Attention

Traditional transformers isolate the `MultiHeadAttention` block from subsequent feed-forward processing. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `Block` class (lines 14-38) composes `MultiHeadAttention` with a `HadamardMLP`, gating mechanisms, and optional Engram memory injection within a single computational unit. This integrated approach differs from the discrete sub-layer structure found in standard transformer implementations.

### HadamardMLP Feed-Forward Networks

Where traditional transformers use a two-layer MLP (`Linear → Activation → Linear`), SAN implements the `HadamardMLP` class (lines 86-103). This component utilizes fast Walsh-Hadamard transforms to drastically reduce the computational cost of feed-forward steps while maintaining expressive capacity. The Hadamard-based approach provides significant FLOP reduction compared to conventional dense layers.

### Zero-Centered RMSNorm

Standard transformers apply LayerNorm after each sub-layer for stabilization. SAN replaces this with `ZCRMSNorm` (lines 46-55), a zero-centered RMS normalization variant specifically designed for low-precision training and inference stability. This modification addresses quantization challenges that standard LayerNorm faces during INT8 or lower precision deployment.

## Memory Augmentation and Context Management

### Engram Memory Tables

Unlike traditional transformers that store all context exclusively in hidden states, SAN introduces **Engram** tables (lines 81-108) that maintain compressed n-gram statistics. During the forward pass, the `_engram_kv` method (lines 109-119) retrieves keys and values from these persistent tables and injects them into the attention stream. This creates a learned retrieval mechanism for long-range patterns without increasing the attention window size.

### KV-Budget Aware Windowing

Rather than using fixed-size attention windows or simple causal masks, SAN implements dynamic KV-budget computation (`kv_budget_window`, lines 98-112). This system automatically adjusts the effective attention range based on model size and hardware memory constraints, shrinking the computational window while preserving output quality through intelligent context selection.

## Auxiliary Heads and Multi-Token Prediction

### Contrastive and Confidence Heads

SAN extends beyond standard output projections with dedicated `ContrastiveHead` (lines 43-61) and `ConfidenceHead` (lines 63-70) classes. These operate on hidden cell states to produce L2-normalized embeddings for retrieval-augmented generation and per-token confidence scores, respectively. Traditional transformers lack these auxiliary output pathways.

### Multi-Token Prediction Pathway

The `SimpleAttentionNetwork` class implements an auxiliary **MTP** (multi-token prediction) pathway within its `__call__` method. After processing through the main block stack, the model concatenates shifted embeddings (`mtp_emb_norm`), combines them via a learned linear projection (`mtp_combine`), processes them through a secondary `Block` (`mtp_block`), and generates additional logits. This architecture improves next-token prediction accuracy and enables more sophisticated decoding strategies than standard single-projection transformers.

## Parameter Sharing and Efficiency

SAN employs strategic parameter sharing not present in vanilla transformers. The multi-token prediction pathway reuses attention parameters across multiple passes through the `mtp_combine` layer, reducing the total parameter count while maintaining model capacity. This differs from traditional architectures where each layer maintains independent parameters.

## Practical Code Examples

The following examples demonstrate how to instantiate and utilize SAN's unique features:

```python

# Building a SAN model

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

cfg = TransformerConfig(
    vocab_size=8192,
    d_model=512,
    num_heads=8,
    num_kv_heads=4,
    num_layers=12,
    max_seq_len=2048,
)
model = SimpleAttentionNetwork(config=cfg)

# Dummy input (batch=1, seq_len=10)

tokens = jnp.arange(10)[None, :]
logits = model(tokens)  # Standard forward pass

```

```python

# Using the contrastive head for retrieval-augmented generation

query = jnp.array([[1, 2, 3, 4]])    # token IDs of a query

tools = jnp.array([[5, 6, 7, 8]])    # token IDs of a tool description

q_emb, t_emb, log_temp = model.forward_contrastive(query, tools)

# q_emb and t_emb are L2-normalized embeddings for cosine similarity comparison

```

```python

# Obtaining per-token confidence scores

tokens = jnp.array([[10, 11, 12, 13, 0]])  # 0 = pad token

confidence = model.forward_confidence(tokens)

# confidence is float32 array (shape: batch × seq_len) with per-token confidence values

```

## Summary

- **HadamardMLP** replaces traditional feed-forward layers with Walsh-Hadamard transforms for reduced computational cost.
- **Engram memory** provides persistent n-gram storage and retrieval, augmenting standard attention with learned long-range patterns.
- **ZCRMSNorm** substitutes LayerNorm for improved low-precision training stability.
- **Contrastive and Confidence heads** enable retrieval-augmented generation and uncertainty quantification absent in standard transformers.
- **MTP pathway** implements multi-token prediction through parameter-shared auxiliary blocks.
- **KV-budget windowing** dynamically adjusts attention scope based on hardware constraints rather than using fixed windows.

## Frequently Asked Questions

### What is the primary architectural difference between Needle 2 SAN and standard transformers?

While both use multi-head self-attention, SAN integrates attention with Hadamard-based feed-forward networks, Engram memory, and gating within unified `Block` classes. Traditional transformers maintain strict separation between attention and feed-forward sub-layers without persistent external memory.

### How does HadamardMLP improve efficiency over standard transformer MLPs?

The `HadamardMLP` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) utilizes fast Walsh-Hadamard transforms instead of dense matrix multiplications. This reduces the computational complexity of the feed-forward step while maintaining model capacity, resulting in fewer FLOPs per token compared to conventional two-layer MLPs.

### What is the purpose of Engram memory in SAN?

**Engram** tables store compressed n-gram statistics that persist across forward passes. Unlike traditional transformers that rely solely on hidden states for context, SAN retrieves keys and values from these tables via `_engram_kv` to inject learned patterns into the attention mechanism, improving long-range dependency modeling without quadratic attention costs.

### How does the Multi-Token Prediction (MTP) pathway work?

The MTP pathway processes shifted embeddings through a secondary `Block` (`mtp_block`) after the main transformer stack, generating additional logits via parameter-shared layers. This allows SAN to predict multiple future tokens simultaneously, improving training efficiency and decoding accuracy compared to standard single-token prediction transformers.