How Much Softmax Do You Need in LLMs?

Papers: Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration (inference) · Pretraining Transformers with Quantized Softmax in Attention (training)

Authors: Shangzhen Zhu, Muyan Hu, Tomasz Kozlowski · University of Illinois Urbana-Champaign · contact: szhu48@illinois.edu

Headlines

1. In a pretrained LLM, softmax can drop a large share of positions and round scores coarsely, but the weights have to keep a roughly exponential shape.
2. When the resolution budget is fixed, spend it near the row maximum.
3. When you train with an approximate softmax, the feedback rule through the approximation matters as much as the approximation itself.

In a transformer, each attention layer lets a token gather information from earlier positions, mixing them with a set of weights that sum to one. Softmax is the step that turns the raw attention scores into those weights. This post asks how much of that step can be simplified. We measure the cost with perplexity, which tracks how well the model predicts the next token (lower is better).

Watch on YouTube

The video above is the short version. This post adds the experimental settings, the full numbers, the kernel’s bit-level construction and the training-side backward rules. Collapsed Details blocks hold the material most readers can skip. Result plots are drawn from the recorded outputs; illustrations and schematics are labeled as such.

Why care about the exponential in softmax?

Attention is

where the mask \(M\) is 0 for allowed positions and \(-\infty\) for future ones. For one query, the row of scores \(s_j\) becomes weights

That is one exponential for every valid score. For Qwen2.5-1.5B (28 layers, 12 query heads), one causal forward pass over 2,048 tokens has

per-score exponentials. This number is counted from the valid positions of the causal mask. It is not a measured instruction count.

The hardware makes this count matter. On an NVIDIA B200, each streaming multiprocessor (SM) can do 8,192 FP16/BF16 matrix-multiply operations per clock, but its special-function unit (MUFU) produces only 16 exponentials per clock. That is a 512× gap. FlashAttention-4 (FA4) already works around it: it computes 10–25% of each row’s exp2 calls with a cubic polynomial on the regular floating-point pipeline instead of the MUFU.

So the question is whether a trained model needs every one of those exponentials at full precision. We split it into three questions about one frozen model, then build a kernel from the answers, then ask what changes in training:

1. How many positions does it need?
2. How precise must the scores be?
3. Must the weights be exponential?

Unless noted, inference numbers come from Qwen2.5-1.5B with all parameters frozen. We report next-token perplexity on about 200K tokens of the WikiText-103 test set, relative to the same model with softmax. Cross-model checks use ten decoder-only models from 0.5B to 72B.

Question 1: How many positions does it need?

Start with one row. Figure 3 shows a real row (layer 16 of 28, head 5, query token “ to”), sorted by weight. In this row, 3 of the 2,048 positions carry 54.5% of the weight, and the smallest 1,536 carry 0.8%.

The whole model shows the same picture. Over all layers, heads and rows of the evaluation text, positions with \(p\ge 0.1\) are 0.15% of valid positions and carry 54% of the attention weight. Positions with \(p<10^{-4}\) are 75% of positions and carry 1.2% in total.

So we drop the smallest weights. In every row of every layer we keep only the largest fraction of valid positions, zero the rest and renormalize. Perplexity changes relative to softmax:

Positions can be cut a lot: keeping half costs almost nothing. But the cut has a limit: at a tenth of the row the damage becomes visible, and below that it grows fast.

Question 2: How precise must the scores be?

This time we keep every position and round the scores instead.

Start with seven scores and a coarse grid (Figure 4). Each score moves to its nearest grid value, so different scores end up sharing a value. Scores on the same grid value also share one exponential: a row needs one exponential per grid value it uses, not one per score. For a row of 2,048 scores, a 16-step grid has at most 17 distinct exponential values.

Terms. For each row, we cut the score range from the row minimum to the row maximum into \(K\) steps, which gives \(K+1\) grid values. Each score is rounded to the nearest grid value, and the weights stay exponential in the rounded score. Step widths shrink geometrically toward the maximum. The ratio

sets the allocation: \(R=1\) gives uniform steps, and \(R>1\) makes steps finer near the maximum.

With \(R=4\), perplexity rises by 0.28% at \(K=16\), 0.09% at \(K=32\) and 0.007% at \(K=64\), all relative to softmax. In other words, each score only needs to land on one of a few dozen values. That is a statement about grid values, not storage: the paper reports \(\log_2(K+1)\) only as a descriptive index budget and says explicitly that it is not a storage specification. Across all ten models, the \(K=32, R=4\) grid stays below a 1% perplexity increase in every point estimate. The largest is Gemma-2-2B at 0.85%, whose 95% CI spans 1%.

