How Needle 2's Hadamard MLP Differs from Standard Feed-Forward Networks

Needle 2 replaces conventional two-layer dense feed-forward blocks with a Hadamard MLP that uses the Walsh-Hadamard transform and learnable diagonal scalings instead of full weight matrices, cutting parameters from O(d_model²) to O(n) while reducing computation to O(n log n).

Standard transformer architectures rely on dense feed-forward networks (FFNs) that consume significant memory and compute. Needle 2, an open-source JAX-based language model from Cactus Compute, takes a fundamentally different approach. Its Hadamard MLP reimagines the feed-forward layer through structured linear algebra—trading learned matrices for fixed orthogonal transforms and lightweight diagonal parameters. This article breaks down exactly how this architectural choice differs from standard FFNs, with direct reference to the implementation in needle/model/architecture.py.

Core Architecture Comparison

Standard Feed-Forward Networks

A conventional FFN in transformers follows a simple pattern: Linear → Activation → Linear. Mathematically, for input x, this computes:

FFN(x) = W₂ · σ(W₁ · x + b₁) + b₂

Where W₁ and W₂ are dense matrices of shape (d_model, d_ff) and (d_ff, d_model), typically with d_ff = 4 × d_model. This requires O(d_model²) parameters and O(d_model²) matrix multiplications that dominate runtime and memory bandwidth.

Hadamard MLP in Needle 2

According to the Needle 2 source code, the HadamardMLP class replaces dense matrices with a composition of diagonal scalings and Walsh-Hadamard transforms:

from needle.model.architecture import HadamardMLP

# Example: a tiny Hadamard MLP with d_model=64

d_model = 64
mlp = HadamardMLP(d_model)

# Dummy input: (batch, seq_len, d_model)

x = jnp.ones((2, 10, d_model), dtype=jnp.bfloat16)

# Forward pass – padding to next power-of-two (64 → 64) is automatic

y = mlp(x)   # shape == (2, 10, d_model)

print(y.shape)   # (2, 10, 64)

The forward pass executes three operations in sequence (from HadamardMLP.__call__, lines 92-103):

1. First diagonal scaling and Hadamard: z = (d1 * x) @ H
2. Activation between transforms: z = nn.silu(d2 * z) @ H
3. Final scaling and crop: output = (d3 * z)[..., :d_model]

Where H is the Walsh-Hadamard matrix—a fixed, orthogonal matrix with entries of only +1 and -1.

Parameter Efficiency: From Matrices to Vectors

The most dramatic difference lies in parameterization. As implemented in needle/model/architecture.py (lines 95-98), the Hadamard MLP stores only:

• Three diagonal vectors: d1, d2, d3 each of size n
• Where n is the next power-of-two ≥ d_model

Component	Standard FFN	Hadamard MLP
Trainable weights	Two dense matrices: ~2 × d_model × d_ff	Three diagonal vectors: 3n
Parameter count	O(d_model²)	O(n) where n ≈ d_model
Storage for d_model=512, d_ff=2048	~4.2 million	1,536 (with n=512)

The Hadamard matrix itself requires zero storage—it's generated on-the-fly by _walsh_matrix(n) (lines 80-85) using only bit manipulations to produce the +1/-1 pattern.

Computational Complexity: Faster Transforms

Standard FFNs bottleneck on dense matrix multiplication. The Hadamard MLP exploits a crucial property: the Walsh-Hadamard transform computes x @ H in O(n log n) using butterfly operations—only additions and subtractions, no multiplications.

import jax.numpy as jnp
from needle.model.architecture import Block, TransformerConfig

cfg = TransformerConfig(d_model=512, num_heads=8, num_kv_heads=4,
                        num_layers=1, dtype="bfloat16")
block = Block(num_heads=cfg.num_heads,
              num_kv_heads=cfg.num_kv_heads,
              d_model=cfg.d_model,
              num_layers=cfg.num_layers,
              dtype=jnp.bfloat16)

x = jnp.ones((1, 16, cfg.d_model), dtype=jnp.bfloat16)
out = block(x)          # runs self-attention then HadamardMLP internally

print(out.shape)        # (1, 16, 512)

This lower arithmetic intensity reduces memory bandwidth pressure—a critical optimization for large-scale training.

Mathematical Properties: Orthogonality and Stability

Dense layers in standard FFNs can represent any linear map, offering maximum flexibility. The Hadamard MLP constrains this expressivity deliberately:

• Orthogonality: The Hadamard matrix H satisfies H @ H.T = n × I, preserving vector norms
• Structured linear maps: The composition D₃ · H · D₂ · σ · H · D₁ produces a rich but controlled family of transformations
• Gradient benefits: Orthogonal transforms avoid exploding or vanishing singular values, improving training stability

As noted in the Block.__call__ implementation (lines 35-38), the Hadamard MLP sits inside a residual path with pre-normalization:

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

Implementation in Needle 2

The key files implementing this architecture:

• needle/model/architecture.py — Contains the HadamardMLP class, _walsh_matrix() helper, and Block integration (lines 35-38, 80-103)
• needle/model/__init__.py — Exports model components for external use

The automatic padding to power-of-two dimensions (handled transparently in HadamardMLP.__call__) ensures the butterfly algorithm works efficiently even when d_model isn't naturally a power of two.

Summary

• Parameter reduction: Hadamard MLP uses O(n) diagonal parameters versus O(d_model²) dense weights
• Speed: O(n log n) butterfly operations replace O(d_model²) matrix multiplications
• Stability: Orthogonal Hadamard transforms preserve norms and improve gradient flow
• Integration: Seamlessly substitutes into standard transformer blocks via Block.__call__ in needle/model/architecture.py

Frequently Asked Questions

What activation function does Needle 2's Hadamard MLP use?

The Hadamard MLP uses SiLU (also known as Swish) as its non-linearity. Specifically, in HadamardMLP.__call__ (line 99), the computation applies nn.silu(d2 * z) between the two Walsh-Hadamard transforms. This sigmoid-weighted linear unit provides smooth, non-monotonic activation that pairs well with the linear transform structure.

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

The Walsh-Hadamard transform achieves its O(n log n) complexity through recursive butterfly decomposition, which requires the dimension to be a power of two. When d_model isn't a power of two, Needle 2 automatically zero-pads to the next power-of-two n, computes the transform, then crops back to d_model. This happens transparently in HadamardMLP.__call__ (lines 92-103).

Can the Hadamard MLP match the expressivity of dense FFNs?

While dense FFNs can represent arbitrary linear maps, the Hadamard MLP's composition of diagonal scalings and fixed orthogonal transforms provides a structured but expressive function class. In practice, this trade-off often improves training stability and generalization while maintaining competitive performance—especially when scaled to large model sizes where parameter efficiency matters more.

Is the Hadamard matrix learned during training?

No. The Hadamard matrix contains no trainable parameters—it consists entirely of +1 and -1 values in a fixed pattern. Only the three diagonal vectors (d1, d2, d3) are learned. The matrix is generated on-the-fly by _walsh_matrix(n) (lines 80-85) using Sylvester's construction, requiring zero storage regardless of model size.

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