# What Is a Simple Attention Network (SAN)? A Deep Dive into Needle's Core Architecture

> Explore the Simple Attention Network (SAN), a transformer architecture offering high-performance inference on edge devices. Discover its unique Hadamard-based MLPs and grouped-query attention for minimal memory use.

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

---

**The Simple Attention Network (SAN) is a transformer architecture that replaces standard feed-forward networks with Hadamard-based MLPs and employs grouped-query attention to enable high-performance inference on edge devices with minimal memory footprint.**

The **Simple Attention Network (SAN)** serves as the architectural foundation for **Needle 2**, a 45 M‑parameter foundation model developed in the `cactus-compute/needle` repository. Unlike conventional transformers that rely on standard dense feed-forward networks, SAN leverages a fixed Walsh-Hadamard transform and multi-lane routing to deliver tool-calling and structured extraction capabilities within a 14 MB binary footprint.

## Core Components of a Simple Attention Network

### Hadamard MLP

The **Hadamard MLP** replaces the traditional feed-forward network (FFN) in each transformer block. Instead of learnable weight matrices, it applies a fixed Walsh-Hadamard transform (`H`) that operates in *n log n* time complexity. The transform is applied twice with learned diagonal scaling vectors (`d1`, `d2`, `d3`) interspersed between them.

This implementation appears in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) at lines 287‑303:

```python
from needle.model.architecture import HadamardMLP
import jax.numpy as jnp

# The Hadamard MLP uses fixed transforms with learned diagonal scaling

mlp = HadamardMLP(d_model=256, dtype=jnp.bfloat16)
out = mlp(x)  # x shape: (batch, seq_len, d_model)

```

### Grouped-Query Attention (GQA)

**Grouped-Query Attention** reduces the quadratic computational cost of self-attention by sharing key and value heads across groups of query heads. This maintains expressive power while significantly cutting memory bandwidth requirements during inference.

The GQA mechanism is implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) at lines 230‑272, where the `Block` class configures `num_heads` and `num_kv_heads` separately to enable head grouping.

### Engram Key-Value Memory

The **Engram KV Memory** provides a long-range, learnable memory system that can be queried using the current hidden state. This creates "hyper-connections" across layers, allowing the model to access persistent information beyond the standard context window.

The integration occurs in the `Block` forward pass within [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 19‑25), where engram states are mixed into the hidden state before the attention computation.

### Multi-Lane Hyper-Connections

**Multi-lane hyper-connections** split each transformer layer into several parallel "lanes" (`mhc_lanes`). Per-lane routing matrices (`phi_pre`, `phi_res`, and `phi_post`) are normalized using Sinkhorn iteration, creating dense inter-layer communication without introducing additional parameters.

This routing logic resides in the `Stack` module at [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) lines 93‑107:

```python
from needle.model.architecture import Block

# Create a block with multi-lane hyper-connections

san_block = Block(
    num_heads=8,
    num_kv_heads=2,
    d_model=256,
    num_layers=1,
    dtype=jnp.bfloat16,
)

```

### Sandwich Normalization and Gating

Every residual branch—both attention and MLP—employs **Sandwich-Norm** wrapped with **RMSNorm** (`ZCRMSNorm`) and multiplied by a learned gate. This stabilizes training and inference for very small models (under 100M parameters).

The gating mechanism appears in `Block.__call__` at lines 32‑37 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where each branch undergoes normalization, transformation, and gated residual addition.

## Data Flow Through a SAN Block

The forward pass through a Simple Attention Network block follows this strict sequence:

1. **Input** → RMS-normalization → **GQA Self-Attention** → RMS-normalization → gated residual addition
2. **RMS-normalization** → **Hadamard-MLP** → gated residual addition
3. **Engram KV** (optional) mixes into the hidden state before the attention step
4. **Multi-lane routing** blends the output across lanes before feeding the next block

This architecture enables dense computation while maintaining the memory efficiency required for edge deployment.

