# Simple Attention Network (SAN) Architecture in Needle: Complete Component Breakdown

> Explore the Simple Attention Network SAN architecture in Needle. Understand its core components including rotary embeddings Engram memory Hadamard MLPs and auxiliary heads.

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

---

**The Simple Attention Network (SAN) is a Flax-based transformer with rotary embeddings, Engram memory, Hadamard MLPs, and auxiliary heads for contrastive learning and confidence estimation.**

This guide examines the SAN architecture defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the core model powering Needle's language modeling capabilities. SAN implements a modern transformer with several efficiency innovations: learned n-gram memory via **Engram**, fast feed-forward computation through **HadamardMLP**, and optional multi-token prediction for improved training signal.

## Core Architecture Components

### SimpleAttentionNetwork: Top-Level Module

The `SimpleAttentionNetwork` class (lines 78-83) serves as the main entry point, orchestrating all sub-components during forward passes. It initializes embeddings, the transformer stack, auxiliary heads, and optional multi-token-prediction pathways.

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

cfg = TransformerConfig(
    vocab_size=32000,
    d_model=768,
    num_heads=12,
    num_layers=27,
    max_seq_len=2048,
)
model = SimpleAttentionNetwork(config=cfg)

```

### Embedding Layer

Token IDs map to dense vectors via `nn.Embed`, scaled by √d_model (lines 83-85). This scaling stabilizes early training by matching the magnitude of subsequent residual stream updates.

## Transformer Stack: Stack and Block

### Stack: Layer-Wise Computation

The `Stack` module (lines 79-86) implements the deep transformer using a **scanned `Block`** with optional reversible checkpointing. Scanning reduces compilation time and memory overhead for deep networks. The stack returns final hidden states and optionally exposes per-layer hidden cells for auxiliary heads.

### Block: Single Transformer Layer

Each `Block` (lines 105-138) contains:

- **Multi-head self-attention** via `MultiHeadAttention`
- **HadamardMLP** for feed-forward processing
- **Gating and residual connections**
- **Optional Engram KV injection** when `engram_kv` is supplied

The modular design allows selective memory augmentation at specific layers.

## Attention Mechanism: MultiHeadAttention

`MultiHeadAttention` (lines 18-78) implements standard multi-head attention with several optimizations:

- **Flash attention** support for memory-efficient computation
- **RMS normalization** on queries and keys
- **Rotary positional embeddings** (RoPE) applied to Q/K
- **Per-head gating** for adaptive attention weighting

The `_rope` method (lines 4-7 of `SimpleAttentionNetwork`) pre-computes cosine/sine tables for position encoding.

## Novel Components: Engram and HadamardMLP

### Engram: Learned N-Gram Memory

The `Engram` class (lines 81-87) provides **long-range memory** through a learned table producing additional key/value pairs from token n-grams. Configure injection points via `engram_layers`:

```python
cfg = TransformerConfig(
    engram_layers=(2, 15),   # inject at layers 2 and 15

    engram_slots=8192,       # memory table size

)

```

Engram enables the model to attend to n-gram patterns without extending context length.

### HadamardMLP: Efficient Feed-Forward

`HadamardMLP` (lines 87-103) replaces conventional feed-forward networks with a **Walsh-Hadamard transform** for fast dense mixing. This reduces parameter count and computation while maintaining expressivity—critical for scaling to larger models.

## Auxiliary Heads and Training Objectives

### ContrastiveHead

Projects pooled hidden states into a **retrieval embedding space** (lines 44-61). Use `encode_contrastive()` for document retrieval or semantic similarity tasks:

```python
emb = model.encode_contrastive(tokens)  # (batch, seq, contrastive_dim)

```

### ConfidenceHead

Produces **per-position uncertainty estimates** (lines 63-77) via `forward_confidence()`:

```python
scores = model.forward_confidence(tokens)  # (batch, seq)

