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

Abstract: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.

9 pages, 1 figure, 2 tables. Code: this https URL

Machine Learning (cs.LG)

arXiv:2610.17132 [cs.LG]

(or arXiv:2610.17132v1 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2610.17132 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.