Simple Attention Network Architecture in Needle: Hadamard MLP and GQA Explained

Needle's Simple Attention Network replaces standard transformer feed-forward layers with a Hadamard MLP and uses Grouped-Query Attention (GQA) to achieve efficient inference on tiny devices.

The Simple Attention Network is the core architecture powering Needle, a 45 million parameter foundation model designed for tool-calling on resource-constrained hardware. Introduced in the Needle 2 paper, this design ditches conventional dense matrix multiplications in favor of orthonormal transforms and query sharing across attention heads. The result is a transformer that maintains competitive capacity while dramatically reducing parameter count and compute requirements.

What Makes the Simple Attention Network "Simple"

Traditional transformers rely on two expensive operations: full attention over all query/key/value projections and dense feed-forward networks with large intermediate dimensions. Needle's Simple Attention Network attacks both bottlenecks through architectural substitutions that preserve expressiveness without learned parameters for the transforms themselves.

Grouped-Query Attention (GQA)

Grouped-Query Attention reduces memory bandwidth and computation by sharing key and value projections across multiple query heads. Instead of every attention head maintaining independent K and V matrices, a smaller set of KV heads serves a larger group of query heads.

As implemented in needle/model/architecture.py, the Block class instantiates MultiHeadAttention with separate num_heads (total query heads) and num_kv_heads (shared KV heads). For Needle-2, this ratio is typically 4:1 — eight query heads share two KV heads.

The GQA mechanism works as follows:

1. Queries project to (batch, seq, num_heads, head_dim)
2. Keys and values project to (batch, seq, num_kv_heads, head_dim)
3. Each query head attends to its assigned KV head via broadcasting

This cuts the KV cache memory by the grouping ratio and reduces the parameter count for K/V projections proportionally.

Hadamard MLP

The Hadamard MLP replaces the standard two-layer feed-forward network with a structure built on the Walsh-Hadamard transform. Located at lines 287-302 of needle/model/architecture.py, the HadamardMLP class implements:

# Conceptual flow (simplified from architecture.py)

def __call__(self, x):
    # Pre-gate RMS normalization

    x = self.norm(x)
    
    # First Hadamard transform

    h = hadamard_transform(x)
    
    # Element-wise gated linear unit

    h = h * gelu(h @ W_gate)
    
    # Second Hadamard transform

    out = hadamard_transform(h)
    
    # Output projection with learned diagonal only

    return out @ W_out

Key characteristics of the Hadamard MLP:

• O(n log n) complexity: The Hadamard transform reduces from O(n²) dense matrix multiplication to O(n log n) via fast Walsh-Hadamard algorithms
• No learned transform weights: The Hadamard matrix H is fixed and orthonormal (HᵀH = I), requiring no gradient updates
• Learned diagonals only: Only the element-wise gates and output projection contain trainable parameters
• Hardware-friendly: The transform consists entirely of additions and subtractions, avoiding expensive multiplication in the core path

The hadamard_transform function applies the recursive butterfly pattern that defines the Walsh-Hadamard matrix, computable efficiently via JAX's lax primitives for XLA compilation.

Block Structure and Residual Flow

Each transformer block in the Simple Attention Network follows a strict gated residual pattern defined in Block.__call__. The implementation at needle/model/architecture.py sequences operations as:

Input
  ↓
RMS-Norm ──→ GQA ──→ RMS-Norm ──→ @ attention_gate ──→ + (residual)
  │                                                            ↓
  └────────────────────────────────────────────────────────────┘
  ↓
RMS-Norm ──→ HadamardMLP ──→ @ mlp_gate ──→ + (residual)
  │                                                            ↓
  └────────────────────────────────────────────────────────────┘
Output

Critical implementation details:

• ZCRMSNorm: Root-mean-square normalization without centering (zero-centering optional), applied before each sub-layer in the pre-norm style
• Learned gates: Both attention and MLP branches multiply by a sigmoid-activated _gate parameter, enabling gradient-based gating of information flow
• Residual connections: Clean addition after each gated block, with the scanned Stack class handling layer-wise carry through jax.lax.scan for efficient compilation

Full Model Composition

The Stack class composes multiple blocks into the complete Simple Attention Network. Running the full forward pass:

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

# Match the Needle-2 checkpoint configuration

cfg = TransformerConfig(
    num_layers=12,
    num_heads=8,
    num_kv_heads=2,       # GQA: 4 query heads per KV head

    d_model=512,
    mhc_lanes=4,
    engram_layers=[3, 7], # Optional engram memory layers

    jax_dtype=jnp.bfloat16,
    flash=True,           # XLA flash attention fusion

)