```

These scores identify low-confidence generations for rejection sampling or active learning.

## Multi-Token Prediction (MTP) Pathway

SAN supports an **auxiliary MTP stream** that predicts multiple future tokens. When active (lines 94-100, 124-138), the pathway:

1. Concatenates current hidden state with next-token embedding
2. Passes through `mtp_block` (small transformer)
3. Applies dual normalization layers
4. Produces secondary logits

Enable with `return_mtp=True` during training for denser supervision signal.

## Masking and Quantization

### Attention Masks

Utility functions create causal and padding masks:

- `make_causal_mask` — prevents attending to future positions
- `make_padding_mask` — excludes padding tokens (lines 88-95)

### 8-Bit Inference Hooks

Optional fake-quantization via `_aq` and `_quantize.maybe_quant_*` (lines 22-26, 42-44) enables 8-bit inference without graph modifications. Set `quant=True` in forward calls:

```python
logits = model(tokens, quant=True)

```

## Complete Forward Pass Example

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

# Configuration with Engram and MTP

cfg = TransformerConfig(
    vocab_size=32000,
    d_model=768,
    num_heads=12,
    num_kv_heads=6,
    num_layers=27,
    max_seq_len=2048,
    engram_layers=(2, 15),
    engram_slots=8192,
)

model = SimpleAttentionNetwork(config=cfg)

# Forward pass with all outputs

tokens = jnp.ones((2, 128), dtype=jnp.int32)
logits = model(tokens)                           # Primary logits

contrastive = model.encode_contrastive(tokens)   # Retrieval embeddings

confidence = model.forward_confidence(tokens)    # Uncertainty scores

# 8-bit quantized inference

logits_q = model(tokens, quant=True)

```

## Source File Reference

| File | Contents |
|------|----------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | SAN definition, all sub-modules, heads, MTP logic |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Inference and generation utilities |
| [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) | Text-to-token conversion |
| [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) | Quantization primitives |
| `tests/*.py` | Forward pass, contrastive head, MTP tests |

## Summary

- **SimpleAttentionNetwork** orchestrates the full SAN architecture in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)
- **Stack** and **Block** implement the transformer core with scanning and reversible checkpointing
- **MultiHeadAttention** combines flash attention, RMS-norm, rotary embeddings, and per-head gating
- **Engram** provides learned n-gram memory at configurable layers
- **HadamardMLP** replaces standard FFNs with efficient Walsh-Hadamard transforms
- **ContrastiveHead** and **ConfidenceHead** enable retrieval and uncertainty estimation
- **MTP pathway** adds multi-token prediction for improved training
- **Quantization hooks** support 8-bit inference without model changes

## Frequently Asked Questions

### What makes SAN different from standard transformer architectures?

SAN incorporates three key innovations: **Engram** for learned n-gram memory injection, **HadamardMLP** for parameter-efficient feed-forward computation, and native **multi-token prediction** support. These components are modular—Engram layers are configurable, HadamardMLP replaces standard FFNs, and MTP can be disabled—allowing flexible trade-offs between capability and efficiency.

### How does Engram memory work in practice?

Engram maintains a learned table of size `engram_slots` that produces additional key/value pairs from token n-grams. During forward passes in specified `engram_layers`, these KV pairs concatenate with standard self-attention keys and values, effectively extending the model's receptive field to n-gram patterns without increasing sequence length or attention complexity.

### When should I use the ContrastiveHead vs. standard hidden states?

Use `encode_contrastive()` when you need **task-agnostic embeddings** for retrieval, clustering, or semantic search. The ContrastiveHead projects pooled hidden cells through a dedicated transformation optimized for similarity metrics. For downstream classification, standard hidden states or confidence-weighted representations may perform better.

### Does quantization affect model quality in SAN?

The fake-quantization implementation in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) uses learnable clipping ranges during training, so quantization-aware fine-tuning preserves quality. For 8-bit inference without retraining, expect minor degradation—test on your specific task using the `forward_confidence()` outputs to flag potentially unreliable predictions.