# How the Hadamard MLP Works in Needle 2's Architecture

> Discover how the Hadamard MLP in Needle 2's architecture optimizes transformers. Learn how Walsh-Hadamard transforms and learned scales reduce memory bandwidth and boost performance. Read now for insights!

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

---

**The HadamardMLP layer replaces traditional dense feed-forward networks with Walsh-Hadamard transforms and learned diagonal scales, reducing memory bandwidth while maintaining expressive power in Needle 2's transformer blocks.**

The HadamardMLP is implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as a core component of the Needle 2 transformer. Unlike conventional MLPs that rely on dense weight matrices, this custom feed-forward block utilizes fast Walsh-Hadamard transforms combined with learnable diagonal scaling to create a memory-efficient pathway. According to the cactus-compute/needle source code, this approach significantly reduces GPU memory footprint while preserving model expressiveness through three carefully initialized diagonal parameters.

## Power-of-2 Dimension Alignment

The layer first guarantees dimension compatibility with fast transform algorithms by rounding the hidden dimension up to the nearest power of two. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) around lines 93-94, the implementation calculates the padded size `n` and applies zero-padding when `d_model` is not already a power of two:

```python
n = 1 << (self.d_model - 1).bit_length()                     # line 93-94

pad = n - self.d_model
z = jnp.pad(x, ((0, 0), (0, 0), (0, pad))) if pad else x   # line 98-99

```

This padding ensures the Walsh-Hadamard matrix operates on compatible tensor dimensions while maintaining the original model width after truncation.

## Walsh-Hadamard Matrix Generation

The normalized Walsh-Hadamard matrix `H` of size `n × n` is generated once per forward pass via the `_walsh_matrix` helper function. This orthogonal matrix satisfies `H·Hᵀ = I` and enables O(n log n) fast Hadamard transforms rather than O(n²) dense multiplication:

```python
H = _walsh_matrix(n).astype(self.dtype)                     # line 94-95

```

Because `H` is fixed and orthogonal, it requires no gradient updates and can be computed efficiently using bit manipulation operations inherent to fast Hadamard transforms.

## Learned Diagonal Parameters

Instead of learning full weight matrices, the HadamardMLP uses three trainable diagonal vectors—`d1`, `d2`, and `d3`—to provide the layer's adaptive capacity. These are initialized in lines 95-98 of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) using JAX's parameter system, with `d3` specifically initialized to a small constant for training stability:

```python
d1 = self.param("d1", jinit.ones, (n,)).astype(self.dtype)   # line 95-96

d2 = self.param("d2", jinit.ones, (n,)).astype(self.dtype)   # line 96-97

d3 = self.param("d3", jinit.constant(0.02), (n,)).astype(self.dtype) # line 97-98

```

Each vector is broadcast across the hidden dimension `n`, allowing the layer to scale inputs, intermediates, and outputs independently while maintaining only O(n) trainable parameters.

## Dual Transform with SiLU Activation

The forward pass applies two Hadamard transforms interleaved with diagonal scaling and non-linearity. First, the padded input is element-wise multiplied by `d1` and matrix-multiplied with `H`. The result is then scaled by `d2`, passed through the SiLU activation (`x·σ(x)`), and multiplied by `H` again. Finally, `d3` scales the output before slicing back to the original dimension:

```python
z = (d1 * z) @ H                                            # line 100-101

z = nn.silu(d2 * z) @ H                                     # line 101-102

return (d3 * z)[..., : self.d_model]                       # line 102-103

```

This sequence effectively approximates a two-layer MLP without requiring dense weight matrices, using the fixed `H` matrices to mix features while the diagonal parameters control the transformation magnitude.

## Memory and Compute Efficiency

The HadamardMLP achieves **O(n log n)** computational complexity compared to the **O(n²)** operations required by standard dense MLPs. Because the Walsh-Hadamard matrix `H` is fixed and orthogonal, the two matrix multiplications can be implemented using fast Hadamard transforms rather than general matrix multiplication kernels.

