Provedllmlarge-language-modelsformal-proofattention

Rescaling by sqrt(d) recovers the raw dot product

For all $\htmlData{mx=1}{d \in \htmlData{mx=0}{\mathbb{N}}}$ and $\htmlData{mx=5}{q, k : \htmlData{mx=2}{\htmlData{mx=0}{\operatorname{Fin}}\,\htmlData{mx=1}{d}} \mathrel{\htmlData{mx=4}{\to}} \htmlData{mx=3}{\mathbb{R}}}$ , if $\htmlData{mx=0}{d} \mathrel{\htmlData{mx=1}{\neq}} 0$ , then

\htmlData{mx=8}{\htmlData{mx=4}{\htmlData{mx=0}{\operatorname{logit}}(\htmlData{mx=1}{d}, \htmlData{mx=2}{q}, \htmlData{mx=3}{k})} \cdot \htmlData{mx=7}{\sqrt{\htmlData{mx=6}{\htmlData{mx=5}{d}}}}} \mathrel{\htmlData{mx=17}{=}} \htmlData{mx=16}{\sum_{i} \htmlData{mx=15}{\htmlData{mx=11}{\htmlData{mx=9}{q}(\htmlData{mx=10}{i})} \cdot \htmlData{mx=14}{\htmlData{mx=12}{k}(\htmlData{mx=13}{i})}}}

What the notation means · or point at any symbol

NNatural numbers: The natural numbers 0, 1, 2, …. In Lean they start at 0. (ℕ)
≠Not equal: d and 0 are different. (d ≠ 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 d. Lean writes f x; papers write f(x). (Fin d)
⟹Implies: If Fin(d), then ℝ. (Between types, A → B is the type of functions from A to B.) (Fin d → ℝ)
x : TVariables: q, k range over Fin d → ℝ: the statement must hold for every choice. ((q k : Fin d → ℝ))
logit: LLMAttention.logit: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (LLMAttention.logit)
·Multiplication: logit(d, q, k) times √(d). (LLMAttention.logit d q k * Real.sqrt (d…)
=Equality: logit(d, q, k) · √(d) and ∑ i, q(i) · k(i) are the same. (LLMAttention.logit d q k * Real.sqrt (d…)
RReal numbers: The real numbers: every point on the number line. (ℝ)
√Square root: The non-negative number whose square is d. (Real.sqrt (d : ℝ))
(x : T)As real numbers: d read as an element of real numbers. Mathematicians leave this implicit; Lean needs it to know which kind of number is meant. ((d : ℝ))
∑Summation: Add up q(i) · k(i) for every i. (∑ i, q i * k i)

theorem attn-logit-scale-identity (d : ℕ) (hd : d ≠ 0) (q k : Fin d → ℝ) :
    LLMAttention.logit d q k * Real.sqrt (d : ℝ) = ∑ i, q i * k i

For d ≠ 0, logit d q k * sqrt d = Σ qᵢkᵢ: the 1/√d factor is exactly a division.

How it was proved

Proved: 28 September 2026
By: Ajax, with the agent “ompmuse”
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: 15 lines of Lean

Rescaling by sqrt(d) recovers the raw dot product

How it was proved

More proved statements