# Softmax over one token has zero gradient, and Adam moves the weights anyway

> With a single token, attention's query and key weights cannot change the output, so their true gradient is zero. A fused kernel returned tiny nonzero gradients instead, and when they outgrew Adam's epsilon the weights drifted until the attention saturated.

- Author: Ajish Pradeep (https://ajishpradeep.com/about/)
- Published: 2026-10-07
- Area: numerics
- Canonical: https://ajishpradeep.com/blog/softmax-over-one-token-has-zero-gradient-and-adam-moves-the-weights-anyway/
- Cite as: Pradeep, A. (2026). Softmax over one token has zero gradient, and Adam moves the weights anyway. https://ajishpradeep.com/blog/softmax-over-one-token-has-zero-gradient-and-adam-moves-the-weights-anyway/

Attention blocks sometimes see sequences of length one. A model may be trained on single items before it is trained on sequences, a cache may hold one entry, or a block built for many tokens may be run in a mode that feeds it one. In that case the block's query and key projections have nothing to do. I wanted to know what actually happens to them during training, because "nothing" is what the mathematics says, and I wasn't sure the hardware agrees.

## On paper, the gradient is exactly zero

Attention, as Vaswani et al. (2017) defined it, takes its weights from a softmax over the scores $s_j = q^\top k_j / \sqrt{d}$:

$$p_i = \frac{e^{s_i}}{\sum_j e^{s_j}}, \qquad \frac{\partial p_i}{\partial s_j} = p_i(\delta_{ij} - p_j)$$

With one token there is one score and one weight, $p_1 = 1$, whatever the score is. The Jacobian is $p_1(1 - p_1) = 0$. The output is $p_1 v_1 = v_1$, which depends on the value projection and not at all on the query or the key. By the chain rule every gradient that has to pass through the score is multiplied by that zero, so the query and key weights receive exactly no gradient. An optimizer that only moves weights in proportion to their gradient should leave them where they started.

## On hardware, it depends on the kernel

I checked with one attention block, 64 dimensions and a batch of 256 single-token sequences, and computed the gradient into the query, key and value weights four ways: a hand-written softmax and PyTorch's fused attention function, each on CPU and GPU, in float32 and with bfloat16 autocast.

| Implementation | Device | Precision | Largest query gradient | Largest key gradient | Largest value gradient |
|---|---|---|---:|---:|---:|
| Hand-written softmax | CPU and GPU | float32 and bfloat16 | 0 | 0 | about 0.011 |
| Fused attention | CPU | float32 | $7 \times 10^{-10}$ | $7 \times 10^{-10}$ | 0.011 |
| Fused attention | CPU | bfloat16 | $5 \times 10^{-10}$ | $5 \times 10^{-10}$ | 0.011 |
| Fused attention | GPU | float32 | $1 \times 10^{-8}$ | $1 \times 10^{-8}$ | 0.012 |
| Fused attention | GPU | bfloat16 | $2 \times 10^{-9}$ | $2 \times 10^{-9}$ | 0.012 |

I had expected low precision to be the culprit. It wasn't. The hand-written softmax returned exact zeros even in bfloat16. Every fused path returned small nonzero gradients, including float32 on the CPU. Fused kernels in the style of FlashAttention (Dao et al., 2022) avoid storing the attention matrix and rebuild what the backward pass needs from other quantities. My reading, which I haven't traced through the kernel source, is that with one token the two routes to what should be the same number round differently, and their difference is what arrives as a gradient. It is about eight orders of magnitude smaller than the real gradient into the value weights.

## What Adam does with a gradient of $10^{-9}$

SGD multiplies a gradient by the learning rate, so a $10^{-9}$ gradient moves a weight by nothing measurable. Adam normalises each coordinate by a running estimate of its gradient's size:

$$u_i = \eta \, \frac{\hat m_i}{\sqrt{\hat v_i} + \epsilon}$$

Kingma and Ba (2015) designed it to be "invariant to diagonal rescaling of the gradients": if every gradient is multiplied by a constant, the update doesn't change. For a gradient made of round-off, that invariance is the problem. Whether the round-off is $10^{-9}$ or $10^{-3}$, the normalised step is about $\eta$, as long as the gradient is large compared with $\epsilon$. The only thing standing between a meaningless gradient and a full-size step is $\epsilon$, whose default is $10^{-8}$.

That turns the question into a comparison: how large is the round-off gradient, relative to $\epsilon$? The round-off scales with the gradient flowing into the block from above, so I varied that directly by multiplying the loss by a constant. Under Adam this changes nothing for the parameters with real gradients. It only changes how large the round-off is next to $\epsilon$.

I trained the single-token block with the fused kernel and bfloat16 autocast, with Adam at learning rate $10^{-3}$, to fit a fixed linear map.

| Loss scale | Steps | Mean round-off gradient on query weights | Spread of query weights vs initial | Relative change of query weights |
|---|---:|---:|---:|---:|
| 1 | 5,000 | $6 \times 10^{-12}$ | 1.00× | 1.3% |
| 100 | 5,000 | $6 \times 10^{-10}$ | 1.27× | 59% |
| 10,000 | 5,000 | $8 \times 10^{-8}$ | 2.56× | 199% |
| 10,000 | 50,000 | $4 \times 10^{-7}$ | 61× | 6,100% |
| 10,000, $\epsilon = 10^{-5}$ | 50,000 | $2 \times 10^{-8}$ | 1.02× | 15% |

At the natural loss scale the round-off sits far below $\epsilon$ and the weights barely move. Raise the gradient reaching the block and the round-off crosses $\epsilon$; the weights start taking real steps. Over a long run the drift feeds itself. As the query and key weights grow, the scores they produce grow, the round-off grows with them, and the mean round-off gradient rose from about $8 \times 10^{-8}$ to $4 \times 10^{-7}$ while the spread of the weights reached 61 times its initial value. Throughout, the block's output was untouched: multiplying the query weights by five changed it by exactly zero. SGD, run at the natural scale for the same 5,000 steps, moved them not at all.

## The drift is harmless until there is a second token

Weights that cannot affect the output are harmless while they still cannot. The trouble starts when the block meets longer sequences. I took the drifted query and key weights and asked the block to solve a 16-token retrieval task: the first token carries a key, one of the other fifteen carries the same key and a value, and the output must be that value. A fresh block learns this in a few hundred steps.

The drifted block never did. Over 3,000 steps its loss stayed around 1.0, the variance of the target, which is what a model that has learned nothing scores. A fresh block reached 0.0005. Looking at the attention explained why:

| Block | Spread of the scores | Attention entropy (16 tokens, max 2.77 nats) | Average weight on the top token |
|---|---:|---:|---:|
| Fresh | 0.06 | 2.77 nats | 0.07 |
| After the long drift | 874 | 0.006 nats | 0.998 |

Scores with a spread of 874 put essentially all the weight on one token, and which token is decided by the drifted weights, not by the keys. The softmax Jacobian $p_i(\delta_{ij} - p_j)$ is close to zero everywhere when one $p$ is close to 1 and the rest are close to 0, so the gradient that should teach the block where to look barely reaches the query and key weights. The block is stuck in a saturated softmax it walked into while it could not see anything.

*(An interactive figure is on the web page; every number it shows is stated in the text.)*

## What stops it

Three things stopped the drift in my runs. A larger $\epsilon$ did: at $10^{-5}$ the round-off stayed below the floor, and after 50,000 steps the weights' spread was 1.02 times the initial value. The unfused softmax did, because it returns exact zeros and zero times anything stays zero. And SGD did, because it never normalises, so a $10^{-9}$ gradient produces a $10^{-11}$ step.

The fix I would reach for first is simpler than any of these and I haven't tested it: if a block will see single tokens for a while, stop training the parameters that cannot matter, by switching their gradients off for that period. Weight decay would also pull the drifted weights back towards zero, but I don't know whether it would win against Adam's full-size steps.

## What I would watch

The loss can't show this, because the output doesn't depend on the drifting weights. These would:

- the ratio of update size to weight size, per parameter group, for every group that should not be learning;
- the size of the gradient each group receives, next to the optimizer's $\epsilon$;
- the entropy of the attention weights on a fixed batch of longer sequences, logged even while training on short ones.

Adam's indifference to the size of a gradient shows up in two other places I measured: its first step after a reset, in [restarting Adam moves every weight by the full learning rate](https://ajishpradeep.com/blog/restarting-adam-moves-every-weight-by-the-full-learning-rate/), and its reaction to one rare spike, in [gradient clipping does different jobs under SGD and Adam](https://ajishpradeep.com/blog/gradient-clipping-does-different-jobs-under-sgd-and-adam/).

## What I still don't know

Multiplying the loss by 10,000 is my stand-in for whatever makes the gradient into a block large in a real model: a large output, a summed loss, a deep stack above it. I haven't measured how large that gradient usually is in practice, so I can't say how common the dangerous regime is. I also tested only PyTorch's fused function, on one GPU and one CPU. Other backends may round to exact zero, or leak more.

The feedback loop interests me most. The round-off grew as the weights grew, which is why the long run went from a 2.6-fold to a 61-fold spread instead of growing steadily. I expected a random walk, and a random walk whose step size grows with the distance travelled is a different object. I'd like to understand its growth rate.

## References

- Adam: A Method for Stochastic Optimization — Diederik P. Kingma and Jimmy Ba, ICLR 2015 — https://arxiv.org/abs/1412.6980
- Attention Is All You Need — Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser and Illia Polosukhin, NeurIPS 2017 — https://arxiv.org/abs/1706.03762
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré, NeurIPS 2022 — https://arxiv.org/abs/2205.14135