This design eliminates the storage of large weight matrices entirely—the only trainable parameters are the three diagonal vectors of length `n`. According to the `Block` class implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 34-36), the HadamardMLP integrates directly into each transformer block as the feed-forward component, contributing to Needle 2's high throughput and low-memory characteristics during both training and inference.

## Practical Implementation

You can instantiate the HadamardMLP directly for custom architectures or rely on the integrated `SimpleAttentionNetwork` class:

```python

# Direct usage of HadamardMLP module

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

# Dummy input: (batch, seq_len, d_model)

x = jnp.ones((2, 16, 768), dtype=jnp.bfloat16)

# Initialize the module

mlp = HadamardMLP(d_model=768)

# Initialize parameters with a random PRNG key

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

# Apply the MLP

output = mlp.apply(params, x)          # shape → (2, 16, 768)

```

For full model integration, the HadamardMLP is automatically configured within the transformer blocks:

```python

# Usage within complete Needle 2 model

from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

cfg = TransformerConfig(d_model=768, num_layers=27, num_heads=12)
model = SimpleAttentionNetwork(config=cfg)

# Initialize parameters

rng = jax.random.PRNGKey(42)
params = model.init(rng, jnp.ones((1, 32), dtype=jnp.int32))

# Run forward pass

logits = model.apply(params, jnp.ones((1, 32), dtype=jnp.int32))

```

## Summary

- **Matrix-free architecture** eliminates dense Linear layers, replacing them with fast Walsh-Hadamard transforms that require no trainable weights.
- **Triple diagonal learning** uses parameters `d1`, `d2`, and `d3` to maintain expressiveness with only O(n) trainable values rather than O(n²).
- **Automatic power-of-2 alignment** handles padding via bitwise operations, ensuring compatibility with fast transform algorithms without manual dimension management.
- **SiLU-gated dual pass** applies non-linearity between two Hadamard transforms, enabling complex function approximation despite the linear nature of the transforms.
- **Block integration** places the HadamardMLP within each transformer Block in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), serving as the primary feed-forward mechanism in Needle 2's architecture.

## Frequently Asked Questions

### How does the Hadamard MLP reduce memory usage compared to standard MLPs?

The HadamardMLP eliminates dense weight matrices entirely. While standard MLPs store two matrices of shape `(d_model, d_model)` requiring O(n²) parameters, the HadamardMLP uses only three diagonal vectors of length `n` (the next power of two), reducing trainable parameters to O(n). Additionally, the fixed Walsh-Hadamard matrix `H` requires no storage or gradients, further reducing GPU memory footprint during training.

### Why is power-of-2 padding necessary for the Hadamard transform?

Walsh-Hadamard matrices inherently require dimensions that are powers of two. The implementation automatically handles mismatched dimensions by computing `n = 1 << (self.d_model - 1).bit_length()` in line 93 of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py), rounding up to the nearest power of two, then zero-padding the input tensor. This ensures compatibility with fast Hadamard transform algorithms while preserving the original model width through post-operation truncation.

### What role does the SiLU activation play in the HadamardMLP?

The SiLU activation (implemented as `nn.silu` in line 101-102 of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py)) provides the non-linear transformation between the two Hadamard passes. Without this non-linearity, the composition of two linear transforms would collapse into a single linear operation. By applying SiLU (`x * sigmoid(x)`) after the `d2` scaling, the layer gains the capacity to approximate complex functions despite using fixed orthogonal matrices for feature mixing.

### How is the HadamardMLP integrated into Needle 2's transformer blocks?

According to the `Block` class definition in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 34-36), the HadamardMLP is invoked within each transformer block's `__call__` method immediately after the attention mechanism. It serves as the feed-forward network (FFN) component, replacing the traditional Linear-ReLU-Linear stack. This integration maintains the standard transformer architecture pattern while providing the computational efficiency benefits of the Hadamard transform approach.