Softmax is one of those operators you carry around without thinking about, like a paperclip. It does its small job: it turns a vector of real numbers into something that sums to one and is everywhere differentiable. The cost is hidden inside the exponential, which makes the numerics fussy and makes the kernel writers earn their pay.
A few weeks ago I went back to a question I’d half-asked myself in 2023: how much of what attention does is the softmax, and how much is just the rest of it — the dot products, the projections, the residual? The answer turns out to be more interesting than I expected, but mostly because the question is wrong.
The setup
I took a small decoder-only transformer (110M params, the kind that fits on a single 4090) and replaced the softmax inside each attention head with a family of simpler maps:
- The identity over a clamped range,
- An L1 normalization of the absolute values,
- A ReLU + L1 normalization,
- A learned per-head temperature with a sigmoid envelope.
Everything else stayed put. Same tokenizer, same training data, same optimizer. The point wasn’t to invent a new architecture — it was to push on a single variable and see what wobbled.
What breaks
The identity map breaks immediately and obviously: the variance of the attention scores compounds through the depth of the network until the residual stream is just noise. This is the boring failure case. It tells you that some normalization is doing real work; it does not tell you what.
The L1-of-absolute-values variant is more interesting. It trains. It even trains reasonably well on the first 8B tokens. Then, somewhere around step 10,000, a few heads collapse into a state where every key attends almost equally to every value. Once you see this in one layer, you tend to see it in the next two layers up.
What softmax is buying you, mostly, is the right to have heads with strong opinions. Take it away and you do not get a broken network; you get a well-mannered, uninteresting one.
What holds
The ReLU + L1 variant holds up the best. It loses about 2.4% on
held-out perplexity vs. softmax, but it’s stable across seeds and it
trains faster per step on the kernels I wrote (no exp, no max
subtraction trick, just a comparison and a sum). For inference on cheap
hardware, that’s an honest trade.
The learned-temperature variant tied softmax on perplexity, which surprised me. Looking at the temperatures after training, the model had landed on something close to a soft-max with a per-head bias term — it had rediscovered the operator I’d tried to replace, basically. Sometimes the right answer to “can we remove this?” is “you can, but the gradient will put it back.”
The one place the trick stops being free
On a single GPU running batch size 1, all four variants come out roughly even on wall-clock. The reason is unromantic: at small batch sizes the attention operator isn’t the bottleneck. The matmuls are.
Crank batch size up to 64, push the context to 16K tokens, and the picture
changes. The ReLU+L1 kernel beats a tuned softmax kernel by about
1.6× on memory bandwidth, because the absence of exp means
you can fuse the normalization into the QKTV pipeline without
spilling. This is the place the trick stops being free for the people
on the other side of it: the inference-cost line on a P&L.
What I took away
Three things, in no particular order:
- Softmax is not load-bearing in the way I thought. It’s load-bearing for character — the ability of a head to be opinionated — and that turns out to matter for quality but not for shape.
- The cleanest substitute (ReLU + L1) is well-known. The literature on this is decent and largely ignored, including by me.
- The reason “just remove softmax” isn’t a free lunch is that the math of attention is the easy part. The hard part is the bookkeeping that lets a 16K-context kernel run at half its theoretical throughput, and softmax happens to be near the center of that bookkeeping.
If you’ve done this experiment yourself, I’d like to hear what your seeds did. The code, such as it is, is up on GitHub. It is not pretty.