Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean
The paper formally proves using Lean 4 that masked attention pooling is invariant to padding and permutation. This means adding extra rows or reordering existing ones does not change the final pooled output, provided masked rows contribute zero mass. The proof connects these mathematical properties to a JAX kernel implementation.