# What Is a Hadamard MLP and How Does It Replace Standard FFNs in Needle 2?

> Discover how Hadamard MLPs replace standard FFNs in Needle 2. Learn about parameter efficiency, reduced memory bandwidth, and maintained model capacity with fixed Walsh-Hadamard transforms and learned scalings.

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

---

**A Hadamard MLP is a parameter-efficient replacement for feed-forward networks that uses fixed Walsh-Hadamard transforms and learned diagonal scalings instead of dense weight matrices, reducing memory bandwidth and parameter count while maintaining model capacity.**

The Needle 2 architecture from the cactus-compute/needle repository implements a **Simple Attention Network** that substitutes traditional feed-forward networks (FFNs) with this specialized layer. By combining orthonormal Walsh-Hadamard matrices with only three trainable scaling vectors, the architecture achieves significant efficiency gains without sacrificing expressive power.

## How the Hadamard MLP Works

At its foundation, the Hadamard MLP replaces the two dense linear transformations of a standard FFN with a structured, fixed matrix that requires no gradient updates.

### The Walsh-Hadamard Transform Foundation

The implementation constructs a **Walsh-Hadamard matrix** `H` that operates in $O(n \log n)$ time without learned weights. This matrix is orthonormal, satisfying $H^T H = I$, and serves as the fixed backbone for all linear transformations within the layer.

### Architecture Implementation in needle/model/architecture.py

The `HadamardMLP` class defined at lines 287-303 implements the forward pass using three learned diagonal parameters (`d1`, `d2`, `d3`) alongside the fixed matrix `H`:

```python
class HadamardMLP(nn.Module):
    d_model: int
    dtype: jnp.dtype = jnp.bfloat16

    @nn.compact
    def __call__(self, x):
        n = 1 << (self.d_model - 1).bit_length()
        H = _walsh_matrix(n).astype(self.dtype)               # ← fixed orthonormal matrix

        d1 = self.param("d1", jinit.ones, (n,)).astype(self.dtype)
        d2 = self.param("d2", jinit.ones, (n,)).astype(self.dtype)
        d3 = self.param("d3", jinit.constant(0.02), (n,)).astype(self.dtype)
        pad = n - self.d_model
        z = jnp.pad(x, ((0, 0), (0, 0), (0, pad))) if pad else x
        z = (d1 * z) @ H                                      # ← first linear + Hadamard

        z = nn.silu(d2 * z) @ H                               # ← non-linearity + second Hadamard

        return (d3 * z)[..., : self.d_model]                  # ← final scaling & truncate

```

The architecture applies **SILU activation** between two Hadamard transformations, with each transformation preceded by element-wise multiplication with learned diagonal scalings (`d1` and `d2`). The final output applies a third scaling vector (`d3`) before truncating to the original model dimension.

## Hadamard MLP vs. Standard FFN Architecture

Replacing a conventional FFN with a Hadamard MLP fundamentally alters the computational and memory characteristics of transformer blocks.

### Parameter Efficiency and Memory Bandwidth

A standard FFN relies on two dense matrices ($W_1$ and $W_2$) each containing $O(d_{model}^2)$ parameters. In contrast, the Hadamard MLP stores only three diagonal vectors (`d1`, `d2`, `d3`) with $O(d_{model})$ total parameters. This design eliminates learned weight matrices from the linear steps, drastically reducing memory bandwidth requirements during both training and inference phases.

### Computational Complexity

While standard FFNs perform dense matrix multiplications costing $O(d_{model}^2)$ operations, the Hadamard MLP exploits the fast Walsh-Hadamard transform structure. The fixed matrix `H` applies in $O(n \log n)$ time where $n$ represents the next power of two greater than or equal to `d_model`, making the linear transforms asymptotically faster than dense alternatives for large dimensions.

## Integration into Needle 2's Transformer Block

Within Needle 2's `Block` class, the Hadamard MLP replaces the traditional FFN sub-layer immediately following the attention mechanism. Line 336 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) demonstrates this integration:

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

