Provedunbiased-attentiongradientone-passtraining

One-pass attention gradient: from N ≥ 2 exact samples, (1/N)Σ U_i K_i − (1/(N(N−1)))Σ_{i≠j} U_i K_j is unbiased for Cov(U, K)

For all $\htmlData{mx=1}{\Omega \in \htmlData{mx=0}{\operatorname{Type}}}$ , $\htmlData{mx=1}{N \in \htmlData{mx=0}{\mathbb{N}}}$ and $\htmlData{mx=3}{w, U, K : \htmlData{mx=0}{\Omega} \mathrel{\htmlData{mx=2}{\to}} \htmlData{mx=1}{\mathbb{R}}}$ , if $\htmlData{mx=3}{\sum_{\omega} \htmlData{mx=2}{\htmlData{mx=0}{w}(\htmlData{mx=1}{\omega})}} \mathrel{\htmlData{mx=4}{=}} 1$ and $2 \mathrel{\htmlData{mx=1}{\leq}} \htmlData{mx=0}{N}$ (given $\htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fintype}}\,\htmlData{mx=1}{\Omega}}$ ), then

\begin{aligned} & \htmlData{mx=53}{\sum_{f : \htmlData{mx=3}{\htmlData{mx=1}{\operatorname{Fin}}\,\htmlData{mx=2}{N}} \mathrel{\htmlData{mx=5}{\to}} \htmlData{mx=4}{\Omega}} \htmlData{mx=52}{\left(\htmlData{mx=11}{\prod_{m} \htmlData{mx=10}{\htmlData{mx=6}{w}(\htmlData{mx=9}{\htmlData{mx=7}{f}(\htmlData{mx=8}{m})})}}\right) \cdot \left(\htmlData{mx=51}{\htmlData{mx=27}{\htmlData{mx=14}{\frac{1}{\htmlData{mx=13}{\htmlData{mx=12}{N}}}} \cdot \htmlData{mx=26}{\sum_{i} \htmlData{mx=25}{\htmlData{mx=19}{\htmlData{mx=15}{U}(\htmlData{mx=18}{\htmlData{mx=16}{f}(\htmlData{mx=17}{i})})} \cdot \htmlData{mx=24}{\htmlData{mx=20}{K}(\htmlData{mx=23}{\htmlData{mx=21}{f}(\htmlData{mx=22}{i})})}}}} - \htmlData{mx=50}{\htmlData{mx=33}{\frac{1}{\htmlData{mx=32}{\htmlData{mx=29}{\htmlData{mx=28}{N}} \cdot (\htmlData{mx=31}{\htmlData{mx=30}{N} - 1})}}} \cdot \htmlData{mx=49}{\sum_{i} \htmlData{mx=48}{\sum_{j \in \htmlData{mx=36}{\htmlData{mx=34}{\operatorname{erase}}(\htmlData{mx=35}{i})}} \htmlData{mx=47}{\htmlData{mx=41}{\htmlData{mx=37}{U}(\htmlData{mx=40}{\htmlData{mx=38}{f}(\htmlData{mx=39}{i})})} \cdot \htmlData{mx=46}{\htmlData{mx=42}{K}(\htmlData{mx=45}{\htmlData{mx=43}{f}(\htmlData{mx=44}{j})})}}}}}}\right)}} \\ & \quad \htmlData{mx=0}{=} \htmlData{mx=83}{\htmlData{mx=65}{\sum_{\omega} \htmlData{mx=64}{\htmlData{mx=56}{\htmlData{mx=54}{w}(\htmlData{mx=55}{\omega})} \cdot (\htmlData{mx=63}{\htmlData{mx=59}{\htmlData{mx=57}{U}(\htmlData{mx=58}{\omega})} \cdot \htmlData{mx=62}{\htmlData{mx=60}{K}(\htmlData{mx=61}{\omega})}})}} - \htmlData{mx=82}{\left(\htmlData{mx=73}{\sum_{\omega} \htmlData{mx=72}{\htmlData{mx=68}{\htmlData{mx=66}{w}(\htmlData{mx=67}{\omega})} \cdot \htmlData{mx=71}{\htmlData{mx=69}{U}(\htmlData{mx=70}{\omega})}}}\right) \cdot \left(\htmlData{mx=81}{\sum_{\omega} \htmlData{mx=80}{\htmlData{mx=76}{\htmlData{mx=74}{w}(\htmlData{mx=75}{\omega})} \cdot \htmlData{mx=79}{\htmlData{mx=77}{K}(\htmlData{mx=78}{\omega})}}}\right)}} \end{aligned}

What the notation means · or point at any symbol

Type: Type: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (Type)
Fintype: Fintype: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (Fintype)
f(x)Function application: Fintype applied to Ω. Lean writes f x; papers write f(x). (Fintype Ω)
NNatural numbers: The natural numbers 0, 1, 2, …. In Lean they start at 0. (ℕ)
⟹Implies: If Ω, then ℝ. (Between types, A → B is the type of functions from A to B.) (Ω → ℝ)
x : TVariables: w, U, K range over Ω → ℝ: the statement must hold for every choice. ((w U K : Ω → ℝ))
∑Summation: Add up w(ω) for every ω. (∑ ω, w ω)
=Equality: ∑ ω, w(ω) and 1 are the same. (∑ ω, w ω = 1)
≤At most: 2 is less than or equal to N. (2 ≤ N)
RReal numbers: The real numbers: every point on the number line. (ℝ)
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)
∏Product: Multiply together w(f(m)) for every m. (∏ m, w (f m))
a/bDivision: 1 divided by N. In Lean, dividing by 0 gives 0. (1 / (N : ℝ))
(x : T)As real numbers: N read as an element of real numbers. Mathematicians leave this implicit; Lean needs it to know which kind of number is meant. ((N : ℝ))
·Multiplication: U(f(i)) times K(f(i)). (U (f i) * K (f i))
-Subtraction: N minus 1. On ℕ, Lean's subtraction stops at 0: 2 − 5 = 0. (N - 1)
erase: Finset.univ.erase: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (Finset.univ.erase)

theorem op-one-pass-gradient-unbiased {Ω : Type} [Fintype Ω] {N : ℕ} (w U K : Ω → ℝ) (hsum : ∑ ω, w ω = 1) (hN : 2 ≤ N) :
    ∑ f : Fin N → Ω, (∏ m, w (f m)) * ((1 / (N : ℝ)) * ∑ i, U (f i) * K (f i) - (1 / ((N : ℝ) * (N - 1))) * ∑ i, ∑ j ∈ Finset.univ.erase i, U (f i) * K (f j)) = ∑ ω, w ω * (U ω * K ω) - (∑ ω, w ω * U ω) * (∑ ω, w ω * K ω)

The query-gradient coordinate is the covariance Cov_π(u, k). Given N exact samples from the attention row (exact rejection from a snapshot or an INT8 proposal, qc-/tp-), the U-statistic above is exactly unbiased, so the backward pass reuses the forward samples and needs no new key reads. The unbiased sample covariance is textbook (Bessel/U-statistics, credited); what is new is its role in the training step: backward attention at zero extra KV traffic under exact sampling. Part of 'Unbiased randomized attention', gradient round.

How it was proved

Proved: 5 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: 84 lines of Lean

One-pass attention gradient: from N ≥ 2 exact samples, (1/N)Σ U_i K_i − (1/(N(N−1)))Σ_{i≠j} U_i K_j is unbiased for Cov(U, K)

How it was proved

More proved statements