Where the steps go matters more

Now fix the number of steps at \(K=21\) and change only \(R\), from ¼ to 8:

Pushing steps toward the bottom of the row makes things worse. Squeezing them toward the row maximum makes things better, and the improvement is monotonic over the tested range.

The surprising part is what happens to the error. From \(R=1\) to \(R=4\):

• the unweighted error (root-mean-square error of the log weights, every position counted equally) rises by 28%;
• the attention-weighted error (each position weighted by its softmax probability) falls by 29%;
• the extra loss caused by the approximation falls by 72%.

Measured uniformly over positions, the \(R=4\) grid is the worse approximation. The model says it is the better one. The effect holds across models: the extra loss is lower at \(R=4\) than at \(R=1\) in all ten point estimates, and the 95% CI excludes zero in nine of them (Figure 6 in the Details below). Gemma-2-2B is unresolved.

The reason is Question 1: attention is concentrated near the maximum, so resolution spent there does far more than resolution spread evenly. Put the limited resolution near the row maximum.

Question 3: Must the weights be exponential?

Now we change how a score is converted into a weight.

Start with the extreme case. Keep the same top half of each row, but give every kept position the same weight. Perplexity is multiplied by more than 5,000 (×5,268). Because that intervention runs in every layer, it also changes the scores that later layers see, so a cleaner control changes one layer at a time. Both conditions reach that layer with identical scores and keep the same positions; they differ only in equal versus softmax weights. Equal weights make the model worse in every one of its 28 layers, and each layer’s 95% CI excludes zero.

Weights that grow along a straight line are almost as bad. With weight proportional to the grid index, perplexity is multiplied by 4,172 at 32 steps, and by about 4,100 to 5,200 at every grid size tested (4 to 1,024 steps). Exponential weights on the same 32-step grid cost about 1%.

You can cut positions, but you can’t lose the shape.

How precise does the shape need to be?

Rewrite the exponential in base 2:

This step is exact, and FlashAttention-4 already computes softmax this way, with exp2. Measure each score in octaves below the row maximum,

so that one step down in \(x\) halves the weight. That step is one octave. The approximation is to put a limited number of levels in each octave.

• 1 level per octave, \(w=2^{\mathrm{round}(x)}\): +0.30%.
• 2 levels per octave, \(w\in\{1,\,1.5\}\times 2^{n}\) (this is Rowmax-H15, below): +0.105%.
• 4 levels per octave, \(\{1,\,1.25,\,1.5,\,1.75\}\times 2^{n}\): +0.033%.
• 8 levels per octave: +0.008%.
• Linear interpolation (no rounding), the curve \(S(x)=2^{\lfloor x\rfloor}\,(1+x-\lfloor x\rfloor)\) through the powers of two: −0.0023%. Its 95% CI includes zero. It rounds nothing, but it is still an approximation of \(2^x\).

One detail matters here. The in-between levels lie on straight lines between neighboring powers of two. The half-octave level is therefore \(1.5\times 2^{n}\), not \(\sqrt2\times 2^{n}\).

Why integer exponents are attractive. A power of two with an integer exponent does not need to be computed. A floating-point number already stores an exponent field, so you write the integer into it: one addition and one shift. A half-octave level sets one extra bit, the leading mantissa bit, which means “× 1.5”. Schraudolph (1999) used the same idea for a fast, compact approximation of the exponential. Writing a fraction into the mantissa as well gives exactly the linear interpolation \(S(x)\) above.

The reference point matters too. Take the same octave grid, but place it on a fixed absolute scale: the fixed \(\log_2\) lattice used by post-training quantizers such as FQ-ViT, here at 8 bits. The row maximum then falls between grid values, with a different phase in every row. Perplexity rises by 0.66% instead of 0.30%, so the extra loss is 2.17× larger than with the grid anchored at each row’s maximum.

Answer to Question 3. In these frozen-model tests, the weights need a roughly exponential shape, and anchoring the grid at the row maximum reduces the loss. Two levels per octave already come close.

Putting it together: Rowmax-H15

The grid experiments measure what the model tolerates; turning that into speed needs an implementation. The kernel replaces per-score exponentials only; it does not drop positions as in Question 1.

The answers give four design rules:

1. Keep the exponential shape (Question 3).
2. Anchor to the running row maximum that online softmax already keeps (Question 2 and the anchor control). There is no scale to calibrate.
3. Use two levels per octave, \(\{1, 1.5\}\times 2^n\).
4. Replace each per-score exponential with an add and a shift. No calibration parameters are needed.

