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

> Explore Needle's Simple Attention Network, featuring Hadamard MLP and Grouped-Query Attention GQA. Achieve efficient inference on tiny devices with this innovative architecture.

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

---

**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](https://arxiv.org/abs/2607.18363), 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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `HadamardMLP` class implements:

```python

# 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`](https://github.com/cactus-compute/needle/blob/main/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:

```python
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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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.