# How Needle 2's Hadamard MLP Differs from Standard Feed-Forward Networks

> Discover how Needle 2's Hadamard MLP revolutionizes feed-forward networks by replacing dense layers with efficient Hadamard transforms and learnable scalings. Slash parameters and computation.

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

---

**Needle 2 replaces conventional two-layer dense feed-forward blocks with a Hadamard MLP that uses the Walsh-Hadamard transform and learnable diagonal scalings instead of full weight matrices, cutting parameters from O(d_model²) to O(n) while reducing computation to O(n log n).**

Standard transformer architectures rely on dense feed-forward networks (FFNs) that consume significant memory and compute. Needle 2, an open-source JAX-based language model from Cactus Compute, takes a fundamentally different approach. Its **Hadamard MLP** reimagines the feed-forward layer through structured linear algebra—trading learned matrices for fixed orthogonal transforms and lightweight diagonal parameters. This article breaks down exactly how this architectural choice differs from standard FFNs, with direct reference to the implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

## Core Architecture Comparison

### Standard Feed-Forward Networks

A conventional FFN in transformers follows a simple pattern: **Linear → Activation → Linear**. Mathematically, for input `x`, this computes:

```

FFN(x) = W₂ · σ(W₁ · x + b₁) + b₂

```

Where `W₁` and `W₂` are dense matrices of shape `(d_model, d_ff)` and `(d_ff, d_model)`, typically with `d_ff = 4 × d_model`. This requires **O(d_model²)** parameters and **O(d_model²)** matrix multiplications that dominate runtime and memory bandwidth.

### Hadamard MLP in Needle 2

According to the Needle 2 source code, the `HadamardMLP` class replaces dense matrices with a composition of diagonal scalings and Walsh-Hadamard transforms:

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

# Example: a tiny Hadamard MLP with d_model=64

d_model = 64
mlp = HadamardMLP(d_model)

# Dummy input: (batch, seq_len, d_model)

x = jnp.ones((2, 10, d_model), dtype=jnp.bfloat16)

# Forward pass – padding to next power-of-two (64 → 64) is automatic

y = mlp(x)   # shape == (2, 10, d_model)

print(y.shape)   # (2, 10, 64)

```

The forward pass executes three operations in sequence (from `HadamardMLP.__call__`, lines 92-103):

1. **First diagonal scaling and Hadamard**: `z = (d1 * x) @ H`
2. **Activation between transforms**: `z = nn.silu(d2 * z) @ H`
3. **Final scaling and crop**: `output = (d3 * z)[..., :d_model]`

Where `H` is the **Walsh-Hadamard matrix**—a fixed, orthogonal matrix with entries of only `+1` and `-1`.

## Parameter Efficiency: From Matrices to Vectors

The most dramatic difference lies in parameterization. As implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 95-98), the Hadamard MLP stores only:

- **Three diagonal vectors**: `d1`, `d2`, `d3` each of size `n`
- Where `n` is the next power-of-two ≥ `d_model`

| Component | Standard FFN | Hadamard MLP |
|-----------|-------------|--------------|
| Trainable weights | Two dense matrices: ~`2 × d_model × d_ff` | Three diagonal vectors: `3n` |
| Parameter count | **O(d_model²)** | **O(n)** where `n ≈ d_model` |
| Storage for d_model=512, d_ff=2048 | ~4.2 million | 1,536 (with n=512) |

The Hadamard matrix itself requires **zero storage**—it's generated on-the-fly by `_walsh_matrix(n)` (lines 80-85) using only bit manipulations to produce the `+1`/`-1` pattern.

## Computational Complexity: Faster Transforms

Standard FFNs bottleneck on dense matrix multiplication. The Hadamard MLP exploits a crucial property: the Walsh-Hadamard transform computes `x @ H` in **O(n log n)** using butterfly operations—only additions and subtractions, no multiplications.

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

cfg = TransformerConfig(d_model=512, num_heads=8, num_kv_heads=4,
                        num_layers=1, dtype="bfloat16")
block = Block(num_heads=cfg.num_heads,
              num_kv_heads=cfg.num_kv_heads,
              d_model=cfg.d_model,
              num_layers=cfg.num_layers,
              dtype=jnp.bfloat16)

x = jnp.ones((1, 16, cfg.d_model), dtype=jnp.bfloat16)
out = block(x)          # runs self-attention then HadamardMLP internally

print(out.shape)        # (1, 16, 512)

```

