How Needle 2's Simple Attention Network (SAN) Differs from Traditional Transformers

Needle 2's Simple Attention Network (SAN) retains multi-head self-attention but replaces standard MLP layers with Hadamard-based feed-forward networks, adds persistent Engram memory for learned n-gram retrieval, and introduces auxiliary contrastive and confidence heads for retrieval-augmented generation.

The cactus-compute/needle repository implements a novel language model architecture that modernizes the classic transformer while maintaining its core attention mechanism. Needle 2's Simple Attention Network (SAN) introduces several architectural innovations designed to improve computational efficiency and enable retrieval-augmented capabilities. This analysis examines the specific technical differences between SAN and traditional transformer implementations based on the source code in needle/model/architecture.py.

Core Architectural Modifications

Block Composition and Self-Attention

Traditional transformers isolate the MultiHeadAttention block from subsequent feed-forward processing. In needle/model/architecture.py, the Block class (lines 14-38) composes MultiHeadAttention with a HadamardMLP, gating mechanisms, and optional Engram memory injection within a single computational unit. This integrated approach differs from the discrete sub-layer structure found in standard transformer implementations.

HadamardMLP Feed-Forward Networks

Where traditional transformers use a two-layer MLP (Linear → Activation → Linear), SAN implements the HadamardMLP class (lines 86-103). This component utilizes fast Walsh-Hadamard transforms to drastically reduce the computational cost of feed-forward steps while maintaining expressive capacity. The Hadamard-based approach provides significant FLOP reduction compared to conventional dense layers.

Zero-Centered RMSNorm

Standard transformers apply LayerNorm after each sub-layer for stabilization. SAN replaces this with ZCRMSNorm (lines 46-55), a zero-centered RMS normalization variant specifically designed for low-precision training and inference stability. This modification addresses quantization challenges that standard LayerNorm faces during INT8 or lower precision deployment.

Memory Augmentation and Context Management

Engram Memory Tables

Unlike traditional transformers that store all context exclusively in hidden states, SAN introduces Engram tables (lines 81-108) that maintain compressed n-gram statistics. During the forward pass, the _engram_kv method (lines 109-119) retrieves keys and values from these persistent tables and injects them into the attention stream. This creates a learned retrieval mechanism for long-range patterns without increasing the attention window size.

KV-Budget Aware Windowing

Rather than using fixed-size attention windows or simple causal masks, SAN implements dynamic KV-budget computation (kv_budget_window, lines 98-112). This system automatically adjusts the effective attention range based on model size and hardware memory constraints, shrinking the computational window while preserving output quality through intelligent context selection.

Auxiliary Heads and Multi-Token Prediction

Contrastive and Confidence Heads

SAN extends beyond standard output projections with dedicated ContrastiveHead (lines 43-61) and ConfidenceHead (lines 63-70) classes. These operate on hidden cell states to produce L2-normalized embeddings for retrieval-augmented generation and per-token confidence scores, respectively. Traditional transformers lack these auxiliary output pathways.

Multi-Token Prediction Pathway

The SimpleAttentionNetwork class implements an auxiliary MTP (multi-token prediction) pathway within its __call__ method. After processing through the main block stack, the model concatenates shifted embeddings (mtp_emb_norm), combines them via a learned linear projection (mtp_combine), processes them through a secondary Block (mtp_block), and generates additional logits. This architecture improves next-token prediction accuracy and enables more sophisticated decoding strategies than standard single-projection transformers.

Parameter Sharing and Efficiency

SAN employs strategic parameter sharing not present in vanilla transformers. The multi-token prediction pathway reuses attention parameters across multiple passes through the mtp_combine layer, reducing the total parameter count while maintaining model capacity. This differs from traditional architectures where each layer maintains independent parameters.

Practical Code Examples

The following examples demonstrate how to instantiate and utilize SAN's unique features:

# Building a SAN model

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

cfg = TransformerConfig(
    vocab_size=8192,
    d_model=512,
    num_heads=8,
    num_kv_heads=4,
    num_layers=12,
    max_seq_len=2048,
)
model = SimpleAttentionNetwork(config=cfg)

# Dummy input (batch=1, seq_len=10)

tokens = jnp.arange(10)[None, :]
logits = model(tokens)  # Standard forward pass

# Using the contrastive head for retrieval-augmented generation

query = jnp.array([[1, 2, 3, 4]])    # token IDs of a query

tools = jnp.array([[5, 6, 7, 8]])    # token IDs of a tool description

q_emb, t_emb, log_temp = model.forward_contrastive(query, tools)

# q_emb and t_emb are L2-normalized embeddings for cosine similarity comparison

# Obtaining per-token confidence scores

tokens = jnp.array([[10, 11, 12, 13, 0]])  # 0 = pad token

confidence = model.forward_confidence(tokens)

# confidence is float32 array (shape: batch × seq_len) with per-token confidence values

Summary

• HadamardMLP replaces traditional feed-forward layers with Walsh-Hadamard transforms for reduced computational cost.
• Engram memory provides persistent n-gram storage and retrieval, augmenting standard attention with learned long-range patterns.
• ZCRMSNorm substitutes LayerNorm for improved low-precision training stability.
• Contrastive and Confidence heads enable retrieval-augmented generation and uncertainty quantification absent in standard transformers.
• MTP pathway implements multi-token prediction through parameter-shared auxiliary blocks.
• KV-budget windowing dynamically adjusts attention scope based on hardware constraints rather than using fixed windows.

Frequently Asked Questions

What is the primary architectural difference between Needle 2 SAN and standard transformers?

While both use multi-head self-attention, SAN integrates attention with Hadamard-based feed-forward networks, Engram memory, and gating within unified Block classes. Traditional transformers maintain strict separation between attention and feed-forward sub-layers without persistent external memory.

How does HadamardMLP improve efficiency over standard transformer MLPs?

The HadamardMLP class in needle/model/architecture.py utilizes fast Walsh-Hadamard transforms instead of dense matrix multiplications. This reduces the computational complexity of the feed-forward step while maintaining model capacity, resulting in fewer FLOPs per token compared to conventional two-layer MLPs.

What is the purpose of Engram memory in SAN?

Engram tables store compressed n-gram statistics that persist across forward passes. Unlike traditional transformers that rely solely on hidden states for context, SAN retrieves keys and values from these tables via _engram_kv to inject learned patterns into the attention mechanism, improving long-range dependency modeling without quadratic attention costs.

How does the Multi-Token Prediction (MTP) pathway work?

The MTP pathway processes shifted embeddings through a secondary Block (mtp_block) after the main transformer stack, generating additional logits via parameter-shared layers. This allows SAN to predict multiple future tokens simultaneously, improving training efficiency and decoding accuracy compared to standard single-token prediction transformers.

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