What Is a Hadamard MLP and How Does It Differ From a Standard FFN?

A Hadamard MLP is a parameter-efficient feed-forward block that replaces dense weight matrices with Walsh-Hadamard transforms and learned diagonal scaling vectors, reducing computational complexity from O(n²) to O(n log n) while maintaining expressivity comparable to a standard FFN.

The Hadamard MLP appears in the Needle transformer architecture (cactus-compute/needle) as a specialized replacement for traditional feed-forward networks. According to the implementation in needle/model/architecture.py, this block eliminates conventional dense linear layers entirely, using instead a sequence of orthogonal matrix multiplications and element-wise scaling operations. If you are optimizing large-scale transformers for memory efficiency, understanding the distinction between a Hadamard MLP and a standard FFN is essential for reducing both parameter count and compute overhead.

Architecture of the Hadamard MLP

Unlike a standard FFN that follows a Dense → Activation → Dense pattern, the Hadamard MLP applies a Walsh-Hadamard transform sandwiched between three learned diagonal scaling vectors.

Core Components in needle/model/architecture.py

The implementation lives in needle/model/architecture.py at lines 87-103. The HadamardMLP class inherits from nn.Module and defines three trainable parameters—d1, d2, and d3—representing diagonal scaling vectors, alongside a normalized Walsh-Hadamard matrix H generated via _walsh_matrix.

# 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

The Forward Pass Flow

The __call__ method executes the following steps:

1. Padding: Input tensors are zero-padded to the nearest power-of-two dimension n, ensuring compatibility with the Walsh-Hadamard matrix.
2. First Scaling and Transform: The input is multiplied element-wise by d1, then matrix-multiplied by the orthogonal Hadamard matrix H (an O(n log n) operation).
3. Non-linearity and Second Transform: After scaling by d2, the SiLU activation is applied, followed by a second multiplication by H.
4. Final Scaling and Trimming: The result is scaled by d3 (initialized to 0.02) and trimmed back to the original d_model dimension.

Hadamard MLP vs Standard FFN: Technical Comparison

Aspect	Standard FFN	Hadamard MLP
Parameter Count	Two dense weight matrices (O(n²)) plus biases	Three diagonal vectors (O(n))
Compute Complexity	O(n²) matrix multiplications	Two O(n log n) Hadamard transforms
Memory Footprint	Stores large weight tensors	Stores only three small vectors
Expressivity	Direct linear mapping through learned weights	Implicit linear mapping via orthogonal basis
Primary Operations	Dense → Activation → Dense	Scale → Hadamard → Scale → SiLU → Hadamard → Scale

Why the Hadamard MLP Reduces Parameters

A standard FFN in a transformer typically uses two weight matrices of size d_model × d_model (or 4*d_model in Gated Linear Unit variants), resulting in O(n²) parameters. In contrast, the Hadamard MLP uses only three vectors of length n (the padded dimension), yielding O(n) parameters. Because the Walsh-Hadamard matrix H is fixed (non-learnable) and orthogonal, it mixes all feature dimensions without requiring trainable weights, while the learned diagonal vectors provide per-feature flexibility.

Implementation Details and Usage

Integrating HadamardMLP Into Transformer Blocks

Within the Needle architecture, HadamardMLP replaces the traditional FFN inside transformer blocks. The following excerpt from needle/model/architecture.py (lines 34-38) shows its placement after normalization:

# In a Transformer block (excerpt from architecture.py)

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

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

    @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 block applies ZCRMSNorm normalization, passes the result through HadamardMLP, and adds a residual connection.

Standalone Instantiation

You can import and use HadamardMLP directly for custom architectures:

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

# Dummy input: (batch, seq_len, hidden_dim)

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

# Create the module (same hidden size as input)

hadamard_mlp = HadamardMLP(d_model=512)

# Initialise parameters

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

# Apply the module

y = hadamard_mlp.apply(variables, x)

print("Output shape:", y.shape)   # → (2, 16, 512)

This minimal example demonstrates initialization and forward execution on a random tensor.

Related Files in the Repository

• needle/model/export.py: References Hadamard-specific quantization details, noting that "norms, Hadamard diagonals, and gates stay FP16" during export.
• tests/test_lora.py: Contains unit tests verifying HadamardMLP behavior under LoRA fine-tuning configurations.

Summary

• The Hadamard MLP replaces dense weight matrices with Walsh-Hadamard transforms and three learned diagonal vectors (d1, d2, d3), reducing parameters from O(n²) to O(n).
• It achieves O(n log n) computational complexity versus the O(n²) cost of standard FFN layers by leveraging fast orthogonal matrix multiplication.
• The implementation in needle/model/architecture.py uses SiLU activation between two Hadamard transforms and requires input padding to the nearest power-of-two dimension.
• This architecture is particularly suited for large-scale transformers where memory efficiency and parameter reduction are prioritized over the raw capacity of dense linear layers.

Frequently Asked Questions

What makes a Hadamard MLP more efficient than a standard FFN?

A Hadamard MLP stores only three diagonal vectors (O(n) parameters) compared to the two dense weight matrices (O(n²) parameters) of a standard FFN. Additionally, the Walsh-Hadamard transform operates in O(n log n) time versus the O(n²) matrix multiplications of dense layers, significantly reducing both memory footprint and computational overhead for large hidden dimensions.

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

The Walsh-Hadamard matrix H is inherently a square matrix of size 2^k × 2^k. The implementation pads inputs to the nearest power-of-two (n = 1 << (self.d_model - 1).bit_length()) to ensure compatible matrix dimensions. After processing, the output is trimmed back to the original d_model size, discarding the padded dimensions.

How does the SiLU activation function fit into the Hadamard MLP architecture?

SiLU (Sigmoid Linear Unit) is applied after the second diagonal scaling (d2) but before the second Hadamard transform. Specifically, the operation nn.silu(d2 * z) @ H introduces non-linearity between the two orthogonal transforms, enabling the block to approximate complex functions despite using fixed, non-learnable mixing matrices.

Can the Hadamard MLP be used with mixed-precision training?

Yes. The reference implementation in needle/model/architecture.py defaults to jnp.bfloat16 for both the Hadamard matrix and the diagonal parameters. The export.py module specifically handles quantization scenarios where Hadamard diagonals maintain FP16 precision while other weights may be quantized, ensuring numerical stability during mixed-precision training and inference.

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