# Simple Attention Network (SAN) Architecture: How It Differs from Standard Transformers

> Explore the Simple Attention Network SAN architecture a transformer variant. Discover how it uses Walsh-Hadamard transforms and Engram memory to improve efficiency and context retrieval.

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

---

**The Simple Attention Network (SAN) is a transformer variant implemented in the needle repository that replaces standard feed-forward layers with Walsh-Hadamard transforms, integrates Engram memory modules for n-gram context retrieval, and adds auxiliary heads for contrastive learning and confidence scoring.**

The Simple Attention Network (SAN) serves as the centerpiece sequence model in the cactus-compute/needle codebase, offering a lightweight alternative to conventional transformer architectures while retaining the core self-attention mechanism. Unlike vanilla transformers, SAN incorporates specialized normalization techniques, hardware-optimized attention paths, and novel memory extensions that reduce parameter count and improve inference efficiency. This article examines the SAN implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), detailing its architectural innovations and practical deployment patterns.

## Architectural Differences from Standard Transformers

The SAN architecture diverges from standard transformers across seven key dimensions, from normalization strategies to memory management.

### Zero-Centered RMS Normalization (ZCRMSNorm)

Where standard transformers apply LayerNorm before attention and feed-forward blocks, SAN uses **ZCRMSNorm** (zero-centered RMSNorm) throughout the stack. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the normalization layer normalizes both embeddings and intermediate tensors using a learnable scale added to the input, providing more stable gradients for **bfloat16** training. The embedding layer specifically wraps tokens with `ZCRMSNorm` and scales by `self.embed_scale = math.sqrt(cfg.d_model)`, contrasting with typical learned embeddings that rely solely on √d_model scaling.

### Optimized Attention with Gating and Flash Support

The **`MultiHeadAttention`** module in SAN extends the standard multi-head dot-product attention with three critical enhancements:

- **ZCRMSNorm** applied independently to queries, keys, and values before the attention computation
- **Flash Attention** support for GPU-optimized memory access patterns (optional path)
- **Gating projection** (`gate_proj`) on the attention output that modulates residual connections

Unlike standard transformers that use `nn.MultiHeadDotProductAttention` with post-attention LayerNorm, SAN's attention block supports integrated **fake quantization** (`_aq`, `maybe_quant_kv`) for 8-bit activation and weight inference without external wrappers.

### HadamardMLP: Walsh-Hadamard Transform Feed-Forward

SAN replaces the standard position-wise MLP (Dense-ReLU-Dense) with **`HadamardMLP`**, implementing a two-stage **Walsh-Hadamard transform** via `_walsh_matrix`. This architecture operates in-place and significantly reduces parameter count compared to dense feed-forward networks. The transform uses learned diagonal scalings (`d1`, `d2`, `d3`) to maintain expressivity while enabling faster computation, making it a core differentiator from standard transformer blocks.

### Engram Memory for Extended Context

While standard transformers rely solely on the attention window for long-range dependencies, SAN introduces **Engram** modules that provide n-gram-like KV caches. Located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), Engram generates key-value pairs from token sequences (`engram_indices`, `engram_kv`) and allows the model to attend to past tokens via convolutional taps. This mechanism is controlled by the `engram_layers` configuration parameter, enabling selective memory injection at specific layer depths (e.g., layers 2-10).

### Auxiliary Prediction Heads

SAN attaches specialized heads to hidden representations that standard transformers typically lack:

- **`ContrastiveHead`**: Enables contrastive pre-training for retrieval tasks via `forward_contrastive` and `encode_contrastive` methods
- **`ConfidenceHead`**: Provides token-level confidence scoring through `forward_confidence`, useful for uncertainty quantification during inference

These heads expose intermediate representations via `hidden_cells` and `hidden_states` methods, facilitating downstream tasks without requiring separate encoder models.

### Advanced Masking and Packing

Beyond the standard causal mask (`jnp.tril`), SAN implements **`make_causal_packing_mask`** for segment-aware batch packing and **windowed masking** for KV-budget management. These utilities allow efficient handling of variable-length sequences during training and inference, addressing limitations in standard transformer padding approaches.

## Implementation Structure

The SAN architecture follows a hierarchical Flax module structure defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

**`SimpleAttentionNetwork`** – The top-level `nn.Module` that orchestrates the architecture. It initializes the token `Embedding` wrapped with ZCRMSNorm, constructs the main **`Stack`** of layers, and attaches the contrastive and confidence heads. The `__call__` method computes logits and optionally executes a masked token prediction (MTP) branch.

