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

> Discover how Needle 2's Simple Attention Network enhances Transformer efficiency with ZCRMSNorm, Hadamard layers, Engram memory, and Multi-Token Prediction.

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

---

**Needle 2's Simple Attention Network replaces standard Transformer components with ZCRMSNorm normalization, Hadamard-based feed-forward layers, an external Engram memory system, and Multi-Token Prediction decoding to improve efficiency and context handling.**

The Simple Attention Network (SAN) in the cactus-compute/needle repository retains the core self-attention mechanism of traditional Transformers but reimplements nearly every surrounding component for improved parameter efficiency and inference speed. While vanilla Transformers rely on LayerNorm and dense feed-forward networks, SAN introduces Walsh-Hadamard transforms, learned external memory, and auxiliary task heads directly into its forward pass.

## Normalization: ZCRMSNorm vs. LayerNorm

Standard Transformers apply LayerNorm (mean-variance normalization) between sublayers. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `Block` class instead uses **ZCRMSNorm**, a RMS-norm variant with learned scaling.

The `ZCRMSNorm` class normalizes activations by their root-mean-square rather than mean and variance, then applies a trainable gain parameter. This appears twice in each `Block`: once before the attention mechanism and once after, replacing the traditional post-attention LayerNorm.

## Feed-Forward Networks: HadamardMLP Architecture

Instead of the standard two-dense-layer FFN with GELU activation, SAN implements **`HadamardMLP`** as defined in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py). This module projects inputs into a larger Walsh-Hadamard space, applies element-wise scaling with SiLU activation, and projects back to the model dimension.

This approach reduces parameter count compared to dense layers while maintaining expressiveness. The `Block` class integrates `HadamardMLP` as its final sublayer, gated by a learned residual gate (`self._gate`).

## External Memory: The Engram Module

Unlike standard Transformers that rely solely on the KV cache for context, SAN adds an **`Engram`** external memory system. Located in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py), the `Engram` class creates *engram indices* from the input token stream, retrieves key-value vectors from embedding tables, and fuses these into the attention computation.

This allows the model to access longer-range context without linearly growing the KV cache size, effectively decoupling memory capacity from sequence length during inference.

## Auxiliary Prediction Heads

SAN extends the typical language modeling head with two specialized modules defined in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py):

- **`ContrastiveHead`** – Generates normalized embeddings for retrieval tasks, enabling semantic search capabilities alongside generation.
- **`ConfidenceHead`** – Predicts per-token confidence scores for uncertainty quantification during decoding.

Both heads are optional during the forward pass and can be accessed via `forward_contrastive()` and `forward_confidence()` methods on the main `SimpleAttentionNetwork` class.

## Multi-Token Prediction Decoding

Standard Transformers produce logits directly from the final hidden state. SAN implements an optional **Multi-Token Prediction (MTP)** stage where, after the main forward pass, the model concatenates the final hidden state with shifted embeddings, processes them through `self.mtp_block`, and generates a secondary set of logits.

This two-stage decoding improves next-token prediction accuracy by allowing the model to refine its predictions based on its own initial outputs.

## Built-in Quantization and Flash Attention

The `SimpleAttentionNetwork` includes native support for low-precision inference through functions like `_aq` and `maybe_quant_kv`, enabling on-the-fly fake-quantization of weights and activations. The `MultiHeadAttention` class automatically selects optimized attention kernels, falling back from JAX's `dot_product_attention` to manual implementations when hardware support is unavailable.

## Practical Usage Example

The following code demonstrates how to instantiate and use the Simple Attention Network with its extended features:

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

# Configure the model

cfg = TransformerConfig(
    vocab_size=32768,
    d_model=768,
    num_heads=12,
    num_kv_heads=6,
    num_layers=27,
    max_seq_len=4096,
    engram_layers=(2, 15),
    engram_orders=(2, 3),
    engram_slots=8192,
    flash=True,
    dtype="bfloat16",
)

# Initialize SAN

san = SimpleAttentionNetwork(config=cfg)

# Forward pass

tokens = jnp.arange(16)[None, :]  # (batch=1, seq=16)

logits = san(tokens)             # (1, 16, vocab_size)

# Multi-token prediction

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

# Auxiliary heads

query_emb, pos_emb, log_temp = san.forward_contrastive(
    query_tokens=tokens,
    tool_tokens=tokens,
    quant=False,
)
conf_scores = san.forward_confidence(tokens)

```

## Summary

- **ZCRMSNorm** replaces LayerNorm in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) for stable training dynamics.
- **HadamardMLP** substitutes dense feed-forward layers, using Walsh-Hadamard transforms for parameter efficiency.
- The **Engram** module provides external memory retrieval, extending context beyond the standard KV cache.
- **ContrastiveHead** and **ConfidenceHead** enable retrieval and uncertainty estimation without external models.
- **Multi-Token Prediction** adds a secondary decoding stage for improved accuracy.
- Native **quantization** helpers and automatic **flash attention** selection optimize inference performance.

## Frequently Asked Questions

### What is ZCRMSNorm and why does Needle 2 use it instead of LayerNorm?

ZCRMSNorm is a RMS-normalization layer with learned scaling that appears in the `Block` class of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py). It normalizes by the root-mean-square of activations rather than mean and variance, which reduces computational overhead while maintaining training stability. Needle 2 uses it to simplify the normalization pipeline and improve hardware utilization compared to traditional LayerNorm.

### How does the Engram module improve context handling in SAN?

The `Engram` class in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) creates indices from input tokens and retrieves key-value pairs from learned embedding tables. This external memory allows the model to reference information from earlier in the sequence or from external sources without increasing the quadratic memory cost of the attention mechanism, effectively enabling longer context windows.

### What is Multi-Token Prediction in Needle 2?

Multi-Token Prediction (MTP) is an optional decoding strategy where the model runs a second forward pass through `self.mtp_block` after generating initial logits. By concatenating the main hidden state with shifted embeddings and processing them again, SAN produces refined logits that often yield more accurate predictions than single-pass generation.

### Can Needle 2's Simple Attention Network load standard Transformer checkpoints?

No, the architectural differences—specifically ZCRMSNorm, HadamardMLP, and the Engram module—change the parameter structure and computation graph significantly. Models must be trained from scratch or loaded from Needle 2 specific checkpoints that include these custom components.