← All papers
First page of Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean

Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean

Yusuf Khaled Sadek, Enes Yilmaz

cs.LG Oct 12, 2026 · v1
Proves padding and permutation invariance of masked attention pooling in Lean 4 with mathlib, and links the proofs to traced JAX code via jax-lean certificates.
Attention-based multiple instance learning (MIL) pools a variable-size bag of instance embeddings into one vector. In batched training, bags are padded to a common length, and the pooled output is expected not to depend on how much padding was added or on the order in which instances are stored. This expectation holds only if padded rows are removed from both the numerator and the softmax normalizer. Common implementation idioms meet it only under side conditions. Zeroing padded features alone does not work. Masking logits with $-\infty$ fails on empty bags and on nonfinite padding. Masking with a large negative constant is exact only after floating-point underflow. We define masked attention pooling with a guarded normalizer and prove in Lean 4 with mathlib that it is invariant under any bijective reindexing of rows and under appending any finite number of masked rows with arbitrary real features and logits. We also prove that values at masked positions are irrelevant and that an all-masked bag maps to zero. The theorems hold for arbitrary finite index types and feature dimensions. Using the jax-lean certificate pipeline, we connect these theorems to the Jaxpr traced from the JAX kernel itself. Lean checks translation certificates for two concrete shapes, identifies the generated functions with the generic specification, and transfers both invariances to the imported Jaxpr's real-arithmetic semantics. All audited theorems depend only on Lean's three standard axioms. The results concern ideal real arithmetic. Floating-point execution, XLA, and the JAX tracer and importer remain outside the proof.

Batched attention-based multiple instance learning (MIL) pooling pads bags to a common length. The output should not depend on the amount of padding or on row order. Common masking idioms, such as zeroing padded features or masking logits with -inf or a large negative constant, break this under side conditions such as empty bags or nonfinite padding.

The authors define masked attention pooling with a guarded softmax normalizer. In Lean 4 with mathlib they prove, for arbitrary finite index types and feature dimensions, four results: invariance under bijective reindexing, invariance under appending masked rows with arbitrary contents, irrelevance of masked values, and zero output on all-masked bags. Using jax-lean, they trace the JAX kernel at two shapes and check translation certificates. These certificates identify the generated functions with the generic specification and transfer the invariances to the imported Jaxpr's real-arithmetic semantics.

Seven audited theorems depend only on Lean's three standard axioms, with no sorry, in a development of about 560 lines. Separate JAX tests confirm the expected floating-point behavior. Floating-point execution, XLA, and the JAX tracer and importer remain outside the proof.

IdiomZero padNaN padEmptyExact in R
Zero padded features only(1.3563,0.8921)(1.3563,0.8921)(0,0)no
Logits to -inf, features unmaskedcorrectNaNNaNn/a
Logits to -1e9, features unmaskedcorrectNaN(2,3)no
Softmax of -inf-masked logits, features zeroedcorrectcorrectNaNn/a
Definition 1 (masked_pool)correctcorrect(0,0)yes
Behaviour of common masking idioms in JAX (float32)