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

> Discover Needle 2's Simple Attention Network architecture. It modifies Transformer's self-attention, replaces LayerNorm, swaps FFNs, and adds Engram memory for enhanced performance.

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

---

**Needle 2's Simple Attention Network retains Transformer's core self-attention mechanism but replaces LayerNorm with ZCRMSNorm, swaps dense FFNs for HadamardMLP, adds external Engram memory, and includes auxiliary heads for contrastive learning and confidence scoring.**

The cactus-compute/needle repository implements a novel language model architecture called the **Simple Attention Network (SAN)**. While it preserves the familiar transformer pipeline—token embeddings, multi-head attention, and stacked layers—it introduces significant structural innovations that improve efficiency, context handling, and downstream task performance. This article breaks down every architectural difference with reference to the actual source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

---

## Embedding and Scaling Layer

Both standard transformers and Needle 2 begin with token embeddings followed by dimensional scaling. In `SimpleAttentionNetwork.setup()` (lines 78–85), the model uses `nn.Embed` for token embeddings and sets `self.embed_scale = math.sqrt(cfg.d_model)`, matching the classic √d_model scaling factor.

The divergence begins immediately after this initial layer. Where conventional transformers feed scaled embeddings directly into standard transformer blocks, SAN applies its first **ZCRMSNorm** before any attention computation.

---

## Block Structure: ZCRMSNorm, HadamardMLP, and Gated Residuals

The `Block` class (lines 5–38) reimagines the standard transformer layer. A typical transformer block follows this pattern:

- Multi-Head Self-Attention → LayerNorm → Feed-Forward Network (two dense layers with GELU) → Residual connections

SAN's `Block` replaces this with a streamlined stack:

1. **ZCRMSNorm** pre-attention normalization
2. **MultiHeadAttention** with optional flash attention
3. **ZCRMSNorm** post-attention normalization
4. **HadamardMLP** instead of dense-GELU-dense FFN
5. **Learned gating** on residual paths via `self._gate`

The `ZCRMSNorm` class (lines 46–55) implements **RMS-normalization with a learned scale**—it normalizes by the root-mean-square of activations rather than mean and variance, then applies a trainable gain parameter. This reduces computational overhead while maintaining training stability.

---

## Feed-Forward Network: HadamardMLP vs. Dense Layers

The `HadamardMLP` class (lines 87–103) is one of SAN's most distinctive components. Standard transformers use two `nn.Dense` layers with GELU or SiLU activation between them. SAN's approach:

- Projects input into a larger **Walsh-Hadamard** transform space
- Applies element-wise scaling and SiLU activation
- Projects back to the original dimension

This yields a **parameter-efficient MLP** that operates faster than dense alternatives while maintaining representational capacity. The Hadamard transform's structured nature allows optimized hardware utilization without sacrificing expressiveness.

---

## External Memory: The Engram Module

Standard transformers store key-value pairs only for tokens in the current sequence. SAN's `Engram` class (lines 81–108) introduces **learned external memory**:

- Creates *engram indices* from the token stream
- Retrieves key/value vectors from dedicated embedding tables
- Fuses these into the attention computation

This enables **longer-range context** without linear growth in KV cache size. The `engram_layers`, `engram_orders`, `engram_heads`, and `engram_slots` configuration parameters control where and how external memory integrates into the network.

---

## Auxiliary Prediction Heads

Beyond the standard language modeling head, `SimpleAttentionNetwork` includes two specialized modules defined in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py):

**ContrastiveHead** (lines 43–61)
- Produces **normalized contrastive embeddings** for retrieval tasks
- Enables semantic search and tool retrieval by embedding queries and documents into a shared space

**ConfidenceHead** (lines 63–71)  
- Predicts a **per-token confidence score**
- Useful for uncertainty quantification and selective prediction

These heads allow the same backbone to serve multiple downstream tasks without architectural modification.

---

## Multi-Token Prediction (MTP) Decoding

SAN implements an optional **two-stage decoding** mechanism. When `return_mtp=True` is passed to the forward call (lines 119–138), the model:

1. Computes primary logits through the main network
2. Concatenates final hidden states with **shifted token embeddings**
3. Processes through `self.mtp_block` (an additional transformer block)
4. Produces secondary logits for **multi-token prediction**

This improves next-token accuracy by allowing the model to explicitly condition on its own predicted distribution, similar to speculative decoding approaches but integrated into the architecture itself.

---

## Built-In Efficiency Features

### Quantization Support

Functions `_aq` and `maybe_quant_kv` (lines 22–27) enable **on-the-fly fake quantization** of activations and weights. This prepares models for low-precision inference without external post-training quantization pipelines.

### Adaptive Flash Attention

The `MultiHeadAttention` class automatically selects JAX's `dot_product_attention` implementation, preferring `cudnn` kernels on GPU with fallback to manual scaled-dot-product attention (lines 45–54). This removes the need for manual kernel selection while maximizing hardware utilization.

---

## Practical Usage Example

```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 = SimpleAttentionNetwork(config=cfg)

# Forward pass

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

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

# With 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 for reduced compute and stable training
- **HadamardMLP** substitutes dense FFNs with Walsh-Hadamard transforms
- **Engram memory** provides external KV storage for extended context
- **ContrastiveHead** and **ConfidenceHead** enable retrieval and uncertainty estimation
- **MTP decoding** improves prediction accuracy through two-stage generation
- **Built-in quantization** and **adaptive flash attention** optimize inference efficiency

---

## Frequently Asked Questions

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

ZCRMSNorm is a **RMS-normalization variant with learned scaling** defined in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) lines 46–55. It normalizes activations by their root-mean-square rather than computing mean and variance, reducing floating-point operations. The learned gain parameter provides equivalent expressiveness to LayerNorm's affine parameters. Needle 2 uses this throughout `SimpleAttentionNetwork` for computational efficiency.

### How does HadamardMLP compare to a standard transformer FFN in terms of parameters?

HadamardMLP uses **Walsh-Hadamard transforms** instead of dense weight matrices for its primary expansion (lines 87–103). The transform itself requires no learnable parameters—only element-wise scaling and SiLU activation are learned. This yields fewer total parameters than two dense layers of equivalent width while maintaining comparable representational capacity through the structured transform.

### What tasks benefit from the ContrastiveHead and ConfidenceHead?

**ContrastiveHead** enables dense retrieval by embedding queries and documents into a normalized shared space—useful for semantic search, tool retrieval, and RAG pipelines. **ConfidenceHead** provides per-token uncertainty estimates for applications requiring calibrated predictions, rejection of low-confidence outputs, or active learning. Both heads extend SAN's utility beyond standard autoregressive generation.

### When should Multi-Token Prediction (MTP) be enabled?

Enable MTP (`return_mtp=True`) when **maximizing next-token prediction accuracy** is prioritized over inference speed. The second pass through `self.mtp_block` adds computational overhead but improves prediction quality by explicitly modeling token dependencies. MTP is particularly valuable for offline batch processing or when serving as a teacher model for distillation.