What Is a Hadamard MLP and How Does It Replace FFN in Transformers?

A Hadamard MLP is a low-parameter feed-forward block that replaces dense weight matrices with Walsh-Hadamard transforms and three learned diagonal scaling vectors, reducing memory and compute from O(n²) to O(n log n).

The Hadamard MLP is a specialized feed-forward layer developed for the Needle transformer architecture. Unlike standard transformer feed-forward networks (FFNs) that rely on two large dense layers, this block uses fast orthogonal transforms to mix hidden states with minimal learnable parameters. This makes it attractive for large-scale models where memory efficiency and inference speed matter.

How Hadamard MLP Works: Architecture Breakdown

The implementation lives in needle/model/architecture.py at lines 87-103. The core mechanism sandwiches two Walsh-Hadamard transforms between three element-wise scaling operations:

# https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L87-L103

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

    @nn.compact
    def __call__(self, x):
        # Pad to the nearest power-of-2 length

        n = 1 << (self.d_model - 1).bit_length()
        H = _walsh_matrix(n).astype(self.dtype)          # Normalised 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 Hadamard transform

        z = nn.silu(d2 * z) @ H               # Non-linearity + second transform

        return (d3 * z)[..., : self.d_model] # Trim back to original width

Step 1: Power-of-Two Padding and Hadamard Matrix Generation

The input dimension is padded to the nearest power of two (n). This allows the use of a normalized Walsh-Hadamard matrix H of size n×n, which is orthogonal and supports fast O(n log n) computation. The helper _walsh_matrix(n) generates this fixed transform matrix—no gradients flow through it.

Step 2: Three Learned Diagonal Scales

Only three vectors are learned:

• d1: scales input before first transform (initialized to ones)
• d2: scales between transforms (initialized to ones, followed by SiLU)
• d3: scales output before truncation (initialized to 0.02)

These provide per-feature flexibility at O(n) parameter cost versus O(n²) for dense matrices.

Step 3: Dual Transform with SiLU Activation

Two Hadamard transforms with an intervening SiLU activation (nn.silu) create a non-linear feature mixing pathway. The composition H → SiLU → H approximates the expressivity of dense layers without storing weight matrices.

Step 4: Dimension Truncation

After the final scaling by d3, the padded dimensions are sliced away to restore the original d_model size.

Hadamard MLP vs. Standard FFN: Key Differences

Aspect	Standard FFN	Hadamard MLP
Parameters	Two dense matrices: O(n²)	Three diagonal vectors: O(n)
Compute	Two O(n²) matrix multiplications	Two O(n log n) Hadamard transforms
Memory	Stores large weight tensors	Stores only 3n scalars
Mixing mechanism	Direct learned linear maps	Orthogonal basis with learned scalings
Best for	General flexibility	Large-scale efficiency (Needle's target)

The parameter reduction is dramatic: for d_model=2048, a standard FFN uses ~8.4M parameters (assuming 4× expansion), while Hadamard MLP uses only ~6,144 parameters—roughly 1/1400th the size.

Using Hadamard MLP in Practice

Inside a Needle Transformer Block

The Hadamard MLP replaces the conventional FFN at lines 34-38 of architecture.py:

# https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L34-L38

class Block(nn.Module):
    # … (other fields omitted)

    @nn.compact
    def __call__(self, x, mask=None, rope=None, quant=False, engram_kv=None, site_flags=None):
        # … (self-attention omitted)

        # --- Hadamard MLP insertion point ---

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

The pattern follows standard transformer design: normalize, apply the feed-forward block, add residual. Only the internal computation differs.

Standalone Usage

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

# Batch of sequences: (batch=2, seq_len=16, hidden=512)

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

# Create module

hadamard_mlp = HadamardMLP(d_model=512)

# Initialize parameters

variables = hadamard_mlp.init(jax.random.PRNGKey(0), x)

# Forward pass

y = hadamard_mlp.apply(variables, x)
print(y.shape)  # (2, 16, 512)

This minimal example shows the API surface: pass d_model, initialize with a PRNG key, and apply.

Implementation Files and References

File	Purpose	Key Content
needle/model/architecture.py	Core implementation	HadamardMLP class, _walsh_matrix helper, Block integration
needle/model/export.py	Quantization handling	Hadamard diagonals remain FP16 while weights may quantize
tests/test_lora.py	Verification	LoRA fine-tuning compatibility tests for Hadamard components

According to the Needle source code, the d3 parameter's small initialization (0.02) ensures stable gradients at training start, while d1 and d2 start at identity-like behavior.

Why Replace FFN with Hadamard MLP?

Transformer FFNs typically dominate parameter count (often 2/3 of total weights). The Hadamard MLP removes this bottleneck by:

• Eliminating dense weight storage entirely
• Reducing compute to near-linear complexity
• Preserving full feature mixing via orthogonal transforms
• Enabling higher effective model capacity at fixed memory budgets

This trade-off sacrifices some architectural flexibility for massive efficiency gains—ideal for inference-bound deployment scenarios that Needle targets.

Summary

• Hadamard MLP replaces dense FFN layers with Walsh-Hadamard transforms and three learned diagonal scales.
• Parameter count drops from O(n²) to O(n); compute drops from O(n²) to O(n log n).
• Implementation in needle/model/architecture.py uses power-of-two padding, dual transforms with SiLU, and careful initialization.
• Integration follows standard transformer block patterns with pre-normalization and residual connections.

Frequently Asked Questions

What makes the Walsh-Hadamard transform suitable for neural networks?

The Walsh-Hadamard transform is orthogonal, invertible, and computable in O(n log n) time via fast algorithms. Unlike random or learned projections, it provides guaranteed feature mixing without training instability. According to the Needle implementation, the fixed matrix H requires no gradients and can be generated on-the-fly or cached.

Does padding to power-of-two limit model flexibility?

The padding adds at most ~50% overhead to dimension (worst case when d_model is just above a power of two). In practice, this is negligible compared to the memory saved by eliminating dense weights. The _walsh_matrix function handles arbitrary input sizes transparently.

Why use SiLU specifically in the Hadamard MLP?

SiLU (x * sigmoid(x)) provides smooth, non-monotonic non-linearity that helps the dual-transform architecture approximate complex functions. The source code explicitly uses nn.silu(d2 * z) between transforms, differing from common GELU or ReLU choices in standard FFNs.

Can Hadamard MLP be combined with quantization?

Yes. The export.py file shows Hadamard diagonals remain in FP16 during quantization workflows, preserving their precision while other weights may compress to lower bitwidths. This reflects their outsized importance to model quality relative to their small size.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}