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

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, 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.

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:

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:

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

ConfidenceHead

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

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:

logits = model(tokens, quant=True)

Complete Forward Pass Example

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	SAN definition, all sub-modules, heads, MTP logic
needle/model/run.py	Inference and generation utilities
needle/model/tokenizer.py	Text-to-token conversion
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
• 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}