# Learnable Diagonal Gates in the Hadamard MLP: Architecture and Implementation

> Explore learnable diagonal gates in the Hadamard MLP from cactus-compute/needle. Discover how d1, d2, and d3 gates optimize activations with parameter-sparse alternatives.

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

---

**The Hadamard MLP in `cactus-compute/needle` implements three learnable diagonal gates—`d1`, `d2`, and `d3`—that apply element-wise scaling to activations before and after Walsh-Hadamard transforms, replacing dense feed-forward layers with efficient, parameter-sparse alternatives.**

The `HadamardMLP` class in the needle repository reimagines the standard transformer feed-forward network using fast Walsh-Hadamard transforms and learnable diagonal matrices. This architectural choice reduces computational complexity while maintaining expressivity through three trainable vectors that act as diagonal scaling gates. Understanding these gates is essential for researchers modifying the model architecture or analyzing its capacity for signal propagation.

## The Three Learnable Diagonal Gates

According to the implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the Hadamard MLP contains exactly three learnable parameters, all defined as one-dimensional vectors that function as diagonal matrices when multiplied element-wise. Each gate operates at a distinct point in the forward pass:

### d1: Pre-Transform Input Scaling

The first gate, **`d1`**, scales the padded input tensor immediately before the first Hadamard matrix multiplication. Initialized to ones via `jinit.ones`, this gate learns to amplify or suppress specific dimensions of the incoming signal before the orthogonal transformation.

```python
d1 = self.param("d1", jinit.ones, (n,)).astype(self.dtype)

```

In the forward pass, this manifests as `d1 * z` before the matrix multiplication with the fixed Hadamard matrix `H`.

### d2: Intermediate Non-Linearity Scaling

The second gate, **`d2`**, operates after the first Hadamard transform but before the SiLU activation and second transform. Also initialized to ones, this gate shapes the intermediate representation in the Hadamard domain, effectively learning which frequency components to emphasize before the non-linearity is applied.

```python
d2 = self.param("d2", jinit.ones, (n,)).astype(self.dtype)

```

This gate appears in the computation as `nn.silu(d2 * z) @ H`, making it the primary mechanism for learned selectivity within the transform domain.

### d3: Post-Transform Output Scaling

The third gate, **`d3`**, scales the output of the second Hadamard transform before the tensor is sliced back to the original `d_model` dimension. Unlike the previous gates, **`d3` is initialized to `0.02`** using `jinit.constant(0.02)`, providing a small initial scale that grows during training.

```python
d3 = self.param("d3", jinit.constant(0.02), (n,)).astype(self.dtype)

```

This final scaling occurs at `(d3 * z)[..., :self.d_model]`, ensuring the returned tensor matches the expected hidden dimension after processing.

## Forward Pass Implementation

The complete interaction of these three learnable diagonal gates is visible in lines 92-103 of [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The forward pass pads the input to the nearest power-of-two (`n`), applies the fixed orthonormal Hadamard matrix, and interleaves the learnable gates with the transforms:

```python
def __call__(self, x):
    n = 1 << (self.d_model - 1).bit_length()     # pad to power-of-two

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

    d1 = self.param("d1", jinit.ones, (n,)).astype(self.dtype)
    d2 = self.param("d2", jinit.ones, (n,)).astype(self.dtype)
    d3 = self.param("d3", jinit.constant(0.02), (n,)).astype(self.dtype)

    pad = n - self.d_model
    z = jnp.pad(x, ((0, 0), (0, 0), (0, pad))) if pad else x
    z = (d1 * z) @ H               # ← first diagonal gate

    z = nn.silu(d2 * z) @ H        # ← second diagonal gate + non-linearity

    return (d3 * z)[..., :self.d_model]   # ← third diagonal gate

```

All three gates share the same length **`n`**, where `n` represents the smallest power-of-two greater than or equal to `d_model`. This sizing ensures compatibility with the fast Walsh-Hadamard transform implementation while maintaining the diagonal matrix property—each vector position scales exactly one dimension of the activation tensor.

## Practical Usage and Inspection

When initializing a `HadamardMLP` module, the three diagonal gates are created automatically and become accessible through the parameter dictionary. The following example demonstrates instantiation and parameter inspection:

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

# Instantiate with hidden size 128

mlp = HadamardMLP(d_model=128)

# Create dummy input: (batch, seq_len, hidden_dim)

x = jnp.ones((4, 16, 128), dtype=jnp.bfloat16)

# Initialize parameters

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

# Inspect the learnable diagonal gates

print("d1 shape:", params["hadamard_mlp"]["d1"].shape)   # (128,) padded to power-of-two

print("d2 shape:", params["hadamard_mlp"]["d2"]["=", params["hadamard_mlp"]["d2"].shape)
print("d3 shape:", params["hadamard_mlp"]["d3"].shape)

```

To execute a forward pass through the module with the learnable gates active:

```python

# Apply the MLP (inside a Flax/JAX module context)

y = mlp.apply({"params": params}, x)
print(y.shape)   # → (4, 16, 128)

```

## Summary

- **`d1`**, **`d2`**, and **`d3`** are the only learnable parameters in the Hadamard MLP, each functioning as a diagonal matrix via element-wise multiplication.
- All gates have dimension `(n,)`, where `n` is the next power-of-two greater than `d_model`, as implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).
- **d1** scales inputs before the first Hadamard transform, **d2** scales intermediate activations before the SiLU non-linearity and second transform, and **d3** scales outputs before dimensionality reduction.
- The gates are initialized differently: `d1` and `d2` use `jinit.ones`, while `d3` uses `jinit.constant(0.02)` for conservative initial output scaling.

## Frequently Asked Questions

### How do the diagonal gates differ from standard feed-forward weights?

Standard feed-forward layers use dense weight matrices of shape `(d_model, 4*d_model)` or similar, requiring matrix multiplications with quadratic complexity relative to dimension. The Hadamard MLP's **diagonal gates** are vectors that apply element-wise scaling, reducing the parameter count from $O(d_{model}^2)$ to $O(d_{model})$ while the Hadamard matrices provide mixing without learned parameters.

### Why is d3 initialized to 0.02 while d1 and d2 use ones?

The conservative initialization of **d3** to `0.02` via `jinit.constant(0.02)` prevents initial training instability. Since `d3` scales the final output before it is sliced and returned to the residual stream, starting with small values ensures the Hadamard MLP initially acts as a near-identity transformation, allowing the network to gradually learn to utilize the transform capacity as training progresses.

### Can the diagonal gates be regularized or frozen during training?

Yes, because the gates are standard JAX/Flax parameters accessed via `params["hadamard_mlp"]["d1"]` (and similarly for `d2` and `d3`), they can be targeted by standard JAX transformations. You can apply weight decay to these specific parameters, freeze them using `jax.tree_map` with masks, or apply gradient clipping specifically to these gates while updating other model components normally.