**`Stack`** – Implements a scanned execution over `Block` layers using `nn.scan` for memory-efficient training across depth.

**`Block`** – Represents a single transformer layer containing:
- Self-attention (`MultiHeadAttention`) with pre-normalization
- Hadamard feed-forward (`HadamardMLP`)
- Gating mechanisms (`_gate`) that control residual pathways

## Practical Usage in Needle

The following snippets demonstrate how to instantiate and utilize the Simple Attention Network within the needle framework:

```python

# -------------------------------------------------

# 1️⃣  Import and configure the model

# -------------------------------------------------

from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

# Compact configuration with Engram memory layers

config = TransformerConfig(
    vocab_size=8192,
    d_model=512,
    num_heads=8,
    num_kv_heads=4,
    num_layers=12,
    max_seq_len=2048,
    engram_layers=(2, 10),   # Enable Engram in layers 2-10

    rope_theta=100_000.0,
)

# -------------------------------------------------

# 2️⃣  Create the model

# -------------------------------------------------

model = SimpleAttentionNetwork(config)

# -------------------------------------------------

# 3️⃣  Dummy input (batch of token IDs)

# -------------------------------------------------

import jax.numpy as jnp
tokens = jnp.ones((2, 128), dtype=jnp.int32)   # batch-size 2, seq-len 128

# -------------------------------------------------

# 4️⃣  Forward pass – obtain logits

# -------------------------------------------------

logits = model(tokens)           # shape: (2, 128, vocab_size)

# -------------------------------------------------

# 5️⃣  Contrastive encoding for retrieval

# -------------------------------------------------

query_emb, key_emb, log_temp = model.forward_contrastive(
    query_tokens=tokens,
    tool_tokens=tokens,
)

# -------------------------------------------------

# 6️⃣  Confidence scoring (token-level)

# -------------------------------------------------

conf_scores = model.forward_confidence(tokens)   # shape: (2, 128)

```

**Implementation Details from Source:**

- `TransformerConfig` is a dataclass consumed by `setup()` (lines 81-90 in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py)) to wire all hyperparameters
- The forward pass automatically generates causal masks if none are provided via `make_causal_mask`
- The `forward_contrastive` and `forward_confidence` methods expose SAN's auxiliary capabilities absent in vanilla transformers

## Summary

- **Simple Attention Network (SAN)** replaces standard transformer MLPs with parameter-efficient **HadamardMLP** layers using Walsh-Hadamard transforms.
- **ZCRMSNorm** provides zero-centered normalization optimized for low-precision training, differing from standard LayerNorm.
- **Engram modules** add n-gram-like KV caches for extended memory beyond the standard attention window.
- Integrated **quantization support** (`_aq`, `maybe_quant_kv`) enables 8-bit inference without external wrappers.
- **Auxiliary heads** for contrastive learning and confidence scoring make SAN suitable for retrieval and uncertainty-aware applications.
- All components are implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) using Flax's `nn.Module` system with scanned layers for memory efficiency.

## Frequently Asked Questions

### How does the Simple Attention Network differ from a standard transformer?

While both use multi-head self-attention, SAN replaces dense feed-forward layers with **HadamardMLP** (Walsh-Hadamard transforms), uses **ZCRMSNorm** instead of LayerNorm, adds **Engram** memory modules for n-gram context, and includes native quantization support. According to the needle source code, these changes reduce parameter count and improve inference speed while maintaining modeling capacity.

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

The **Engram** module generates KV caches from token n-grams (`engram_indices`, `engram_kv`) and injects them into specific layers defined by `engram_layers`. This allows the model to access convolutional taps of past context, effectively extending memory beyond the standard attention window without increasing the KV-cache size for every token.

### Why does SAN use HadamardMLP instead of standard dense layers?

**HadamardMLP** implements a fast **Walsh-Hadamard transform** (`_walsh_matrix`) that operates in-place with learned diagonal scalings (`d1`, `d2`, `d3`). This approach reduces the parameter count compared to traditional Dense-ReLU-Dense blocks while maintaining computational efficiency, making it particularly advantageous for deployment on memory-constrained hardware.

### How do I access the contrastive and confidence heads in SAN?

The `SimpleAttentionNetwork` class exposes these through dedicated methods: use `forward_contrastive(query_tokens, tool_tokens)` to obtain query/key embeddings for retrieval tasks, and `forward_confidence(tokens)` to get token-level uncertainty scores. These methods access intermediate representations via `hidden_cells` and `hidden_states`, providing functionality not present in standard transformer implementations.