# Instantiate the stack

stack = Stack(cfg)

# Token embeddings as input

tokens = jnp.zeros((1, 256, cfg.d_model), dtype=jnp.bfloat16)

# Forward pass returns final hidden states

outputs, hidden = stack(tokens)
print(outputs.shape)  # (batch=1, seq=256, d_model=512)

The Stack uses jax.lax.scan to iterate over the 12 layers, carrying the hidden state and optional engram memory through the loop. This pattern enables:

• Constant memory usage regardless of depth (via gradient checkpointing)
• XLA fusion of the entire block sequence
• Optional hyper-connection routing between layers (Sinkhorn-normalized)

Performance and Efficiency Characteristics

The Simple Attention Network design choices yield measurable benefits for edge deployment:

Aspect	Standard Transformer	Simple Attention Network
Feed-forward params	O(d_model × d_ff)	O(d_model) learned diagonals only
Feed-forward compute	O(batch × seq × d_model × d_ff)	O(batch × seq × d_model × log d_model)
KV cache size	2 × num_layers × num_heads × head_dim × seq	2 × num_layers × num_kv_heads × head_dim × seq
Attention memory bandwidth	Proportional to num_heads	Reduced by num_heads / num_kv_heads

The Hadamard MLP's parameter efficiency is particularly stark: where a standard 512-dim model might use 2048-dim FFN intermediates (4× expansion), the Hadamard variant achieves comparable effective capacity with roughly 4× fewer parameters in the feed-forward path.

Integration with Engram Memory and Tool Calling

While the core Simple Attention Network is defined by GQA + Hadamard MLP, Needle extends this base with:

• Engram key-value memory: Additional KV pairs inserted at layers 3 and 7 (configurable) for expanded context capacity without proportional parameter growth
• Tool-calling heads: Specialized output projections for function selection and argument generation, trained with the base model in a unified objective

The engram mechanism reuses the same GQA infrastructure, treating engram entries as additional KV heads that attend to the fixed bank.

Summary

• Grouped-Query Attention (GQA) in Block → MultiHeadAttention reduces KV cache and projection parameters by sharing keys/values across query head groups
• Hadamard MLP at HadamardMLP (lines 287-302, architecture.py) replaces dense feed-forward layers with O(n log n) Walsh-Hadamard transforms and learned diagonals
• Gated residual flow with RMS normalization and learned gates enables stable training of deep stacks with pre-norm architecture
• Full implementation composes through Stack with jax.lax.scan for memory-efficient depth scaling
• 45M parameters total achieves tool-calling capability competitive with larger models through architectural efficiency rather than scale

Frequently Asked Questions

What is the difference between Hadamard MLP and a standard feed-forward network?

A standard feed-forward network computes gelu(x @ W_1) @ W_2 with learned matrices requiring O(d_model × d_ff) parameters and O(batch × seq × d_model × d_ff) operations. The Hadamard MLP substitutes fixed orthonormal Hadamard transforms for W_1 and W_2, reducing learned parameters to output diagonal projections only and operations to O(batch × seq × d_model × log d_model). As implemented in needle/model/architecture.py, this yields comparable expressiveness with roughly 4× parameter reduction.

How does Grouped-Query Attention reduce memory usage?

GQA shares key and value projections across multiple query heads. In Needle's configuration with num_heads=8 and num_kv_heads=2, four query heads reuse the same KV projection. This reduces the KV cache memory footprint by 4× during autoregressive generation and cuts the K/V parameter count from 2 × num_layers × num_heads × head_dim to 2 × num_layers × num_kv_heads × head_dim. The attention computation itself remains identical — queries broadcast to match shared KV dimensions.

Why use the Walsh-Hadamard transform instead of other efficient transforms?

The Walsh-Hadamard transform provides three properties ideal for neural networks: it is orthonormal (preserves norms, no gradient instability), parameter-free (no training needed), and hardware-efficient (butterfly pattern maps well to GPU/TPU shuffle operations). Unlike random projections or learning-based compressions, it guarantees invertibility and energy preservation without optimization. The hadamard_transform implementation in Needle uses JAX primitives that XLA compiles to fused kernels.

Where is the Simple Attention Network defined in the Needle codebase?

The core definitions reside in needle/model/architecture.py: HadamardMLP (lines 287-302) implements the feed-forward replacement, Block contains the GQA instantiation and gated residual flow, and Stack composes blocks into the full model. Configuration via TransformerConfig specifies GQA ratios, layer count, and optional engram layers. The README (assets/architecture.png) provides visual documentation of the data flow.

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