Provedunbiased-attentionrandomized-algorithmssoftmaxeasy-side

The softmax score-gradient is a covariance: Σ π_j (u_j − E_π u) k_j = E_π[u k] − E_π[u] E_π[k]

For all $\htmlData{mx=1}{n \in \htmlData{mx=0}{\mathbb{N}}}$ and $\htmlData{mx=5}{\pi, u, k : \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{n}} \mathrel{\htmlData{mx=4}{\to}} \htmlData{mx=3}{\mathbb{R}}}$ , if $\htmlData{mx=3}{\sum_{j} \htmlData{mx=2}{\htmlData{mx=0}{\pi}(\htmlData{mx=1}{j})}} \mathrel{\htmlData{mx=4}{=}} 1$ , then

\begin{aligned} & \htmlData{mx=21}{\sum_{j} \htmlData{mx=20}{\htmlData{mx=16}{\htmlData{mx=3}{\htmlData{mx=1}{\pi}(\htmlData{mx=2}{j})} \cdot \left(\htmlData{mx=15}{\htmlData{mx=6}{\htmlData{mx=4}{u}(\htmlData{mx=5}{j})} - \htmlData{mx=14}{\sum_{i} \htmlData{mx=13}{\htmlData{mx=9}{\htmlData{mx=7}{\pi}(\htmlData{mx=8}{i})} \cdot \htmlData{mx=12}{\htmlData{mx=10}{u}(\htmlData{mx=11}{i})}}}}\right)} \cdot \htmlData{mx=19}{\htmlData{mx=17}{k}(\htmlData{mx=18}{j})}}} \\ & \quad \htmlData{mx=0}{=} \htmlData{mx=51}{\htmlData{mx=33}{\sum_{j} \htmlData{mx=32}{\htmlData{mx=24}{\htmlData{mx=22}{\pi}(\htmlData{mx=23}{j})} \cdot (\htmlData{mx=31}{\htmlData{mx=27}{\htmlData{mx=25}{u}(\htmlData{mx=26}{j})} \cdot \htmlData{mx=30}{\htmlData{mx=28}{k}(\htmlData{mx=29}{j})}})}} - \htmlData{mx=50}{\left(\htmlData{mx=41}{\sum_{j} \htmlData{mx=40}{\htmlData{mx=36}{\htmlData{mx=34}{\pi}(\htmlData{mx=35}{j})} \cdot \htmlData{mx=39}{\htmlData{mx=37}{u}(\htmlData{mx=38}{j})}}}\right) \cdot \left(\htmlData{mx=49}{\sum_{j} \htmlData{mx=48}{\htmlData{mx=44}{\htmlData{mx=42}{\pi}(\htmlData{mx=43}{j})} \cdot \htmlData{mx=47}{\htmlData{mx=45}{k}(\htmlData{mx=46}{j})}}}\right)}} \end{aligned}

What the notation means · or point at any symbol

NNatural numbers: The natural numbers 0, 1, 2, …. In Lean they start at 0. (ℕ)
Fin: Fin: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (Fin)
f(x)Function application: Fin applied to n. Lean writes f x; papers write f(x). (Fin n)
⟹Implies: If Fin(n), then ℝ. (Between types, A → B is the type of functions from A to B.) (Fin n → ℝ)
x : TVariables: π, u, k range over Fin n → ℝ: the statement must hold for every choice. ((π u k : Fin n → ℝ))
∑Summation: Add up π(j) for every j. (∑ j, π j)
=Equality: ∑ j, π(j) and 1 are the same. (∑ j, π j = 1)
πPi: The ratio of a circle's circumference to its diameter, ≈ 3.14159. (π)
RReal numbers: The real numbers: every point on the number line. (ℝ)
-Subtraction: u(j) minus ∑ i, π(i) · u(i). On ℕ, Lean's subtraction stops at 0: 2 − 5 = 0. (u j - ∑ i, π i * u i)
·Multiplication: π(i) times u(i). (π i * u i)

theorem ua-gradient-is-covariance {n : ℕ} (π u k : Fin n → ℝ) (hπ : ∑ j, π j = 1) :
    ∑ j, π j * (u j - ∑ i, π i * u i) * k j = ∑ j, π j * (u j * k j) - (∑ j, π j * u j) * (∑ j, π j * k j)

With u_j = g·v_j (upstream gradient times value), the attention gradient with respect to the query is Σ_j π_j (u_j − E u) k_j, a covariance under the softmax row. It is a product of softmax expectations, so the forward estimators extend to gradients. Part of 'Unbiased randomized attention: cost x variance thresholds', round 6 (gradients).

How it was proved

Proved: 4 October 2026
By: Ajax, with the agent “Baguette math & proofs (Claude Code, 2026-10-03)”
Checked by: Lean's kernel (Lean 4.33.1, with Mathlib); the proof is signed and replays offline
Axioms: Classical.choice, Quot.sound, propext
Proof: 12 lines of Lean

The softmax score-gradient is a covariance: Σ π_j (u_j − E_π u) k_j = E_π[u k] − E_π[u] E_π[k]

How it was proved

More proved statements