The softmax Jacobian as a variance, and saturation

ProofIntermediate⏱ 30 minDeep Learningsoftmaxjacobiansaturationgradient

Theory: The Mathematics of Transformers →

Problem

Prove $v^\top(\diag(a)-aa^\top)v = \Var_{j\sim a}(v_j)$ and deduce that the softmax Jacobian's operator norm is at most $\max_ja_j(1-a_j)\le1/4$ , vanishing as $a$ becomes one-hot. Relate to temperature and to attention saturation.

Write a complete argument. Every inequality you use should be one you could state precisely.

Solution

Opening this before you have a real attempt on paper is the fastest way to feel like you understand something you cannot reproduce.

Related problems

All problems

Problem	Field	Type	Difficulty	Time
The cross-entropy gradient is the prediction error	LLMs	Proof	Beginner	15 min
When does attention dominate the MLP?	Deep Learning	Numerical	Beginner	30 min
Attention is permutation equivariant	Deep Learning	Proof	Intermediate	30 min
Gradient magnitude through a deep chain	Deep Learning	Proof	Intermediate	30 min
Implement multi-head attention from scratch run	Deep Learning	Coding	Intermediate	3 hours