Provedattentiontransformersoftmaxformal-proof

Swapping both keys and values preserves two-key attention

For all $\htmlData{mx=1}{s_{0}, s_{1}, v_{0}, v_{1} \in \htmlData{mx=0}{\mathbb{R}}}$ :

\htmlData{mx=5}{\htmlData{mx=0}{\operatorname{attend2}}(\htmlData{mx=1}{s_{0}}, \htmlData{mx=2}{s_{1}}, \htmlData{mx=3}{v_{0}}, \htmlData{mx=4}{v_{1}})} \mathrel{\htmlData{mx=12}{=}} \htmlData{mx=11}{\htmlData{mx=6}{\operatorname{attend2}}(\htmlData{mx=7}{s_{1}}, \htmlData{mx=8}{s_{0}}, \htmlData{mx=9}{v_{1}}, \htmlData{mx=10}{v_{0}})}

What the notation means · or point at any symbol

RReal numbers: The real numbers: every point on the number line. (ℝ)
x ∈ TVariables: s₀, s₁, v₀, v₁ range over real numbers: the statement must hold for every choice. ((s₀ s₁ v₀ v₁ : ℝ))
attend2: attend2: a definition from the library, not special notation. Its name is its meaning; open it in Lean to see how it is defined. (attend2)
f(x)Function application: attend2 applied to s₀, s₁, v₀, v₁. Lean writes f x; papers write f(x). (attend2 s₀ s₁ v₀ v₁)
=Equality: attend2(s₀, s₁, v₀, v₁) and attend2(s₁, s₀, v₁, v₀) are the same. (attend2 s₀ s₁ v₀ v₁ = attend2 s₁ s₀ v₁ …)

theorem attention2-swap-equivariance (s₀ s₁ v₀ v₁ : ℝ) :
    attend2 s₀ s₁ v₀ v₁ = attend2 s₁ s₀ v₁ v₀

A property of one-query, two-key scaled dot-product attention with scalar values; scores already include the dot-product scale.

How it was proved

Proved: 26 September 2026
By: Ajax, with the agent “Codex attention proofs 2026-09-26”
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: 16 lines of Lean