The Hadamard MLP: Channel Mixing for Almost No Parameters

On a phone or a microcontroller, a token's cost is not the arithmetic. It is the bytes: every weight the model touches has to come out of flash or DRAM, and in a small transformer two thirds of those bytes belong to the feed-forward layers. Our attention-only study showed that the feed-forward network's parameters matter and its functional form largely does not. That leaves a design question: if what the block needs is channel mixing and a nonlinearity, how few parameters can buy it? Needle's answer is the Hadamard MLP. This post is the construction, with the maths, and the bill.

What a feed-forward layer does

A standard block mixes across positions with attention and across channels with a two-layer network,

$\mathrm{FFN}(x) = W_2\, \sigma(W_1 x), \qquad W_1 \in \mathbb{R}^{4d \times d},\ W_2 \in \mathbb{R}^{d \times 4d},$

so at $d = 768$ it stores $2 \cdot 768 \cdot 3072 \approx 4.7\text{M}$ parameters and spends the same number of multiply-adds on every token. A dense matrix is the most general linear map, and also the most expensive: it lets every channel talk to every other channel, but it pays $d^2$ to do so. The question is whether every-to-every mixing has a cheaper form.

The transform that mixes everything for free

It does. The Walsh-Hadamard transform is an orthogonal matrix of $\pm 1/\sqrt{n}$ entries defined recursively,

$H_1 = [1], \qquad H_{2n} = \frac{1}{\sqrt{2}} \begin{bmatrix} H_n & H_n \\ H_n & -H_n \end{bmatrix},$

and it can be applied with the butterfly network below: $\log_2 n$ stages of pairwise sums and differences, $n \log_2 n$ additions in total, and not a single stored weight.

The Walsh-Hadamard transform. Each stage pairs channels a fixed distance apart and replaces them by their sum and difference, so n channels are fully mixed in n log₂ n additions with no weights at all. Needle's MLP starts from this transform and learns away from it.

Every output channel depends on every input channel, which is exactly the property a channel mixer needs. What it lacks is anything to learn. Needle 2 used the transform as is, with learned diagonal scalings on either side, on the observation that a fixed rotation plus learned per-channel gains is already a surprisingly capable mixer. Needle 3 keeps the idea and makes the rotation itself learnable, without giving back the cost.

Kronecker factors, and why 32 × 32

Write the $n = 1024$ channels of a token, after padding from 768, as a $32 \times 32$ tile $Z$ . A Kronecker-factored matrix $A \otimes B$ with $A, B \in \mathbb{R}^{32 \times 32}$ acts on the tile by mixing its rows with one factor and its columns with the other,

$(A \otimes B)\, \mathrm{vec}(Z) = \mathrm{vec}\!\left(A^{\top} Z\, B\right),$

which is two small matrix products, $2 \cdot 32^3 = 65{,}536$ multiply-adds, against the $1024^2 \approx 10^6$ of a dense matrix of the same size, and $2 \cdot 32^2 = 2{,}048$ parameters against a million. The Walsh matrix is itself a Kronecker product of smaller Walsh matrices, so initialising $A$ and $B$ to $H_{32}$ makes the stage an exact Hadamard transform on the first step of training, and every step after that can move it.

One Monarch stage. A 1,024 × 1,024 mixing matrix is replaced by the Kronecker product of two 32 × 32 matrices, which is applied by folding the token into a tile and multiplying its rows, then its columns. The engine runs each stage this way, in registers.

One Kronecker stage only mixes within rows and within columns of the tile. Two channels in different rows and different columns never meet. The fix is the same one that turns butterflies into a full transform: permute between stages. Needle applies three Kronecker mixes with a fixed random shuffle of the channels after the first and the second, so that by the third stage every channel has had a path to every other. This is the Monarch construction, a product of block-structured matrices with permutations between them, and it is expressive enough to represent the dense mixers a transformer actually learns while staying a few tiny matmuls to apply.

The whole layer

Putting the pieces in order, with $\Pi_1, \Pi_2$ the fixed permutations, $D_1 \dots D_4$ learned diagonals and $b$ a bias, the layer computes

$y = D_4\, M_3\, D_3\, \Pi_2\, M_2\, \mathrm{silu}\!\big(D_2\, c(x) \odot \Pi_1\, M_1\, D_1\, x + b\big), \qquad M_s = A_s \otimes B_s,$

with one addition that earns its keep: the gain on the middle diagonal is conditioned on the input through a rank-8 bottleneck,

$c(x) = 1 + \mathrm{softmax}(x V)\, U, \qquad V \in \mathbb{R}^{768 \times 8},\ U \in \mathbb{R}^{8 \times 1024},$

with $U$ zero at initialisation so the gain starts at exactly one. It costs 14K multiply-adds and lets a token scale its own channels before the nonlinearity, a small amount of the input-dependence a dense feed-forward layer gets from its width, bought with a vector rather than a matrix. $D_4$ starts at 0.02 so each new layer writes gently into the residual stream, the same small-init convention as the attention output.

Three Kronecker mixes, two fixed permutations between them, four learned diagonals, one bias, one activation and a rank-8 gain that lets the input modulate its own channel scaling. Everything mixes; almost nothing is stored.

Counting everything, a layer holds $3 \times 2 \times 32^2$ factor weights, five vectors of length 1,024 and the two rank-8 matrices: 25.6K parameters, and about 0.21M multiply-adds per token.

The bill

At d = 768 with a 4d hidden width, a dense feed-forward layer holds 4.7M parameters and spends 4.7M multiply-adds per token; the Hadamard MLP holds 25.6K and spends about 0.21M. Across the 20-layer model that is the difference between 296 and 100 MFLOPs per token, with attention and the engram unchanged.

Per layer the Hadamard MLP is 180 times smaller than the dense layer it replaces and 22 times cheaper to run. Across the 20 blocks of Needle 3 the difference is 94M parameters, which would have taken the model from 121M to 215M, and 196 MFLOPs per token, which would have taken it from 100 to 296. Attention and the engram memory are unchanged by the swap, and the engram is where Needle keeps the knowledge a feed-forward layer would otherwise have to hold: 70.8M of its 121M parameters live in hashed n-gram tables that are read by gather, a few rows per token, costing no arithmetic at all. The feed-forward budget was spent on memory that is only touched when it is needed, rather than on a matrix that is touched every token.

On the device

The bytes argument is what the engine cashes in. Each Kronecker stage is applied as the two $32 \times 32$ products of the diagram, on a tile that fits in vector registers, so a layer's entire channel mixer streams through the CPU in a few hundred fused multiply-adds per lane with no weight matrix to fetch. The three factor pairs, the diagonals and the rank-8 gain are a little over a hundred kilobytes per layer at full precision, and far less in the shipped 2-bit archive, so the mixer stays resident in cache for the life of a session. Decode on a Raspberry Pi 5 runs at 400 to 4,000 tokens per second across the depth ladder, and the Hadamard MLP is a large part of why the number has that many digits.

What it gives up

A dense feed-forward layer can represent any linear map on its hidden width; three Kronecker mixes with permutations cannot represent all of them, and the rank-8 gain is a narrow substitute for the input-dependence of a wide hidden layer. On the reasoning-dense, tool-shaped data Needle trains on, where the block's job is to route and transform what is already in context, that expressivity was not where the loss lived, and the engram absorbs the storage role. Whether the same trade holds for knowledge-heavy general text at larger scale is the open question, and the same one the attention-only study left on the table. Within Needle's regime the numbers above are the argument: a 121M model that beats models ten times its size on mobile tool calls, spending a third of the compute a transformer of its shape would.