# Hadamard MLP in Needle 2: Parameter-Efficient Feed-Forward Layers Using Walsh-Hadamard Transforms

> Explore the Hadamard MLP in Needle 2, a parameter-efficient feed-forward layer using Walsh-Hadamard transforms for dense-like capacity with O(n log n) complexity and fewer parameters.

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

---

**The Hadamard MLP is a specialized feed-forward layer in Needle 2 that replaces conventional dense weight matrices with three learned diagonal vectors and fast Walsh-Hadamard transforms, achieving dense-like capacity with O(n log n) complexity and minimal parameter count.**

The Hadamard MLP serves as the core feed-forward component in Needle 2’s transformer architecture, offering a structurally efficient alternative to traditional multi-layer perceptrons. Unlike standard MLPs that rely on full matrix multiplications, this layer leverages orthogonal Walsh-Hadamard matrices and diagonal weight tensors to reduce the parameter footprint while preserving model expressiveness. According to the cactus-compute/needle source code, the implementation strikes a balance between computational efficiency and representational power.

## Architecture of the Hadamard MLP

### Structured Diagonal Weights

Instead of conventional weight matrices, the Hadamard MLP learns three diagonal tensors named `d1`, `d2`, and `d3`. These vectors are initialized using `jinit.ones` (with `d3` optionally set to a small constant), replacing the enormous parameter counts of dense linear layers. This diagonal structure reduces storage requirements to **O(n)** per tensor rather than **O(n²)**.

### The Walsh-Hadamard Transform

The layer utilizes a normalized Walsh-Hadamard matrix `H` generated by the `_walsh_matrix` helper function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). Because `H` is orthogonal and its application requires only **O(n log n)** operations via the fast Walsh-Hadamard transform, the layer avoids expensive dense matrix multiplications entirely. The matrix is constructed once and reused across forward passes.

### Power-of-Two Padding Strategy

To ensure compatibility with square Hadamard matrices, the input dimension is padded to the next power of two using the expression `n = 1 << (d_model-1).bit_length()`. This guarantees that a valid n×n Hadamard matrix exists for the computation. After processing, outputs are truncated back to the original `d_model` dimension.

## Implementation Details in Needle 2

The core implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) between lines 287-303. The `HadamardMLP` class encapsulates the forward logic, while the `_walsh_matrix` utility handles construction of the transform matrix. Within transformer blocks, the standard MLP is replaced by calling `HadamardMLP(self.d_model, self.dtype, name="hadamard_mlp")(x)` at lines 335-337.

## Computational Flow

The forward pass follows a precise four-step sequence:

1. **Padding**: Input `x` is padded to size `n` (the next power of two).
2. **First Transform**: The padded input undergoes element-wise multiplication with `d1`, followed by the Walsh-Hadamard transform `H`.
3. **Activation and Second Transform**: After element-wise scaling by `d2`, SiLU activation is applied, followed by a second `H` transform.
4. **Projection and Truncation**: The result is scaled by `d3` and truncated to the original `d_model` dimension.

## Integration in Transformer Blocks

In Needle 2's transformer architecture, the `Block.__call__` method invokes the Hadamard MLP to process attention outputs. This substitution occurs at lines 335-337 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where `HadamardMLP` replaces the conventional `Dense → Activation → Dense` pattern. The layer preserves the model's dtype (typically FP16 or BFloat16) throughout computation.

## Usage Examples

### Inside a Custom Flax Module

```python

# Example: Using HadamardMLP inside a custom Flax module

import flax.linen as nn
import jax.numpy as jnp
from needle.model.architecture import HadamardMLP

class SimpleHadamardBlock(nn.Module):
    d_model: int = 512
    dtype: jnp.dtype = jnp.bfloat16

    @nn.compact
    def __call__(self, x):
        # x shape: (batch, seq_len, d_model)

        # Apply the Hadamard‑based MLP

        y = HadamardMLP(self.d_model, self.dtype, name="hadamard_mlp")(x)
        return y

# Running a forward pass

import jax
key = jax.random.PRNGKey(0)
batch = jax.random.normal(key, (2, 16, 512))
model = SimpleHadamardBlock()
variables = model.init(key, batch)          # initializes d1, d2, d3

output = model.apply(variables, batch)     # shape (2, 16, 512)

print(output.shape)                         # → (2, 16, 512)

```

### Direct Layer Invocation

```python

# Example: Directly calling the layer in a notebook

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

# Create a dummy input (batch=1, seq=8, d_model=256)

x = jnp.arange(1 * 8 * 256, dtype=jnp.float32).reshape(1, 8, 256)

# Instantiate the HadamardMLP

mlp = HadamardMLP(d_model=256, dtype=jnp.float32, name="hadamard_mlp")

# Initialise parameters

rng = jax.random.PRNGKey(42)
params = mlp.init(rng, x)

# Apply the layer

y = mlp.apply(params, x)
print(y.shape)   # → (1, 8, 256)

```

## Summary

- The Hadamard MLP replaces dense weight matrices with three diagonal vectors (`d1`, `d2`, `d3`) and orthogonal Walsh-Hadamard transforms.
- Implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 287-303) uses **O(n log n)** fast transforms rather than **O(n²)** matrix multiplication.
- Input dimensions are padded to power-of-two sizes to accommodate square Hadamard matrices.
- The layer is invoked in transformer blocks at lines 335-337, offering significant parameter reduction while maintaining expressive capacity.
- FP16-friendly design preserves numerical precision through the orthogonal transform structure.

## Frequently Asked Questions

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

A standard MLP uses two dense weight matrices (W1 and W2) requiring **O(n²)** parameters each. The Hadamard MLP in Needle 2 replaces these with three diagonal vectors totaling **O(n)** parameters and utilizes the Walsh-Hadamard matrix for mixing, reducing storage by orders of magnitude while maintaining similar representational capacity through the orthogonal transform structure.

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

Walsh-Hadamard matrices exist only for dimensions that are powers of two (n = 2^k). The implementation automatically pads inputs using `n = 1 << (d_model-1).bit_length()` to satisfy this mathematical constraint, then truncates the output to restore the original dimensionality.

### How does the Hadamard MLP achieve O(n log n) complexity?

Instead of dense matrix multiplication (**O(n²)**), the layer applies the fast Walsh-Hadamard transform algorithm. This recursive decomposition allows the transform to be computed in **O(n log n)** time, similar to the FFT, making the layer significantly faster than conventional MLPs for large dimensions.

### Where is the Hadamard MLP implemented in the Needle 2 codebase?

The `HadamardMLP` class is defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) at lines 287-303, with the `_walsh_matrix` helper function generating the transform matrix. Usage within transformer blocks appears at lines 335-337, where it replaces the traditional feed-forward network in the `Block` class.