# How Does the Hadamard MLP Differ from a Standard FFN in Needle 2?

> Discover how the Hadamard MLP in Needle 2 revolutionizes FFNs by reducing parameters and complexity using diagonal scalings and Walsh-Hadamard transforms. Learn more!

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

---

**The Hadamard MLP in Needle 2 replaces the standard two-layer feed-forward network with three learned diagonal scalings and fixed Walsh-Hadamard transforms, cutting parameter count from O(d_model × d_ff) to O(d_model) and computational complexity from O(d_model·d_ff) to O(d_model·log d_model).**

Needle 2, the lightweight transformer implementation from `cactus-compute/needle`, introduces a **Hadamard MLP** as a drop-in replacement for the conventional feed-forward network (FFN). This architectural change dramatically reduces both memory footprint and compute while preserving expressive capacity. This article breaks down exactly how the Hadamard MLP differs from a standard FFN, with direct reference to the implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).

## Key Differences: Hadamard MLP vs. Standard FFN

The two architectures solve the same problem—adding non-linear transformations after attention—but with fundamentally different approaches to weight structure and computation.

### Layer Composition and Weight Structure

**Standard FFN**: Uses two dense weight matrices (`W₁`, `W₂`) forming a classic bottleneck architecture: project up, apply non-linearity, project down. Parameter count scales quadratically with the expansion factor.

**Hadamard MLP**: Uses **three diagonal parameter vectors** (`d1`, `d2`, `d3`) and a **fixed, non-learned Walsh-Hadamard matrix** `H`. The matrix `H` is generated on-the-fly via `_walsh_matrix()` and is orthogonal, meaning it preserves signal norm without requiring gradient updates.

From [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 87-103):

```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()          # power-of-2 size

        H = _walsh_matrix(n).astype(self.dtype)            # fixed Hadamard matrix

        d1 = self.param("d1", jinit.ones, (n,)).astype(self.dtype)   # diagonal 1

        d2 = self.param("d2", jinit.ones, (n,)).astype(self.dtype)   # diagonal 2

        d3 = self.param("d3", jinit.constant(0.02), (n,)).astype(self.dtype)  # diagonal 3

        pad = n - self.d_model
        z = jnp.pad(x, ((0, 0), (0, 0), (0, pad))) if pad else x
        z = (d1 * z) @ H                     # first scaling + Hadamard

        z = nn.silu(d2 * z) @ H              # SiLU + second scaling + Hadamard

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

```

### Computational Complexity

The computational savings come from the **Fast Walsh-Hadamard Transform (FWHT)**.

- **Standard FFN**: Two heavy matrix multiplications at `O(d_model · d_ff)` each. With typical `d_ff = 4 × d_model`, this is `O(d_model²)`.
- **Hadamard MLP**: The Walsh-Hadamard transform runs in `O(d_model · log d_model)`. The three diagonal multiplications are negligible at `O(d_model)`.

This logarithmic factor becomes significant at scale, especially for Needle 2's deployment targets.

### Memory Usage

Memory reduction is equally dramatic:

- **Standard FFN**: Stores two dense tensors of shape `(d_model, d_ff)` and `(d_ff, d_model)`.
- **Hadamard MLP**: Stores only three vectors of size `≈ d_model`. For `d_model = 768` and `d_ff = 3072`, this cuts learnable parameters from ~4.7M to ~2.3K—a **~2000× reduction**.

### Non-Linearity Placement

Both use **SiLU activation** (`nn.silu`), but placement differs:

- **Standard FFN**: Applied once between the two dense projections.
- **Hadamard MLP**: Applied **after the second diagonal scaling**, between the two Hadamard transforms. This creates a "sandwich" structure: scale → transform → activate-scale → transform → scale.

## Integration in the Transformer Block

The Hadamard MLP replaces the FFN at the block level. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `Block.__call__` method shows this substitution clearly:

```python
class Block(nn.Module):
    ...
    @nn.compact
    def __call__(self, x, mask=None, rope=None, quant=False, engram_kv=None, site_flags=None):
        # ... attention computation omitted ...

        skip = x
        x = ZCRMSNorm(dtype=self.dtype, name="pre_hada_norm")(x)
        # Hadamard MLP replaces traditional FFN

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

```

The normalization layer is renamed `pre_hada_norm` to reflect this architectural shift.

## Contrast with Standard FFN Implementation

For reference, a conventional FFN matching Needle 2's interface would look like:

```python
class StandardFFN(nn.Module):
    d_model: int
    d_ff: int = 4 * d_model          # typical expansion factor

    dtype: jnp.dtype = jnp.bfloat16

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.d_ff, dtype=self.dtype)(x)
        x = nn.silu(x)
        x = nn.Dense(self.d_model, dtype=self.dtype)(x)
        return x

```

The Hadamard MLP achieves comparable expressiveness through its **learned diagonal modulations** of the fixed orthogonal basis, rather than learning full dense transformations.

## Why Needle 2 Uses the Hadamard MLP

The design prioritizes **efficient deployment** without sacrificing quality. The Hadamard matrix's orthogonality ensures stable gradients and preserved signal norms. The learned diagonals `d1`, `d2`, `d3` provide sufficient flexibility for downstream tasks while the fixed `H` eliminates that parameter burden entirely.

This matches Needle 2's goal as a lightweight transformer—minimal FLOPs, minimal memory, maximal portability.

## Summary

- **Weight structure**: Hadamard MLP uses 3 diagonal vectors vs. 2 dense matrices in standard FFN.
- **Parameter count**: Reduced from `O(d_model × d_ff)` to `O(d_model)`.
- **Compute**: Reduced from `O(d_model · d_ff)` to `O(d_model · log d_model)` via FWHT.
- **Implementation**: Located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as `HadamardMLP`, integrated in `Block.__call__`.
- **Non-linearity**: SiLU applied between two Hadamard transforms vs. between dense projections.

## Frequently Asked Questions

### Why does the Hadamard MLP pad to power-of-2 dimensions?

The Fast Walsh-Hadamard Transform requires input lengths that are powers of 2. The implementation computes `n = 1 << (self.d_model - 1).bit_length()` to find the next power of 2, pads the input, and crops the output back to `d_model`. This adds minimal overhead for typical model dimensions.

### Can the Hadamard MLP match a standard FFN's expressiveness?

Yes. The combination of learned diagonal scales and orthogonal transforms can approximate any linear transformation arbitrarily well, given sufficient width. In practice, the three diagonal modulations provide enough flexibility for transformer-level performance at a fraction of the cost.

### Where is the Walsh-Hadamard matrix defined in Needle 2?

The helper function `_walsh_matrix()` generates the fixed Hadamard matrix on-the-fly in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). It produces a deterministic, orthogonal matrix without requiring storage or gradient computation.

### Does the Hadamard MLP work with quantization?

Yes. All operations—diagonal scaling, matrix multiplication with `H`, and SiLU—are compatible with standard quantization schemes. The fixed `H` matrix simplifies quantization-aware training since its values are known constants.