```

Each transformer block applies **ZCRMSNorm** normalization before feeding inputs into the Hadamard MLP. According to the README.md (lines 21-26), this configuration constitutes the core of the "Simple Attention Network" recipe, emphasizing that the Hadamard MLP replacement is fundamental to the architecture's efficiency.

## Practical Implementation Example

You can instantiate and apply the Hadamard MLP directly using the Needle 2 implementation:

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

# Minimal example: apply a Hadamard MLP to a dummy tensor

batch, seq, dim = 2, 16, 128
x = jnp.ones((batch, seq, dim), dtype=jnp.bfloat16)

mlp = HadamardMLP(d_model=dim, dtype=jnp.bfloat16, name="hadamard_mlp")
params = mlp.init(jax.random.PRNGKey(0), x)   # initialise learned diag scalings

y = mlp.apply(params, x)                     # forward pass

print(y.shape)  # → (2, 16, 128)

# Full block usage (attention + Hadamard MLP)

block = Block(
    num_heads=4,
    num_kv_heads=2,
    d_model=dim,
    num_layers=1,
    dtype=jnp.bfloat16,
    flash=True,
)
block_params = block.init(jax.random.PRNGKey(0), x)
out, _ = block.apply(block_params, x)        # returns transformed representation

print(out.shape)  # → (2, 16, 128)

```

The [`export.py`](https://github.com/cactus-compute/needle/blob/main/export.py) and [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py) modules demonstrate how these parameters persist during model serialization and inference, ensuring the diagonal scalings remain available throughout the model lifecycle.

## Summary

- **A Hadamard MLP replaces dense FFN layers** with fixed Walsh-Hadamard transforms and learned diagonal scalings, reducing learnable parameters from $O(d^2)$ to $O(d)$.
- **The implementation resides in** [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) **lines 287-303**, utilizing JAX/Flax with bfloat16 precision support.
- **Three trainable parameters** (`d1`, `d2`, `d3`) provide per-dimension scaling flexibility, while fixed orthonormal matrices handle linear transformations without gradient updates.
- **Integration occurs post-attention** in each transformer block following ZCRMSNorm normalization, as implemented at line 336.
- **Computational complexity improves** to $O(n \log n)$ for the linear components compared to $O(n^2)$ for standard dense layers, significantly reducing memory bandwidth.

## Frequently Asked Questions

### What makes a Hadamard MLP different from a standard FFN?

A standard FFN uses two dense weight matrices with $O(d_{model}^2)$ parameters each and performs general matrix multiplication. A Hadamard MLP uses fixed, orthonormal Walsh-Hadamard matrices requiring no training data and only three diagonal scaling vectors with $O(d_{model})$ parameters. The fixed matrices handle the bulk of linear transformation while the learned scalings (`d1`, `d2`, `d3`) provide task-specific adaptability.

### Why does Needle 2 use fixed Walsh-Hadamard matrices instead of learned weights?

Fixed Walsh-Hadamard matrices are orthonormal and computable in $O(n \log n)$ time without consuming parameter memory or gradient bandwidth. By freezing these transformations and training only the diagonal scalings, Needle 2 achieves substantial memory efficiency while the orthonormal property ensures stable gradient flow through the network.

### How does the Hadamard MLP affect model performance?

The Hadamard MLP reduces memory footprint and computational overhead compared to dense FFNs, allowing Needle 2 to scale to larger contexts efficiently. The Simple Attention Network design leverages these savings while the learned diagonal parameters (`d1`, `d2`, `d3`) preserve sufficient expressive capacity for complex tasks, as evidenced by the architecture's integration in [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py) for inference workloads.

### Can the Hadamard MLP be ported to other transformer architectures?

Yes, the implementation follows standard JAX/Flax patterns and can replace FFN layers in other transformer models. Key requirements include ensuring the hidden dimension aligns with power-of-two constraints (handled automatically via padding in the `__call__` method) and initializing the three diagonal parameters (`d1`, `d2`, `d3`) with appropriate default values (1.0, 1.0, and 0.02 respectively).