## Implementation Example

Here is a complete example demonstrating both high-level tool-calling usage and low-level block instantiation:

```python
import needle
from needle.model.architecture import HadamardMLP, Block
import jax.numpy as jnp

# High-level API: Tool-calling with Needle 2

@needle.tool
def get_weather(city: str):
    """Retrieve current weather for a city."""
    return {"city": city, "temp_c": 22, "sky": "sunny"}

agent = needle.Needle(tools=[get_weather])
result = agent.run("What's the weather in Paris?")
print(result["results"])

# Output: [{'city': 'Paris', 'temp_c': 22, 'sky': 'sunny'}]

# Low-level: Direct SAN block usage

x = jnp.zeros((1, 8, 256), dtype=jnp.bfloat16)
san_block = Block(
    num_heads=8,
    num_kv_heads=2,
    d_model=256,
    num_layers=1,
    dtype=jnp.bfloat16,
    flash=False,
)
y = san_block(x)  # Output maintains input shape

print(y.shape)  # (1, 8, 256)

```

## Memory Efficiency and Quantization

The Simple Attention Network achieves its compact footprint through aggressive quantization while preserving the Hadamard transforms. The quantization hooks in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) (lines 6‑10) demonstrate how the model's weights are compressed while keeping the Hadamard diagonal scales in FP16 precision. This approach yields a **14 MB binary** that operates within approximately **28 MB of RAM**, making it suitable for microcontrollers and mobile devices.

## Summary

- **Simple Attention Networks** replace standard FFNs with **Hadamard MLPs** that use fixed transforms and learned diagonal scaling, implemented in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py#L287)‑303).
- **Grouped-Query Attention** reduces memory overhead by sharing KV heads across query groups, defined in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py#L230)‑272).
- **Multi-lane hyper-connections** enable dense layer-to-layer communication via Sinkhorn-normalized routing matrices in the `Stack` module (L93‑107).
- **Engram KV memory** provides persistent cross-layer memory accessed during the block forward pass (L19‑25).
- The architecture fits within 14 MB on disk and 28 MB RAM while maintaining tool-calling capabilities.

## Frequently Asked Questions

### How does the Hadamard MLP in SAN differ from standard transformer FFNs?

Standard FFNs use two dense linear transformations with a non-linearity between them, requiring *O(d²)* parameters. The Hadamard MLP replaces these with fixed Walsh-Hadamard matrices that require no learned weights, instead learning only three diagonal scaling vectors (`d1`, `d2`, `d3`). This reduces parameter count while maintaining *n log n* computational complexity, as implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L287)‑303).

### What is the purpose of multi-lane hyper-connections in SAN?

Multi-lane hyper-connections split each layer into parallel pathways (`mhc_lanes`) and learn routing matrices (`phi_pre`, `phi_post`, `phi_res`) normalized via Sinkhorn iteration. This creates dense communication between layers without adding parameters, allowing information to flow across multiple computational paths. The implementation resides in the `Stack` class at [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) lines 93‑107.

### Why does SAN use Grouped-Query Attention instead of standard multi-head attention?

Grouped-Query Attention (GQA) shares key and value heads across groups of query heads, reducing the memory bandwidth required for KV cache loading during inference. This is crucial for edge deployment where memory constraints are severe. GQA is configured in the `Block` class at [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L230)‑272) by setting `num_kv_heads` lower than `num_heads`.

### How large is a Simple Attention Network model when deployed?

According to the `cactus-compute/needle` source code, a full 45 M‑parameter SAN model (Needle 2) occupies approximately **14 MB** as a binary file and requires roughly **28 MB of RAM** during inference. This efficiency is achieved through the Hadamard MLP's parameter reduction and specialized quantization hooks in [`export.py`](https://github.com/cactus-compute/needle/blob/main/export.py#L6)‑10) that preserve critical weights in FP16 while compressing others.