We built this into FlashAttention-4 and call it Rowmax-H15. For the base-2 exponent \(x\) that FA4’s online exponential path receives (measured from the running maximum),

so even \(n\) gives \(2^k\) and odd \(n\) gives \(1.5\cdot 2^k\) (RNE is round-to-nearest-even).

What is replaced, and what is not. Only the per-score exponential changes. The running maximum, the rescaling of the accumulators, the log-sum-exp and the pipeline are unchanged, and the row sum and the \(PV\) product use the same approximated weights. FA4 rescales lazily: it skips the rescale while the row maximum grows by less than a threshold (\(\tau=8.0\)). The anchor can therefore lag the final row maximum, which makes \(x\) positive; the formula covers both signs.

For Qwen2.5-1.5B, the 705 million per-score exponentials of a 2K-token forward pass become zero. Each score instead costs three cheap integer and floating-point instructions. That count comes from a static instruction audit of isolated micro-kernels (3 instructions against 8 for FA4’s cubic emulation), not from the B200 kernel itself.

Accuracy was measured inside the kernel, on the BF16 path at 2K context, over the same ~200K WikiText-103 tokens, against stock FA4:

• Qwen2.5-1.5B: perplexity +0.121% (95% CI 0.086–0.156%).
• Five models from three families: Mistral-7B-v0.3 +0.091%, Qwen2.5-72B +0.101%, Qwen2.5-1.5B +0.121%, Llama-3.1-8B +0.246%, Llama-3.1-70B +0.492%. All five 95% CIs lie above zero, so the loss is small but measurable. The five models do not establish a trend with size.

Speed was measured on one B200 with FP8 inputs at 8K tokens, using the standard shape of batch 1, 16 heads and head dimension 128 (not Qwen2.5-1.5B’s own shape). The metric is attention-forward call latency, the pooled median over three rounds, and “faster” means stock latency divided by patched latency, minus one:

• causal: 0.2063 → 0.1835 ms, 12.4% faster;
• non-causal: 0.4879 → 0.3879 ms, 25.8% faster.

These speedups are for the attention forward computation, not the whole language model. Note also that accuracy (BF16, 2K) and speed (FP8, 8K) come from different settings. Kernel accuracy was not measured at FP8 or at 8K.

# y = 2x: base-2 exponent of this score, measured from the running row max, already doubled
y = max(y, -252.0)                  # keep the result a normal float: n >= -252 gives 2^-126 or larger
                                    # (n = -254 gives 0, n = -255 wraps to NaN)
t = y + from_bits(0x4B400000)       # bit pattern of 2^23 + 2^22 = 12582912.0
                                    # 1) FP32 add: rounds to n = round_RNE(2x),
                                    #    which now sits in the low mantissa bits of t
b = bits(t) << 22                   # 2) integer shift: bit 0 of n -> leading mantissa bit (x1.5),
                                    #    the rest -> exponent field (the 0x4B4... bits shift out)
w = from_bits(b + 0x3F800000)       # 3) integer add of the bit pattern of 1.0 supplies the bias
# w == (1 + 0.5*(n % 2)) * 2**(n // 2)

Scalar distortion metrics can mislead

A common way to judge an approximation is a single distortion number, such as a pointwise error bound or the divergence between the approximate and the softmax attention distributions. On Qwen2.5-1.5B such a number does track damage: across 56 interventions, log extra loss against log mean row-wise Jensen–Shannon divergence (JSD) has slope 1.05 and \(R^2=0.945\) (in-sample). But it does not determine the damage. Flattening the scores (scaling them by \(\alpha<1\)) and sharpening them (\(\alpha>1\)), tuned to the same mean JSD, cost different amounts, and which one hurts more depends on the model: flattening in seven of ten models, sharpening in two (Figure 12). Question 2 showed the same thing: the grid with more unweighted error had less extra loss. So approximations are judged here by model loss, including the kernel above, whose per-element error is far above one BF16 unit in the last place.

What changes when you train with an approximate softmax?

The training study uses related K-interval grid operators, not the Rowmax-H15 kernel itself.

At inference time, an approximation only perturbs a function that has already been trained. Training is different. It also works out how each parameter should change, and that feedback (gradient) has to pass through the approximation too.

The training paper puts a grid like the one in Question 2 into every attention layer of a GPT-2-style model trained from scratch. Each row gets \(K\) equal steps and \(K+1\) exponential grid values, simulated in fp32.

Same forward, different feedback