This lower arithmetic intensity reduces memory bandwidth pressure—a critical optimization for large-scale training.

## Mathematical Properties: Orthogonality and Stability

Dense layers in standard FFNs can represent **any linear map**, offering maximum flexibility. The Hadamard MLP constrains this expressivity deliberately:

- **Orthogonality**: The Hadamard matrix `H` satisfies `H @ H.T = n × I`, preserving vector norms
- **Structured linear maps**: The composition `D₃ · H · D₂ · σ · H · D₁` produces a rich but controlled family of transformations
- **Gradient benefits**: Orthogonal transforms avoid exploding or vanishing singular values, improving training stability

As noted in the `Block.__call__` implementation (lines 35-38), the Hadamard MLP sits inside a residual path with pre-normalization:

```python
x = ZCRMSNorm(dtype=self.dtype, name="pre_hada_norm")(x)
x = HadamardMLP(self.d_model, self.dtype, name="hadamard_mlp")(x)
return skip + x

```

## Implementation in Needle 2

The key files implementing this architecture:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** — Contains the `HadamardMLP` class, `_walsh_matrix()` helper, and `Block` integration (lines 35-38, 80-103)
- **[`needle/model/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/__init__.py)** — Exports model components for external use

The automatic padding to power-of-two dimensions (handled transparently in `HadamardMLP.__call__`) ensures the butterfly algorithm works efficiently even when `d_model` isn't naturally a power of two.

## Summary

- **Parameter reduction**: Hadamard MLP uses O(n) diagonal parameters versus O(d_model²) dense weights
- **Speed**: O(n log n) butterfly operations replace O(d_model²) matrix multiplications
- **Stability**: Orthogonal Hadamard transforms preserve norms and improve gradient flow
- **Integration**: Seamlessly substitutes into standard transformer blocks via `Block.__call__` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)

## Frequently Asked Questions

### What activation function does Needle 2's Hadamard MLP use?

The Hadamard MLP uses **SiLU** (also known as Swish) as its non-linearity. Specifically, in `HadamardMLP.__call__` (line 99), the computation applies `nn.silu(d2 * z)` between the two Walsh-Hadamard transforms. This sigmoid-weighted linear unit provides smooth, non-monotonic activation that pairs well with the linear transform structure.

### Why does the Hadamard MLP require power-of-two dimensions?

The Walsh-Hadamard transform achieves its O(n log n) complexity through recursive butterfly decomposition, which requires the dimension to be a power of two. When `d_model` isn't a power of two, Needle 2 automatically zero-pads to the next power-of-two `n`, computes the transform, then crops back to `d_model`. This happens transparently in `HadamardMLP.__call__` (lines 92-103).

### Can the Hadamard MLP match the expressivity of dense FFNs?

While dense FFNs can represent arbitrary linear maps, the Hadamard MLP's composition of diagonal scalings and fixed orthogonal transforms provides a **structured but expressive** function class. In practice, this trade-off often improves training stability and generalization while maintaining competitive performance—especially when scaled to large model sizes where parameter efficiency matters more.

### Is the Hadamard matrix learned during training?

No. The Hadamard matrix contains **no trainable parameters**—it consists entirely of `+1` and `-1` values in a fixed pattern. Only the three diagonal vectors (`d1`, `d2`, `d3`) are learned. The matrix is generated on-the-fly by `_walsh_matrix(n)` (lines 80-85) using Sylvester's construction, requiring zero storage regardless of model size.