Take two copies of the same model with the same grid inside attention. In these runs the grid spans each row from its lowest to its highest score, and weights are linearly interpolated between neighboring grid values (LERP, \(K=32\)). The setup is a 124M model on WikiText-103 for 100M tokens. The forward rule is identical: given the same parameters, both compute exactly the same output.

• Model A also sends feedback through the grid’s ends, because moving the lowest or highest score moves the whole grid.
• Model B ignores that effect. It calls .detach() on the row minimum and maximum, which cuts the calibration derivatives: the feedback that reaches the scores through the grid’s ends. The ends’ own scores still receive their direct feedback; only this calibration path is cut.

Initially, the two train almost identically. After about 25–30 million tokens they separate. At 100M tokens, Model B’s validation perplexity is 2.09× Model A’s, while Model A shows no detectable difference from softmax.

What Model B does has a name: detach. For plain softmax it is harmless, because subtracting the maximum does not change the result and its feedback cancels exactly. So people detach the maximum out of habit. With a grid, the ends set the grid itself, so that feedback must stay.

The result holds up in the paper’s checks. It replicates on two more seeds. With hard rounding instead of LERP, detach also fails against its matched full-feedback run at every \(K\) tested. Forcing Model B’s feedback to sum to zero, the identity that shift invariance implies, does not fix it; that run ends worse. Restoring only the row maximum’s calibration path recovers about 98% of the gap.

Three choices

That was one of three choices that decide whether training works:

1. Feedback through the grid’s ends: keep it; do not detach.
2. Where the grid’s ends come from: the row minimum and maximum (“MinMax”), or a fixed-width window of 6 nats below the maximum, with zero weight below it (“fixed window”). Under MinMax, each step is (max − min)/K wide. Whenever the lowest score drifts further down, every grid value is pulled apart, including those near the top, where most of the weight is. With a fixed window, the lowest score can keep drifting, but the grid near the top stays where it is.
3. The feedback rule for rounding: rounding turns smooth changes into steps. Nudge a score a little and the rounded value does not move, so rounding gives no feedback by itself. Training needs a hand-chosen surrogate: the forward pass uses the rounded weights, and the backward pass uses the interpolated ones. The question is where the surrogate sits, before normalization (on the unnormalized weights) or after it (on the probabilities). In these tests it worked much better after normalization.

A longer run

The next results come from a separate setup. It is longer than the A/B experiment: GPT-2 124M on FineWebEdu for 2.5B training tokens, with hard rounding and only \(K=4\) steps per row. All of these runs keep the feedback through the grid’s ends (choice 1). They start from MinMax calibration with the surrogate before normalization, then change the surrogate, the calibration, or both, all at \(K=4\) (Figure 15).

With MinMax calibration and the pre-normalization surrogate, validation perplexity ends at 2.4× softmax’s (+144%), and the model also loses 3.9 to 19.4 accuracy points on downstream tasks (six benchmarks in seven configurations). It does something specific: it stretches its score range to about seven times the softmax model’s, and the step width reaches 53 nats. It sharpens scores that the grid cannot even resolve (Figure 16). Changing either choice removes most of the gap, and changing both leaves perplexity about 2% above softmax.

Shorter checks (two separate experiments, both shorter than the 2.5B-token run):

• A 1B-parameter model (32 layers, width 1536) on WikiText-103 for 100M tokens, one seed. The paper calls it a probe, not a scale experiment. At \(K=4\), the calibration effect, the surrogate effect and their interaction keep their direction.
• The 124M model on WikiText-103 for 100M tokens with five seeds, over a 5 × 5 grid of methods and \(K\). The ordering of conditions is reproducible across seeds (Kendall’s \(W=0.91\)), and the three effects keep their sign on every seed.

What was not measured. These runs simulate the grid in fp32 attention arithmetic. No training kernel was built, and training speed was not measured.

The common thread with inference: keep the limited resolution near the top scores, which matter most. In training, the feedback rule has to be right too.

Related work and positioning

Citation

If you find this work useful, please cite:

@misc{zhu2026approximating,
  title  = {Approximating Softmax in Pretrained {LLMs}: Model Sensitivity and Kernel Acceleration},
  author = {Zhu, Shangzhen and Hu, Muyan and Kozlowski, Tomasz},
  year   = {2026},
  eprint = {2609.33586},
  archivePrefix = {arXiv}
}
@misc{zhu2026pretraining,
  title  = {Pretraining Transformers with Quantized Softmax in Attention},
  author = {Zhu, Shangzhen and Hu, Muyan and Kozlowski, Tomasz},
  year   = {2026},
  eprint = {2609.33591},
  archivePrefix = {arXiv}
}

Video: How Much Softmax Do LLMs Need?